Update main.py

main.py

@@ -10,6 +10,7 @@ import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
 # Configure your API keys here
 CLIPDROP_API_KEY = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
 STABLE_DIFFUSION_API_KEY = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
+ESRGAN_API_KEY = 'your_esrgan_api_key_here'
 
 # Set up environment variable for Replicate API Token
 os.environ['REPLICATE_API_TOKEN'] = 'r8_Tm3LQMS81QaGXzzdGVRyUCOQ3cuNd1i1sJlqp'  # Replace with your actual API token
@@ -39,15 +40,14 @@ def resize_image(image_bytes, max_size=(256, 256)):
     img.save(buffer, format="PNG")
     return buffer.getvalue()
 
-def upscale_image_stable_diffusion(image_bytes):
+def upscale_image_esrgan(image_bytes):
     # Set up environment variables
-    os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'
-    os.environ['STABILITY_KEY'] = STABLE_DIFFUSION_API_KEY
+    os.environ['ESRGAN_API_KEY'] = ESRGAN_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",
+        key=os.environ['ESRGAN_API_KEY'],
+        upscale_engine="esrgan-v1-x2plus",
         verbose=True,
     )
 
@@ -96,8 +96,8 @@ def main():
         st.success("Resizing image...")
         resized_image_bytes = resize_image(image_bytes)
         
-        st.success("Upscaling image with stable-diffusion-x4-latent-upscaler...")
-        upscaled_image_bytes = upscale_image_stable_diffusion(resized_image_bytes) # Change this line
+        st.success("Upscaling image with ESRGAN...")
+        upscaled_image_bytes = upscale_image_esrgan(resized_image_bytes)
         
         st.success("Further upscaling image with GFPGAN...")
         img = further_upscale_image(upscaled_image_bytes)