Update main.py

main.py

@@ -1,82 +1,83 @@
 import streamlit as st
+import replicate
+import os
 import requests
 from PIL import Image
 from io import BytesIO
-import getpass, os
-import warnings
-from stability_sdk import client
-import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
-import replicate
 
-# API keys
-api_key = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
-os.environ['STABILITY_KEY'] = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
-os.environ['REPLICATE_API_TOKEN'] = 'r8_0a77UG8yfrNXOS6xHhUTLh80dJQ5kxO0CTLmq'  # 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}")
-        return None
-
-# Stability API function
-def upscale_image_stability(img):
-    answers = stability_api.upscale(init_image=img)
-
-    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:
-                return Image.open(io.BytesIO(artifact.binary))
-
-# GFPGAN function
-def upscale_image_gfpgan(image_path):
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+# Set up environment variable for Replicate API Token
+os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'  # Replace with your actual API token
+
+
+def upscale_image(image_path):
+    # Open the image file
+
+
+
+
+
+
+
+
+
     with open(image_path, "rb") as img_file:
+        # Run the GFPGAN model
         output = replicate.run(
             "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
             input={"img": img_file, "version": "v1.4", "scale": 16}
         )
+        
+        # The output is a URI of the processed image
+        # We will retrieve the image data and save it
         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)
+        img = Image.open(BytesIO(response.content))
+        img.save("upscaled.png")  # Save the upscaled image
+        return img
 
-        if img1:
-            st.image(img1, caption="Generated Image", use_column_width=True)
-            img1.save('generated_image.png')
+def main():
+    st.title("Image Upscaling")
+    st.write("Upload an image and it will be upscaled.")
 
-            img2 = upscale_image_stability(img1)
-            st.image(img2, caption="Upscaled Image (Stability API)", use_column_width=True)
-            img2.save('upscaled_image_stability.png')
+    uploaded_file = st.file_uploader("Choose an image...", type="png")
+    if uploaded_file is not None:
+        with open("temp_img.png", "wb") as f:
+            f.write(uploaded_file.getbuffer())
+        st.success("Uploaded image successfully!")
+        if st.button("Upscale Image"):
+            img = upscale_image("temp_img.png")
+            st.image(img, caption='Upscaled Image', use_column_width=True)
 
-            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")
+if __name__ == "__main__":
+    main()