Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from random import randint
 from all_models import models
+
 from externalmod import gr_Interface_load, randomize_seed
 
 import asyncio
@@ -13,18 +14,16 @@ 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 = None  # Ensure failed models are not stored
+                m = gr.Interface(lambda: None, ['text'], ['image'])
             models_load[model] = m
 
 print("Loading models...")
@@ -46,24 +45,16 @@ 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")
@@ -71,7 +62,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' or models_load.get(model_str) is None:
+    if model_str == 'NA':
         return None
     loop = asyncio.new_event_loop()
     result = loop.run_until_complete(infer(model_str, prompt, seed))
@@ -89,10 +80,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)