Spaces:
Running
Running
File size: 1,729 Bytes
2741722 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import gradio as gr
from all_models import models
from externalmod import gr_Interface_load
import asyncio
import os
from datetime import datetime
# Load the models
HF_TOKEN = os.getenv("HF_TOKEN", None)
def load_models(models):
loaded_models = {}
for model in models:
try:
loaded_models[model] = gr_Interface_load(f'models/{model}', hf_token=HF_TOKEN)
except Exception as e:
print(f"Error loading {model}: {e}")
return loaded_models
models_load = load_models(models)
# Generate image function
async def infer(model_str, prompt, seed=-1):
task = asyncio.create_task(
asyncio.to_thread(models_load[model_str].fn, prompt=prompt, seed=seed, token=HF_TOKEN)
)
await asyncio.sleep(0)
result = await asyncio.wait_for(task, timeout=300)
return result
def generate_image(model_name, prompt, seed):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
return loop.run_until_complete(infer(model_name, prompt, seed))
finally:
loop.close()
# Interface
with gr.Blocks() as demo:
with gr.Column():
model_choice = gr.Dropdown(
choices=models, label="Select Model", value=models[0]
)
prompt_input = gr.Textbox(label="Enter your prompt")
seed_input = gr.Slider(
label="Seed", minimum=-1, maximum=100000, step=1, value=-1
)
generate_button = gr.Button("Generate Image")
output_image = gr.Image(label="Generated Image", show_download_button=True)
generate_button.click(
fn=generate_image,
inputs=[model_choice, prompt_input, seed_input],
outputs=[output_image],
)
demo.launch()
|