Update app.py

app.py

@@ -1,15 +1,17 @@
 import gradio as gr
 from random import randint
 from all_models import models
+
 from externalmod import gr_Interface_load, randomize_seed
+
 import asyncio
 import os
 from threading import RLock
 
-# Create a lock for thread safety
+# Create a lock to ensure thread safety when accessing shared resources
 lock = RLock()
 
-# Load Hugging Face token from environment variables
+# Load Hugging Face token from environment variable
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 # Function to load models
@@ -19,79 +21,87 @@ def load_fn(models):
     for model in models:
         if model not in models_load:
             try:
-                print(f"Loading model: {model}")
+                print(f"Attempting to load model: {model}")
                 m = gr_Interface_load(f'models/{model}', hf_token=HF_TOKEN)
-                print(f"Successfully loaded: {model}")
+                print(f"Successfully loaded model: {model}")
             except Exception as error:
-                print(f"Error loading {model}: {error}")
+                print(f"Error loading model {model}: {error}")
                 m = gr.Interface(lambda: None, ['text'], ['image'])
             models_load[model] = m
 
-# Load models
+# Load the models
 print("Loading models...")
 load_fn(models)
 print("Models loaded successfully.")
 
-num_models = min(3, len(models))  # Reduce to 3 models to prevent GPU overloading
+num_models = 6
+
+default_models = models[:num_models]
+inference_timeout = 600
+MAX_SEED = 3999999999
 starting_seed = randint(1941, 2024)
+
 print(f"Starting seed: {starting_seed}")
 
-# Extend choices to match num_models
 def extend_choices(choices):
-    extended = choices[:num_models] + (num_models - len(choices[:num_models])) * ['NA']
-    return extended
-
-# Update image boxes based on selected models
-def update_imgbox(choices):
-    choices_extended = extend_choices(choices)
-    return [gr.Image(None, label=m, visible=(m != 'NA')) for m in choices_extended]
+    return choices[:num_models] + ['NA'] * (num_models - len(choices))
 
-# Asynchronous inference function
-async def infer(model_str, prompt, seed=1, timeout=600):
+# Asynchronous function for inference
+async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
     if model_str == 'NA':
         return None
+    
+    print(f"Starting inference for model: {model_str} with prompt: '{prompt}' and seed: {seed}")
+    
     try:
-        print(f"Running inference on {model_str} with prompt: '{prompt}'")
-        task = asyncio.to_thread(models_load[model_str].fn, prompt=prompt, seed=seed, token=HF_TOKEN)
-        result = await asyncio.wait_for(task, timeout=timeout)
+        result = await asyncio.to_thread(models_load[model_str].fn, prompt, seed=seed, token=HF_TOKEN)
         if result:
-            with lock:
-                image_path = "image.png"
-                result.save(image_path)
-                return image_path
-    except Exception as e:
-        print(f"Error in inference for {model_str}: {e}")
+            return result
+    except (Exception, asyncio.TimeoutError) as e:
+        print(f"Error during inference for model {model_str}: {e}")
     return None
 
-# Wrapper function for inference
 def gen_fnseed(model_str, prompt, seed=1):
     if model_str == 'NA':
         return None
-    return asyncio.run(infer(model_str, prompt, seed))
+    
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        result = loop.run_until_complete(infer(model_str, prompt, seed))
+    except Exception as e:
+        print(f"Error during generation for model {model_str}: {e}")
+        result = None
+    finally:
+        loop.close()
+    return result
 
-# Create Gradio interface
+# Creating the Gradio UI
 print("Creating Gradio interface...")
 with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
-    gr.HTML("<center><h1>Compare-3</h1></center>")
-    with gr.Tab('Compare-3'):
+    gr.HTML("<center><h1>Compare-6</h1></center>")
+    
+    with gr.Tab('Compare-6'):
         txt_input = gr.Textbox(label='Your prompt:', lines=4)
-        gen_button = gr.Button('Generate images')
-        seed = gr.Slider(label="Seed (max 3999999999)", minimum=0, maximum=3999999999, step=1, value=starting_seed)
-        seed_rand = gr.Button("Randomize Seed 🎲")
+        gen_button = gr.Button('Generate up to 6 images')
+        
+        with gr.Row():
+            seed = gr.Slider("Seed", 0, MAX_SEED, step=1, value=starting_seed)
+            seed_rand = gr.Button("Randomize Seed 🎲")
+        
         seed_rand.click(randomize_seed, None, [seed])
         
-        output = [gr.Image(label=m) for m in models[:num_models]]
-        current_models = [gr.Textbox(m, visible=False) for m in models[:num_models]]
+        with gr.Row():
+            output = [gr.Image(label=m, min_width=480) for m in default_models]
+            current_models = [gr.Textbox(m, visible=False) for m in default_models]
         
         for m, o in zip(current_models, output):
-            gen_button.click(gen_fnseed, inputs=[m, txt_input, seed], outputs=[o])
+            gen_button.click(fn=gen_fnseed, inputs=[m, txt_input, seed], outputs=[o])
         
         with gr.Accordion('Model selection'):
-            model_choice = gr.CheckboxGroup(models, label=f'Choose up to {num_models} models', value=models[:num_models])
-            model_choice.change(update_imgbox, model_choice, output)
-            model_choice.change(extend_choices, model_choice, current_models)
+            model_choice = gr.CheckboxGroup(models, label=f'Choose up to {num_models} models', value=default_models)
+            model_choice.change(lambda c: extend_choices(c), model_choice, current_models)
 
-# Reduce concurrency to avoid T4 overload
-demo.queue(default_concurrency_limit=50, max_size=100)
 print("Launching Gradio interface...")
-demo.launch(show_api=False, max_threads=50, debug=True)
+demo.queue(default_concurrency_limit=50, max_size=100)
+demo.launch(share=True, max_threads=50)