Update app.py

app.py

@@ -1,12 +1,21 @@
+import os
 import subprocess
 import sys
 
-# Check if Flask is installed
+import gradio as gr
+from random import randint
+from all_models import models
+from flask import Flask, request, send_file
+from io import BytesIO
+from PIL import Image
+
+app = Flask(__name__)
+
+# Install Flask if not already installed
 try:
     import flask
 except ImportError:
     print("Flask is not installed. Installing Flask...")
-    # Install Flask using subprocess
     result = subprocess.run([sys.executable, "-m", "pip", "install", "flask"], capture_output=True, text=True)
     if result.returncode != 0:
         print("Failed to install Flask.")
@@ -14,18 +23,8 @@ except ImportError:
         sys.exit(1)
     else:
         print("Flask installed successfully.")
-    # Try importing Flask again
     import flask
 
-import gradio as gr
-from random import randint
-from all_models import models
-from flask import Flask, request, send_file
-from io import BytesIO
-from PIL import Image
-
-app = Flask(__name__)
-
 def load_model(model_name):
     try:
         # Load model within Gradio context
@@ -58,7 +57,7 @@ def gen_fn(model_str, prompt, negative_prompt=None):
         with gr.Blocks() as block:
             result = model(full_prompt)
         
-        if isinstance(result, str):  # Assuming result might be a file path
+        if isinstance(result, str):
             if os.path.exists(result):
                 return Image.open(result)
             else:
@@ -67,7 +66,7 @@ def gen_fn(model_str, prompt, negative_prompt=None):
         elif isinstance(result, Image.Image):
             return result
         else:
-            print("Result is not an image:", type(result))
+            print("Unexpected result type:", type(result))
             return None
     except Exception as e:
         print("Error generating image:", e)
@@ -85,7 +84,7 @@ def home():
     if prompt:
         # Generate the image
         image = gen_fn(model, prompt, negative_prompt)
-        if isinstance(image, Image.Image):  # Ensure the result is a PIL image
+        if isinstance(image, Image.Image):
             # Save image to BytesIO object
             img_io = BytesIO()
             image.save(img_io, format='PNG')