Update main.py

main.py

@@ -27,6 +27,18 @@ def generate_image_from_text(prompt):
     else:
         r.raise_for_status()
 
+def resize_image(image_bytes, max_size=(1024, 1024)):
+    # 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'
@@ -81,8 +93,11 @@ def main():
         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(image_bytes)
+        upscaled_image_bytes = upscale_image_stable_diffusion(resized_image_bytes)
         
         st.success("Further upscaling image with GFPGAN...")
         img = further_upscale_image(upscaled_image_bytes)