Create main.py

main.py

@@ -0,0 +1,104 @@
+import streamlit as st
+import os
+import requests
+from PIL import Image
+from io import BytesIO
+
+# Set up environment variables for API keys
+os.environ['CLIPDROP_API_KEY'] = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
+os.environ['STABILITY_API_KEY'] = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
+os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'
+
+# Importing Replicate and Stability SDK libraries
+import replicate
+import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
+
+def generate_image(prompt):
+    # Make a POST request to the ClipDrop text-to-image API
+    url = 'https://clipdrop-api.co/text-to-image/v1'
+    headers = {'x-api-key': os.environ['CLIPDROP_API_KEY']}
+    data = {'prompt': prompt}
+    response = requests.post(url, headers=headers, data=data)
+
+    if response.status_code == 200:
+        # Get the generated image from the response
+        img = Image.open(BytesIO(response.content))
+        return img
+    else:
+        st.error('Failed to generate image from text prompt.')
+
+def upscale_image(img):
+    # Set up Stability API connection
+    stability_api = replicate.StabilityInference(
+        key=os.environ['STABILITY_API_KEY'],
+        upscale_engine="stable-diffusion-x4-latent-upscaler"
+    )
+
+    # Upscale the image using the Stability API
+    upscale_responses = stability_api.upscale(init_image=img)
+
+    if upscale_responses:
+        # Get the upscaled image from the response
+        upscaled_img = None
+        for resp in upscale_responses:
+            for artifact in resp.artifacts:
+                if artifact.type == generation.ARTIFACT_IMAGE:
+                    upscaled_img = Image.open(BytesIO(artifact.binary))
+                    break
+            if upscaled_img:
+                break
+        return upscaled_img
+    else:
+        st.error('Failed to upscale the image.')
+
+def gfpgan_upscale_image(img):
+    # Set up Replicate API connection
+    replicate_api = replicate.ReplicateAPI(token=os.environ['REPLICATE_API_TOKEN'])
+
+    # Upscale the image using GFPGAN model
+    output = replicate_api.run(
+        "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
+        input={"img": img, "version": "v1.4", "scale": 16}
+    )
+
+    if output:
+        # Get the upscaled image URI from the output
+        response = requests.get(output)
+        upscaled_img = Image.open(BytesIO(response.content))
+        return upscaled_img
+    else:
+        st.error('Failed to upscale the image using GFPGAN.')
+
+def main():
+    st.title("Image Upscaling")
+    st.write("Enter a text prompt to generate and upscale an image.")
+
+    # Get user input for the text prompt
+    prompt = st.text_input("Enter a text prompt:", max_chars=1000)
+
+    if st.button("Generate and Upscale"):
+        # Generate the image from text prompt using ClipDrop API
+        generated_img = generate_image(prompt)
+
+        if generated_img:
+            st.image(generated_img, caption='Generated Image', use_column_width=True)
+
+            # Upscale the generated image using Stability API
+            upscaled_img = upscale_image(generated_img)
+
+            if upscaled_img:
+                st.image(upscaled_img, caption='Upscaled Image (Stability API)', use_column_width=True)
+
+                # Further upscale the image using GFPGAN
+                gfpgan_upscaled_img = gfpgan_upscale_image(upscaled_img)
+
+                if gfpgan_upscaled_img:
+                    st.image(gfpgan_upscaled_img, caption='Upscaled Image (GFPGAN)', use_column_width=True)
+                    st.success("Image generation and upscaling completed successfully.")
+                else:
+                    st.error("Failed to further upscale the image using GFPGAN.")
+            else:
+                st.error("Failed to upscale the generated image.")
+    
+if __name__ == "__main__":
+    main()