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-sd3.5-medium',
    'Prgckwb/trpfrog-sdxl',
    '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(model_id, device=device, 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()