app.py

import gradio as gr
import spaces
import torch
import os
from diffusers import DiffusionPipeline


model_name = os.environ.get('MODEL_REPO_ID')
pipe = DiffusionPipeline.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
)

def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, num_samples):
    return pipe(
        prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        width=width,
        height=height,
        num_images_per_prompt=num_samples
    ).images

if torch.cuda.is_available():
    pipe.to('cuda')
    generate = spaces.GPU(generate)


gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="Prompt"),
        gr.Text("", label="Negative Prompt"),
        gr.Number(7, label="Number inference steps"),
        gr.Number(3, label="Guidance scale"),
        gr.Number(512, label="Width"),
        gr.Number(512, label="Height"),
        gr.Number(1, label="# images"),
    ],
    outputs=gr.Gallery(),
).launch()