Update app.py

app.py

@@ -1,42 +1,50 @@
 import gradio as gr
-from random import randint
-from all_models import models
-from externalmod import gr_Interface_load, randomize_seed
+import torch
 import asyncio
 import os
+from random import randint
 from threading import RLock
 from pathlib import Path
+from all_models import models
+from externalmod import gr_Interface_load, randomize_seed
 
 # Create a lock for thread safety
 lock = RLock()
 
-# Load Hugging Face token from environment variable (if available)
+# Load Hugging Face token from environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Function to load models
+# Function to load models with optimized settings
 def load_fn(models):
     global models_load
     models_load = {}
+
     for model in models:
         if model not in models_load:
             try:
                 print(f"Loading model: {model}")
-                m = gr_Interface_load(f'models/{model}', hf_token=HF_TOKEN)
+                m = gr_Interface_load(
+                    f'models/{model}',
+                    hf_token=HF_TOKEN,
+                    torch_dtype=torch.float16  # Reduce memory usage
+                )
+                m.enable_model_cpu_offload()  # Offload to CPU when not in use
                 models_load[model] = m
             except Exception as e:
                 print(f"Error loading model {model}: {e}")
-                models_load[model] = gr.Interface(lambda: None, ['text'], ['image'])
+                models_load[model] = None
 
 print("Loading models...")
 load_fn(models)
 print("Models loaded successfully.")
 
+# Constants
 num_models = 1
 starting_seed = randint(1941, 2024)
 MAX_SEED = 3999999999
-MAX_SEED = int(MAX_SEED)
 inference_timeout = 600
 
+# Update UI components
 def extend_choices(choices):
     return choices[:num_models] + ['NA'] * (num_models - len(choices))
 
@@ -44,58 +52,69 @@ def update_imgbox(choices):
     choices_extended = extend_choices(choices)
     return [gr.Image(None, label=m, visible=(m != 'NA')) for m in choices_extended]
 
-async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
-    if model_str not in models_load:
+# Async inference function
+async def infer(model_str, prompt, seed=1):
+    if model_str not in models_load or models_load[model_str] is None:
+        print(f"Model {model_str} is unavailable.")
         return None
-    
+
     kwargs = {"seed": seed}
     try:
         print(f"Running inference for model: {model_str} with prompt: '{prompt}'")
         result = await asyncio.to_thread(models_load[model_str].fn, prompt=prompt, **kwargs, token=HF_TOKEN)
+
         if result:
             with lock:
                 png_path = "image.png"
                 result.save(png_path)
                 return str(Path(png_path).resolve())
+    except torch.cuda.OutOfMemoryError:
+        print(f"CUDA memory error for {model_str}. Try reducing image size.")
     except Exception as e:
         print(f"Error during inference for {model_str}: {e}")
+
     return None
 
+# Synchronous wrapper
 def gen_fnseed(model_str, prompt, seed=1):
     if model_str == 'NA':
         return None
+
     try:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
-        result = loop.run_until_complete(infer(model_str, prompt, seed, inference_timeout))
+        result = loop.run_until_complete(infer(model_str, prompt, seed))
     except Exception as e:
         print(f"Error generating image for {model_str}: {e}")
         result = None
     finally:
         loop.close()
+
     return result
 
+# Gradio UI
 print("Creating Gradio interface...")
 with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
     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 up to 6 images')
         seed = gr.Slider(label="Seed (0 to MAX)", minimum=0, maximum=MAX_SEED, value=starting_seed)
         seed_rand = gr.Button("Randomize Seed 🎲")
-        
+
         seed_rand.click(randomize_seed, None, [seed], queue=False)
-        
+
         output = [gr.Image(label=m) for m in models[:num_models]]
         current_models = [gr.Textbox(m, visible=False) for m in models[:num_models]]
-        
+
         for m, o in zip(current_models, output):
             gen_button.click(gen_fnseed, inputs=[m, txt_input, seed], outputs=[o], queue=False)
-        
+
         with gr.Accordion('Model selection'):
             model_choice = gr.CheckboxGroup(models, label=f'Choose up to {num_models} models')
             model_choice.change(update_imgbox, model_choice, output)
             model_choice.change(extend_choices, model_choice, current_models)
 
-demo.queue(default_concurrency_limit=50, max_size=100)
+demo.queue(default_concurrency_limit=20, max_size=50)  # Adjusted for better stability
 demo.launch(show_api=False)