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 lock = RLock() # Load Hugging Face token from environment variables HF_TOKEN = os.environ.get("HF_TOKEN") # Function to load models 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) print(f"Successfully loaded: {model}") except Exception as error: print(f"Error loading {model}: {error}") m = gr.Interface(lambda: None, ['text'], ['image']) models_load[model] = m # Load 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 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] # Asynchronous inference function async def infer(model_str, prompt, seed=1, timeout=600): if model_str == 'NA': return None 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) 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 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)) # Create Gradio interface print("Creating Gradio interface...") with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo: gr.HTML("