Update app.py

app.py

@@ -3,8 +3,9 @@ from PIL import Image
 import requests
 import io
 import base64
+import threading
 
-# List of models to process
+# List of models
 models = [
     "dreamlike-art/dreamlike-photoreal-2.0",
     "stabilityai/stable-diffusion-xl-base-1.0",
@@ -69,59 +70,75 @@ models = [
     "Yntec/epiCPhotoGasm",
 ]
 
-# Initialize a dictionary to track models and their status
-model_results = {"success": {}, "failed": []}
+# Load models
+loaded_models = []
+available_models = []
 
-def generate_images(prompt):
-    output_images = []
-    failed_models = []
-    for model_name in models:
+for model in models:
+    try:
+        loaded_models.append(gr.load(f'models/{model}'))
+        available_models.append(model)
+    except Exception as e:
+        print(f"Model {model} could not be loaded: {e}")
+
+print(loaded_models)
+print(available_models)
+
+def generate_image(prompt, model_index):
+    try:
+        model = loaded_models[model_index]
+        out_img = model(prompt)
+        print(out_img)
+        return out_img
+    except Exception as e:
+        print(e)
+        return None
+
+def run_all_models(prompt):
+    out_box = []
+    out_html = ""
+
+    for i, model in enumerate(loaded_models):
         try:
-            # Attempt to load and generate an image with the model
-            model = gr.Interface.load(f"models/{model_name}")
-            img_path = model.predict(prompt)
+            out_img = generate_image(prompt, i)
+            if out_img:
+                raw = Image.open(out_img)
+                raw = raw.convert('RGB')
+                out_box.append(raw)
 
-            # Fetch image
-            response = requests.get(img_path, stream=True)
-            if response.status_code == 200:
-                img = Image.open(io.BytesIO(response.content))
-                output_images.append((model_name, img))
-            else:
-                failed_models.append(model_name)
+                # Convert image to base64 for HTML display
+                buffered = io.BytesIO()
+                raw.save(buffered, format="PNG")
+                img_str = base64.b64encode(buffered.getvalue()).decode()
+                img_tag = f"<img src='data:image/png;base64,{img_str}'/>"
+                out_html += f"<div class='img_class'><a href='https://huggingface.co/models/{available_models[i]}'>{available_models[i]}</a><br>{img_tag}</div>"
         except Exception as e:
-            print(f"Error with model {model_name}: {e}")
-            failed_models.append(model_name)
+            print(e)
+            out_html += f"<div class='img_class'>{available_models[i]}: {str(e)}</div>"
 
-    return output_images, failed_models
+    html_out = f"<div class='grid_class'>{out_html}</div>"
+    return out_box, html_out
 
-def app_interface(prompt):
-    output_images, failed_models = generate_images(prompt)
-    images_html = ""
-    for model_name, img in output_images:
-        # Convert images for display in Gradio
-        img_buffer = io.BytesIO()
-        img.save(img_buffer, format="PNG")
-        img_str = base64.b64encode(img_buffer.getvalue()).decode()
-        images_html += f"""
-        <div>
-            <strong>{model_name}</strong>
-            <img src='data:image/png;base64,{img_str}' style="max-width: 200px;" />
-        </div>
-        """
-    
-    failed_html = "<ul>" + "".join(f"<li>{model}</li>" for model in failed_models) + "</ul>"
-    return images_html, f"<h3>Failed Models:</h3>{failed_html}"
+css = """
+.grid_class {
+    display: flex;
+    flex-wrap: wrap;
+    gap: 10px;
+}
+.img_class {
+    min-width: 200px;
+    margin: 5px;
+}
+"""
 
-# Define Gradio UI
-with gr.Blocks() as app:
+with gr.Blocks(css=css, theme="Nymbo/Nymbo_Theme") as app:
     with gr.Row():
         with gr.Column():
-            prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt")
-            generate_button = gr.Button("Generate Images")
-        with gr.Column():
-            output_html = gr.HTML(label="Generated Images")
-            failed_html = gr.HTML(label="Failed Models")
+            inp = gr.Textbox(label="Prompt")
+            btn = gr.Button("Generate Images")
 
-    generate_button.click(app_interface, inputs=[prompt_input], outputs=[output_html, failed_html])
+    out_html = gr.HTML()
+    outp = gr.Gallery()
+    btn.click(run_all_models, [inp], [outp, out_html])
 
 app.launch()