Update main.py

main.py

@@ -1,104 +1,42 @@
 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_Tm3LQMS81QaGXzzdGVRyUCOQ3cuNd1i1sJlqp'
-
-# Importing Replicate and Stability SDK libraries
-import replicate
-import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
+API_KEY = "1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea"
+API_URL = "https://clipdrop-api.co/text-to-image/v1"
 
 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)
+    headers = {"x-api-key": API_KEY}
+    files = {"prompt": (None, prompt, "text/plain")}
 
-    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.')
+    try:
+        response = requests.post(API_URL, files=files, headers=headers)
+        response.raise_for_status()
 
-def gfpgan_upscale_image(img):
-    # Set up Replicate API connection
-    replicate_api = replicate.ReplicateAPI(token=os.environ['REPLICATE_API_TOKEN'])
+        # Get the generated image
+        image = Image.open(BytesIO(response.content))
 
-    # 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.')
+        return image
+    except requests.exceptions.RequestException as e:
+        st.error(f"Error occurred during image generation: {str(e)}")
+        return None
 
 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)
+    st.title("Text-to-Image Generator")
 
-            # Upscale the generated image using Stability API
-            upscaled_img = upscale_image(generated_img)
+    # Text prompt input
+    prompt = st.text_input("Enter a text prompt")
 
-            if upscaled_img:
-                st.image(upscaled_img, caption='Upscaled Image (Stability API)', use_column_width=True)
+    if prompt:
+        # Generate image when the "Generate Image" button is clicked
+        if st.button("Generate Image"):
+            st.write("Generating image...")
+            image = generate_image(prompt)
 
-                # Further upscale the image using GFPGAN
-                gfpgan_upscaled_img = gfpgan_upscale_image(upscaled_img)
+            if image:
+                # Display the generated image
+                st.image(image, caption="Generated Image", use_column_width=True)
 
-                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()