app.py

import gradio as gr
import numpy as np
import random

import spaces  # [uncomment to use ZeroGPU]
from diffusers import DiffusionPipeline
import torch


model_ids = [
    "Prgckwb/trpfrog-sd3.5-large",
    "Prgckwb/trpfrog-diffusion",
]

if torch.cuda.is_available():
    torch_dtype = torch.float16
    device = "cuda"
else:
    torch_dtype = torch.float32
    device = "cpu"

pipelines = {
    model_id: DiffusionPipeline.from_pretrained(
        model_id, torch_dtype=torch_dtype
    )
    for model_id in model_ids
}


@spaces.GPU()
def inference(
    model_id: str,
    prompt: str,
    width: int,
    height: int,
    progress=gr.Progress(track_tqdm=True),
):
    pipe = pipelines[model_id].to(device)

    image = pipe(
        prompt=prompt,
        width=width,
        height=height,
    ).images[0]

    return image


if __name__ == "__main__":
    theme = gr.themes.Ocean()

    demo = gr.Interface(
        fn=inference,
        inputs=[
            gr.Dropdown(label="Model", choices=model_ids, value=model_ids[0]),
            gr.Textbox(label="Prompt", placeholder="an icon of trpfrog"),
            gr.Slider(label="Width", minimum=64, maximum=1024, step=64, value=1024),
            gr.Slider(label="Height", minimum=64, maximum=1024, step=64, value=1024),
        ],
        outputs=[
            gr.Image(label="Output"),
        ],
        theme=theme,
    )
    demo.queue().launch()