Update app.py

app.py

@@ -1,145 +1,93 @@
 import gradio as gr
+import asyncio
+import os
 from random import randint
+from threading import RLock
 from all_models import models
-
 from externalmod import gr_Interface_load, randomize_seed
 
-import asyncio
-import os
-from threading import RLock
+# Lock to prevent concurrent access issues
 lock = RLock()
-HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None # If private or gated models aren't used, ENV setting is unnecessary.
+HF_TOKEN = os.getenv("HF_TOKEN", None)  # Hugging Face token for private models
 
-
-def load_fn(models):
+# Load models
+def load_models(models):
     global models_load
     models_load = {}
     
     for model in models:
-        if model not in models_load.keys():
+        if model not in models_load:
             try:
-                m = gr_Interface_load(f'models/{model}', hf_token=HF_TOKEN)
+                models_load[model] = gr_Interface_load(f'models/{model}', hf_token=HF_TOKEN)
             except Exception as error:
-                print(error)
-                m = gr.Interface(lambda: None, ['text'], ['image'])
-            models_load.update({model: m})
-
-
-load_fn(models)
-
-
-num_models = 9
-
-default_models = models[:num_models]
-inference_timeout = 600
-MAX_SEED=666666666
-starting_seed = randint(666666000, 666666666)
-
-def extend_choices(choices):
-    return choices[:num_models] + (num_models - len(choices[:num_models])) * ['NA']
-
-
-def update_imgbox(choices):
-    choices_plus = extend_choices(choices[:num_models])
-    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):
-    from pathlib import Path
-    kwargs = {}
-    noise = ""
-    kwargs["seed"] = seed
-    task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn,
-                               prompt=f'{prompt} {noise}', **kwargs, token=HF_TOKEN))
-    await asyncio.sleep(0)
+                print(f"Error loading model {model}: {error}")
+                models_load[model] = None  # Handle failed model load
+
+load_models(models)
+
+# Constants
+NUM_MODELS = 9
+DEFAULT_MODELS = models[:NUM_MODELS]
+INFERENCE_TIMEOUT = 600
+MAX_SEED = 666666666
+starting_seed = randint(666666000, MAX_SEED)
+
+# Async inference function
+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:
+        return "https://huggingface.co/spaces/Yntec/ToyWorld/resolve/main/error.png"
+    
     try:
-        result = await asyncio.wait_for(task, timeout=timeout)
-    except (Exception, asyncio.TimeoutError) as e:
-        print(e)
-        print(f"Task timed out: {model_str}")
-        if not task.done(): task.cancel()
-        result = None
-    if task.done() and result is not None:
-        with lock:
-            png_path = "image.png"
-            result.save(png_path)
-            image = str(Path(png_path).resolve())
-            return image
-    return None
-
-
-
-
-def gen_fnseed(model_str, prompt, seed=1):
+        result = await asyncio.wait_for(
+            asyncio.to_thread(models_load[model_str].fn, prompt=prompt, seed=seed, token=HF_TOKEN), 
+            timeout=timeout
+        )
+        return result if result else "https://huggingface.co/spaces/Yntec/ToyWorld/resolve/main/error.png"
+    except asyncio.TimeoutError:
+        print(f"Timeout error: {model_str}")
+    except Exception as e:
+        print(f"Error in inference: {e}")
+    return "https://huggingface.co/spaces/Yntec/ToyWorld/resolve/main/error.png"
+
+# Synchronous wrapper
+def generate_image(model_str, prompt, seed=1):
     if model_str == 'NA':
         return None
+    
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
     try:
-        loop = asyncio.new_event_loop()
-        result = loop.run_until_complete(infer(model_str, prompt, seed, inference_timeout))
-    except (Exception, asyncio.CancelledError) as e:
-        print(e)
-        print(f"Task aborted: {model_str}")
-        result = None
-        with lock:
-            image = "https://huggingface.co/spaces/Yntec/ToyWorld/resolve/main/error.png"
-        result = image
+        return loop.run_until_complete(infer(model_str, prompt, seed))
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        return "https://huggingface.co/spaces/Yntec/ToyWorld/resolve/main/error.png"
     finally:
         loop.close()
-    return result
 
+# Gradio UI
 with gr.Blocks(theme='Yntec/HaleyCH_Theme_craiyon') as demo:
