Update main.py

main.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, render_template
+import streamlit as st
 import os
 import requests
 from PIL import Image
@@ -12,9 +12,7 @@ CLIPDROP_API_KEY = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49b
 STABLE_DIFFUSION_API_KEY = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
 
 # Set up environment variable for Replicate API Token
-os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'
-
-app = Flask(__name__)
+os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'  # Replace with your actual API token
 
 def generate_image_from_text(prompt):
     r = requests.post('https://clipdrop-api.co/text-to-image/v1',
@@ -30,22 +28,36 @@ def generate_image_from_text(prompt):
         r.raise_for_status()
 
 def resize_image(image_bytes, max_size=(320, 320)):
+    # Open the image from bytes
     img = Image.open(BytesIO(image_bytes))
+    
+    # Resize the image
     img.thumbnail(max_size)
+    
+    # Save it back to bytes
     buffer = BytesIO()
     img.save(buffer, format="PNG")
     return buffer.getvalue()
 
 def upscale_image_stable_diffusion(image_bytes):
+    # 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:
@@ -54,30 +66,43 @@ def upscale_image_stable_diffusion(image_bytes):
                 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):
+    # Run the GFPGAN model
     output = replicate.run(
-        "tencentarc/gfpgan:your_model_version",
+        "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
         input={"img": BytesIO(image_bytes), "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)
     img = Image.open(BytesIO(response.content))
+    img.save("upscaled.png")  # Save the upscaled image
     return img
 
-@app.route('/', methods=['GET', 'POST'])
-def index():
-    if request.method == 'POST':
-        prompt = request.form.get('prompt')
-        if prompt:
-            image_bytes = generate_image_from_text(prompt)
-            resized_image_bytes = resize_image(image_bytes)
-            upscaled_image_bytes = upscale_image_stable_diffusion(resized_image_bytes)
-            img = further_upscale_image(upscaled_image_bytes)
-            img_path = os.path.join('static', 'upscaled_image.png')
-            img.save(img_path) 
-            return render_template('index.html', image_url=img_path)
-    return render_template('index.html')
+def main():
+    st.title("Image Generation and Upscaling")
+    st.write("Enter a text prompt and an image will be generated and upscaled.")
+
+    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("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)
+        
+        st.success("Further upscaling image with GFPGAN...")
+        img = further_upscale_image(upscaled_image_bytes)
+        
+        st.image(img, caption='Upscaled Image', use_column_width=True)
 
 if __name__ == "__main__":
-    app.run(port=8501, debug=True)
+    main()