Update main.py

main.py

@@ -1,47 +1,77 @@
+# main.py
+
 import streamlit as st
-import replicate
 import os
 import requests
 from PIL import Image
 from io import BytesIO
+import replicate
+
+# Configure your API keys here
+CLIPDROP_API_KEY = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
+STABLE_DIFFUSION_API_KEY = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
 
-# Set up environment variables for Replicate and Stability AI API Tokens
-os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'  # Replace with your actual Replicate API token
-os.environ['STABILITY_KEY'] = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'  # Replace with your actual Stability AI API key
-
-def upscale_image(image_path):
-    # Open the image file
-    with open(image_path, "rb") as img_file:
-        # Upscale the image using Stability AI's API and Stable Diffusion 4x Upscaler
-        stability_api_url = "https://api.stability.ai/v1/image-upscaling"
-        headers = {
-            'Authorization': f'Bearer {os.environ["STABILITY_KEY"]}',
-        }
+# Set up environment variable for Replicate API Token
+os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'  # Replace with your actual API token
+
+def generate_image_from_text(prompt):
+    r = requests.post('https://clipdrop-api.co/text-to-image/v1',
         files = {
-            'image': img_file,
-            'upscaler': 'stable-diffusion-x4-latent-upscaler',
-        }
-        response = requests.post(stability_api_url, headers=headers, files=files)
+            'prompt': (None, prompt, 'text/plain')
+        },
+        headers = { 'x-api-key': CLIPDROP_API_KEY }
+    )
+    
+    if r.ok:
+        return r.content
+    else:
+        r.raise_for_status()
+
+def upscale_image_stable_diffusion(image_bytes):
+    url = 'https://stable-diffusion-x4-latent-upscaler.com/v1/upscaling' # Update this with correct API endpoint
+    headers = { 'x-api-key': STABLE_DIFFUSION_API_KEY }
+    files = {
+        'image': ('image.png', image_bytes, 'image/png')
+    }
+    
+    r = requests.post(url, headers=headers, files=files)
+    
+    if r.ok:
+        return r.content
+    else:
+        r.raise_for_status()
 
-        # Retrieve the upscaled image from the response
-        img = Image.open(BytesIO(response.content))
-        img.save("upscaled.png")  # Save the upscaled image
-        return img
+def further_upscale_image(image_bytes):
+    # Run the GFPGAN model
+    output = replicate.run(
+        "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
+        input={"img": BytesIO(image_bytes), "version": "v1.4", "scale": 16}
+    )
+    
+    # The output is a URI of the processed image
+    # We will retrieve the image data and save it
+    response = requests.get(output)
+    img = Image.open(BytesIO(response.content))
+    img.save("upscaled.png")  # Save the upscaled image
+    return img
 
 def main():
-    st.title("Image Upscaling")
-    st.write("Upload an image and enter a prompt to generate and upscale an image.")
-
-    uploaded_file = st.file_uploader("Choose an image...", type="png")
-    prompt = st.text_input("Enter a prompt for image generation")
-
-    if uploaded_file is not None and prompt:
-        with open("temp_img.png", "wb") as f:
-            f.write(uploaded_file.getbuffer())
-        st.success("Uploaded image successfully!")
-        if st.button("Generate and Upscale Image"):
-            img = upscale_image("temp_img.png")
-            st.image(img, caption='Generated and Upscaled Image', use_column_width=True)
+    st.title("Image Generation and Upscaling")
+    st.write("Enter a text prompt and an image will be generated and upscaled.")
+
+    prompt = st.text_input("Enter a textual prompt to generate an image...")
+    
+    if prompt:
+        st.success("Generating image from text prompt...")
+        image_bytes = generate_image_from_text(prompt)
+        
+        st.success("Upscaling image with stable-diffusion-x4-latent-upscaler...")
+        upscaled_image_bytes = upscale_image_stable_diffusion(image_bytes)
+        
+        st.success("Further upscaling image with GFPGAN...")
+        img = further_upscale_image(upscaled_image_bytes)
+        
+        st.image(img, caption='Upscaled Image', use_column_width=True)
 
 if __name__ == "__main__":
     main()