-    gr.HTML(
-    """
-        <div>
-        <p> <center><img src="https://huggingface.co/spaces/Yntec/open-craiyon/resolve/main/open_craiyon.png" style="height:79px; width:367px; margin-top: -22px; margin-bottom: -44px;" span title="Best free ai art image generator open craiyon"></center>
-        </p>
-    """
-)
-    with gr.Tab('🖍️ AI models drawing images from any prompt! 🖍️'): 
-        with gr.Row():
-            txt_input = gr.Textbox(label='Your prompt:', lines=4, scale=3)
-            gen_button = gr.Button('Draw it! 🖍️', scale=1)
-        with gr.Row():
-            seed = gr.Slider(label="Use a seed to replicate the same image later (maximum 666666666)", minimum=0, maximum=MAX_SEED, step=1, value=starting_seed, scale=3)
-            seed_rand = gr.Button("Randomize Seed 🎲", size="sm", variant="secondary", scale=1)    
-        seed_rand.click(randomize_seed, None, [seed], queue=False)
-        #stop_button = gr.Button('Stop', variant = 'secondary', interactive = False)
-               
-        gen_button.click(lambda s: gr.update(interactive = True), None)
-        gr.HTML(
-        """
-            <div style="text-align: center; max-width: 1200px; margin: 0 auto;">
-              <div>
-                <body>
-                <div class="center"><p style="margin-bottom: 10px; color: #5b6272;">Scroll down to see more images and select models.</p>
-                </div>
-                </body>
-              </div>
-            </div>
-        """
-               )
-        with gr.Row():
-            output = [gr.Image(min_width=480) for m in default_models]
-            #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_event = gr.on(triggers=[gen_button.click, txt_input.submit], fn=gen_fnseed,
-                            inputs=[m, txt_input, seed], outputs=[o], concurrency_limit=None, queue=False)
-                #stop_button.click(lambda s: gr.update(interactive = False), None, stop_button, cancels = [gen_event])
-        with gr.Accordion('Model selection'):
-            model_choice = gr.CheckboxGroup(models, label = 'Untick the models you will not be using', value=default_models, interactive=True)
-            #model_choice = gr.CheckboxGroup(models, label = f'Choose up to {num_models} different models from the 2 available! Untick them to only use one!', value = default_models, multiselect = True, max_choices = num_models, interactive = True, filterable = False)
-            model_choice.change(update_imgbox, model_choice, output)
-            model_choice.change(extend_choices, model_choice, current_models)
-        with gr.Row():
-            gr.HTML(
-    """
-        <div class="footer">
-        <p> For more than a hundred times more models (that's not a typo) check out <a href="https://huggingface.co/spaces/Yntec/ToyWorld">Toy World</a>!</a>
-        </p>
-    """
-)
+    gr.HTML("""<center><img src='https://huggingface.co/spaces/Yntec/open-craiyon/resolve/main/open_craiyon.png' height='79'></center>""")
+    
+    with gr.Tab('🖍️ AI Image Generator 🖍️'):
+        txt_input = gr.Textbox(label='Enter your prompt:', lines=4)
+        gen_button = gr.Button('Generate Image 🖍️')
+        
+        seed_slider = gr.Slider(label="Seed (for reproducibility)", minimum=0, maximum=MAX_SEED, step=1, value=starting_seed)
+        random_seed_btn = gr.Button("Randomize Seed 🎲")
+        
+        random_seed_btn.click(randomize_seed, None, [seed_slider])
+        
+        output_images = [gr.Image(label=m) for m in DEFAULT_MODELS]
+        model_inputs = [gr.Textbox(m, visible=False) for m in DEFAULT_MODELS]
+        
+        for model, img_output in zip(model_inputs, output_images):
+            gen_button.click(generate_image, inputs=[model, txt_input, seed_slider], outputs=[img_output])
+        
+        gr.Accordion("Model Selection", open=False):
+            model_choice = gr.CheckboxGroup(models, label="Select Models", value=DEFAULT_MODELS)
+            model_choice.change(lambda selected: [gr.Image(visible=m in selected) for m in models], inputs=[model_choice], outputs=[output_images])
+    
+    gr.HTML("""<p>Check out more models at <a href='https://huggingface.co/spaces/Yntec/ToyWorld'>Toy World</a>!</p>""")
 
 demo.queue(default_concurrency_limit=200, max_size=200)
-demo.launch(show_api=False, max_threads=400)
- 
+demo.launch(show_api=False, max_threads=400)