Update main.py

main.py

@@ -1,4 +1,4 @@
-import streamlit as st
+from flask import Flask, request, render_template
 import os
 import requests
 from PIL import Image
@@ -12,7 +12,9 @@ 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'  # Replace with your actual API token
+os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'
+
+app = Flask(__name__)
 
 def generate_image_from_text(prompt):
     r = requests.post('https://clipdrop-api.co/text-to-image/v1',
@@ -27,37 +29,23 @@ def generate_image_from_text(prompt):
     else:
         r.raise_for_status()
 
-def resize_image(image_bytes, max_size=(640, 640)):
-    # Open the image from bytes
+def resize_image(image_bytes, max_size=(320, 320)):
     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:
@@ -66,43 +54,30 @@ 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:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
+        "tencentarc/gfpgan:your_model_version",
         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
 
-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)
+@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')
 
 if __name__ == "__main__":
-    main()
+    app.run(port=8501, debug=True)