Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from random import randint
 from all_models import models
-
 from externalmod import gr_Interface_load, randomize_seed
 
 import asyncio
@@ -14,16 +13,18 @@ HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None
 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)
+                if not hasattr(m, 'predict'):
+                    raise ValueError(f"Model {model} does not have a 'predict' method.")
                 print(f"Loaded model: {model}")
             except Exception as error:
                 print(f"Error loading model {model}: {error}")
-                m = gr.Interface(lambda: None, ['text'], ['image'])
+                m = None  # Ensure failed models are not stored
             models_load[model] = m
 
 print("Loading models...")
@@ -45,16 +46,24 @@ def update_imgbox(choices):
     return [gr.Image(None, label=m, visible=(m != 'NA')) for m in choices_plus]
 
 async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
+    if model_str not in models_load or models_load[model_str] is None:
+        print(f"Model {model_str} is not available.")
+        return None
+    
+    model = models_load[model_str]
     kwargs = {"seed": seed}
+
     print(f"Starting inference: {model_str} | Prompt: '{prompt}' | Seed: {seed}")
-    task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn, prompt=prompt, **kwargs, token=HF_TOKEN))
+    
     try:
+        task = asyncio.create_task(asyncio.to_thread(model.predict, prompt, **kwargs))
         result = await asyncio.wait_for(task, timeout=timeout)
     except Exception as e:
         print(f"Error during inference: {e}")
         if not task.done():
             task.cancel()
         return None
+
     if task.done() and result:
         with lock:
             result.save("image.png")
@@ -62,7 +71,7 @@ async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
     return None
 
 def gen_fnseed(model_str, prompt, seed=1):
-    if model_str == 'NA':
+    if model_str == 'NA' or models_load.get(model_str) is None:
         return None
     loop = asyncio.new_event_loop()
     result = loop.run_until_complete(infer(model_str, prompt, seed))
@@ -80,10 +89,10 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
         seed_rand.click(randomize_seed, None, [seed])
         output = [gr.Image(label=m) 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, [m, txt_input, seed], o)
-        
+
         with gr.Accordion('Model selection'):
             model_choice = gr.CheckboxGroup(models, label=f'Choose up to {num_models} models', value=default_models)
             model_choice.change(update_imgbox, model_choice, output)