app.py

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

lock = RLock()
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)
                print(f"Loaded model: {model}")
            except Exception as error:
                print(f"Error loading model {model}: {error}")
                m = None  # Avoid using gr.Interface here
            models_load[model] = m

print("Loading models...")
load_fn(models)
print("Models loaded successfully.")

num_models = 6
default_models = models[:num_models]
inference_timeout = 600
MAX_SEED = 3999999999
starting_seed = randint(1941, 2024)
print(f"Starting seed: {starting_seed}")

def extend_choices(choices):
    return choices[:num_models] + (num_models - len(choices)) * ['NA']

def update_imgbox(choices):
    choices_plus = extend_choices(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
    
    kwargs = {"seed": seed}
    print(f"Starting inference: {model_str} | Prompt: '{prompt}' | Seed: {seed}")

    try:
        result = await asyncio.wait_for(
            asyncio.to_thread(models_load[model_str].fn, prompt=prompt, **kwargs),
            timeout=timeout
        )
        if result:
            save_path = "image.png"
            with lock:
                result.save(save_path)
            return save_path
    except Exception as e:
        print(f"Error during inference: {e}")
    
    return None

def gen_fnseed(model_str, prompt, seed=1):
    if model_str == 'NA':
        return None
    return asyncio.run(infer(model_str, prompt, seed))

print("Creating Gradio interface...")
with gr.Blocks(theme="gradio/soft") as demo:
    gr.HTML("<center><h1>TEXT-IMAGE-USING-MULTIMODELS</h1></center>")

    with gr.Tab():
        txt_input = gr.Textbox(label='Your prompt:', lines=4)
        gen_button = gr.Button('Generate')
        seed = gr.Slider("Seed", minimum=0, maximum=MAX_SEED, step=1, value=starting_seed)
        seed_rand = gr.Button("Randomize Seed 🎲")
        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)
            model_choice.change(extend_choices, model_choice, current_models)

demo.queue(default_concurrency_limit=500, max_size=500)
demo.launch(show_api=False, max_threads=400)