Update main.py

main.py

@@ -1,132 +1,82 @@
 import streamlit as st
-import os
 import requests
-import base64
 from PIL import Image
 from io import BytesIO
-import replicate
+import getpass, os
+import warnings
 from stability_sdk import client
 import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
+import replicate
 
-Image.MAX_IMAGE_PIXELS = None
-
-
-# Configure your API keys here
-CLIPDROP_API_KEY = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
-STABLE_DIFFUSION_API_KEY = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
-ESRGAN_API_KEY = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
+# API keys
+api_key = 'YOUR_API_KEY'
+os.environ['STABILITY_KEY'] = 'YOUR_API_KEY'
+os.environ['REPLICATE_API_TOKEN'] = 'REPLICATE_API_TOKEN'  # Replace with your actual API token
 
-# Set up environment variable for Replicate API Token
-os.environ['REPLICATE_API_TOKEN'] = 'r8_Tm3LQMS81QaGXzzdGVRyUCOQ3cuNd1i1sJlqp'  # Replace with your actual API token
+# Increase the pixel limit
+Image.MAX_IMAGE_PIXELS = None
 
-def generate_image_from_text(prompt):
-    r = requests.post('https://clipdrop-api.co/text-to-image/v1',
-        files = {
-            'prompt': (None, prompt, 'text/plain')
-        },
-        headers = { 'x-api-key': CLIPDROP_API_KEY }
-    )
-    
-    if r.ok:
-        return r.content
+# Establish connection to Stability API
+stability_api = client.StabilityInference(
+    key=os.environ['STABILITY_KEY'], 
+    upscale_engine="esrgan-v1-x2plus",
+    verbose=True,
+)
+
+# ClipDrop API function
+def generate_image(prompt):
+    headers = {'x-api-key': api_key}
+    body_params = {'prompt': (None, prompt, 'text/plain')}
+    response = requests.post('https://clipdrop-api.co/text-to-image/v1', files=body_params, headers=headers)
+
+    if response.status_code == 200:
+        return Image.open(BytesIO(response.content))
     else:
-        r.raise_for_status()
-
-def upscale_image_esrgan(image_bytes):
-    # Set up environment variables
-    os.environ['ESRGAN_API_KEY'] = ESRGAN_API_KEY
-
-    # Set up the connection to the API
-    stability_api = client.StabilityInference(
-        key=os.environ['ESRGAN_API_KEY'],
-        upscale_engine="esrgan-v1-x2plus",
-        verbose=True,
-    )
+        st.write(f"Request failed with status code {response.status_code}")
+        return None
 
-    # Open the image from bytes
-    img = Image.open(BytesIO(image_bytes))
-
-    # Call the upscale API
+# Stability API function
+def upscale_image_stability(img):
     answers = stability_api.upscale(init_image=img)
 
-    # Process the response
-    upscaled_img_bytes = None
     for resp in answers:
         for artifact in resp.artifacts:
+            if artifact.finish_reason == generation.FILTER:
+                warnings.warn(
+                    "Your request activated the API's safety filters and could not be processed."
+                    "Please submit a different image and try again.")
             if artifact.type == generation.ARTIFACT_IMAGE:
-                upscaled_img = Image.open(BytesIO(artifact.binary))
-                upscaled_img_bytes = BytesIO()
-                upscaled_img.save(upscaled_img_bytes, format='PNG')
-                upscaled_img_bytes = upscaled_img_bytes.getvalue()
-    
-    return upscaled_img_bytes
-
-def further_upscale_image(image_bytes):
-    # Ensure environment variable is set correctly
-    print("Replicate API token: ", os.environ['REPLICATE_API_TOKEN'])
+                return Image.open(io.BytesIO(artifact.binary))
 
-    # Save the image bytes to a temporary file
-    temp_file_name = "temp.png"
-    with open(temp_file_name, 'wb') as temp_file:
-        temp_file.write(image_bytes)
-
-    # Run the GFPGAN model
-    try:
-        print("Running GFPGAN model...")
+# GFPGAN function
+def upscale_image_gfpgan(image_path):
+    with open(image_path, "rb") as img_file:
         output = replicate.run(
             "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
-            input={"img": open(temp_file_name, "rb"), "version": "v1.4", "scale": 16}
+            input={"img": img_file, "version": "v1.4", "scale": 16}
         )
-        print("Model output: ", output)
-    except Exception as e:
-        print("Error running GFPGAN model: ", e)
-        raise e
-
-    # Get the image data from the output URI
-    try:
-        print("Fetching image data from output URI...")
         response = requests.get(output)
-    except Exception as e:
-        print("Error fetching image data from output URI: ", e)
-        raise e
+        return Image.open(BytesIO(response.content))
 
-    # Open and save the image
-    try:
-        print("Saving upscaled image...")
-        img = Image.open(BytesIO(response.content))
-        output_file = "upscaled.png"
-        img.save(output_file)  # Save the upscaled image
-    except Exception as e:
-        print("Error saving upscaled image: ", e)
-        raise e
+# Streamlit UI
+st.title("Image Generator and Upscaler")
 
-    # Create a function to make download link
-    def create_download_link(file, filename):
-        with open(file, 'rb') as f:
-            bytes = f.read()
-            b64 = base64.b64encode(bytes).decode()
-            href = f'<a href="data:file/octet-stream;base64,{b64}" download="{filename}">Download File</a>'
-            return href
+prompt = st.text_input("Enter a prompt for the image generation")
 
-    return create_download_link(output_file, "upscaled_image.png")
+if st.button("Generate and Upscale"):
+    if prompt:
+        img1 = generate_image(prompt)
 
-def main():
-    st.title("Image Generation and Upscaling")
-    st.write("Enter a text prompt and an image will be generated and upscaled.")
+        if img1:
+            st.image(img1, caption="Generated Image", use_column_width=True)
+            img1.save('generated_image.png')
 
-    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 ESRGAN...")
-        upscaled_image_bytes = upscale_image_esrgan(image_bytes)
-        
-        st.success("Further upscaling image with GFPGAN...")
-        download_link = further_upscale_image(upscaled_image_bytes)
-        
-        st.markdown(download_link, unsafe_allow_html=True)
+            img2 = upscale_image_stability(img1)
+            st.image(img2, caption="Upscaled Image (Stability API)", use_column_width=True)
+            img2.save('upscaled_image_stability.png')
 
-if __name__ == "__main__":
-    main()
+            img3 = upscale_image_gfpgan('upscaled_image_stability.png')
+            st.image(img3, caption="Upscaled Image (GFPGAN)", use_column_width=True)
+            img3.save('upscaled_image_gfpgan.png')
+    else:
+        st.write("Please enter a prompt")