Update main.py

main.py

@@ -1,42 +1,51 @@
-import streamlit as st
+# main.py
+
 import requests
-from PIL import Image
-from io import BytesIO
-import getpass, os
+import os
 import warnings
+import io
+import getpass
+
+from PIL import Image
 from stability_sdk import client
 import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
 import replicate
 
-# 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
-
-# Increase the pixel limit
-Image.MAX_IMAGE_PIXELS = None
-
-# 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:
-        st.write(f"Request failed with status code {response.status_code}")
+def generate_and_upscale_image(text_prompt, clipdrop_api_key, stability_api_key, replicate_api_token):
+
+    headers = {'x-api-key': clipdrop_api_key}
+    body_params = {'prompt': (None, text_prompt, 'text/plain')}
+
+    response = requests.post('https://clipdrop-api.co/text-to-image/v1',
+                             files=body_params,
+                             headers=headers)
+
+    if response.status_code != 200:
+        print(f"Request failed with status code {response.status_code}")
         return None
 
-# Stability API function
-def upscale_image_stability(img):
+    with open('generated_image.png', 'wb') as f:
+        f.write(response.content)
+
+    os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'
+    os.environ['STABILITY_KEY'] = stability_api_key
+
+    stability_api = client.StabilityInference(
+        key=os.environ['STABILITY_KEY'],
+        upscale_engine="esrgan-v1-x2plus",
+        verbose=True,
+    )
+
+    max_pixels = 1048576
+    img = Image.open('generated_image.png')
+    width, height = img.size
+
+    if width * height > max_pixels:
+        scale_factor = (max_pixels / (width * height))**0.5
+        new_width = int(width * scale_factor)
+        new_height = int(height * scale_factor)
+        img = img.resize((new_width, new_height))
+
     answers = stability_api.upscale(init_image=img)
 
     for resp in answers:
@@ -46,37 +55,20 @@ def upscale_image_stability(img):
                     "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:
-                return Image.open(io.BytesIO(artifact.binary))
+                upscaled_img = Image.open(io.BytesIO(artifact.binary))
+                upscaled_img.save("upscaled_image.png")
+
+    os.environ['REPLICATE_API_TOKEN'] = replicate_api_token
+    Image.MAX_IMAGE_PIXELS = None
 
-# GFPGAN function
-def upscale_image_gfpgan(image_path):
-    with open(image_path, "rb") as img_file:
+    with open("upscaled_image.png", "rb") as img_file:
         output = replicate.run(
             "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
             input={"img": img_file, "version": "v1.4", "scale": 16}
         )
-        response = requests.get(output)
-        return Image.open(BytesIO(response.content))
-
-# Streamlit UI
-st.title("Image Generator and Upscaler")
-
-prompt = st.text_input("Enter a prompt for the image generation")
-
-if st.button("Generate and Upscale"):
-    if prompt:
-        img1 = generate_image(prompt)
-
-        if img1:
-            st.image(img1, caption="Generated Image", use_column_width=True)
-            img1.save('generated_image.png')
 
-            img2 = upscale_image_stability(img1)
-            st.image(img2, caption="Upscaled Image (Stability API)", use_column_width=True)
-            img2.save('upscaled_image_stability.png')
+    response = requests.get(output)
+    final_img = Image.open(io.BytesIO(response.content))
+    final_img.save("gfpgan_upscaled_image.png")
 
-            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")
+    return final_img