Update main.py

main.py

@@ -1,11 +1,11 @@
-# main.py
-
 import streamlit as st
 import os
 import requests
 from PIL import Image
 from io import BytesIO
 import replicate
+from stability_sdk import client
+import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
 
 # Configure your API keys here
 CLIPDROP_API_KEY = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
@@ -28,18 +28,34 @@ def generate_image_from_text(prompt):
         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)
+    # Set up environment variables
+    os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'
+    os.environ['STABILITY_KEY'] = STABLE_DIFFUSION_API_KEY
+
+    # Set up the connection to the API
+    stability_api = client.StabilityInference(
+        key=os.environ['STABILITY_KEY'],
+        upscale_engine="stable-diffusion-x4-latent-upscaler",
+        verbose=True,
+    )
+
+    # Open the image from bytes
+    img = Image.open(BytesIO(image_bytes))
+
+    # Call the upscale API
+    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.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()
     
-    if r.ok:
-        return r.content
-    else:
-        r.raise_for_status()
+    return upscaled_img_bytes
 
 def further_upscale_image(image_bytes):
     # Run the GFPGAN model