app.py

# app.py
import torch
import gradio as gr
from diffusers import DiffusionPipeline

print("Loading pipeline...")

pipe = DiffusionPipeline.from_pretrained(
    "GGPENG/StyleDiffusion",
    torch_dtype=torch.bfloat16,
)

device = "cuda" if torch.cuda.is_available() else "cpu"

pipe.to("cuda")

pipe.unet.load_attn_procs(
"./pytorch_custom_diffusion_weights.bin"
)

def generate(prompt, steps, guidance):

    image = pipe(
        prompt,
        num_inference_steps=steps,
        guidance_scale=guidance,
        eta=1
    ).images[0]

    return image


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(
            label="Prompt",
            value="A <new1> reference. New Year image with a rabbit as the main element"
        ),
        gr.Slider(10, 320, value=100, label="Steps"),
        gr.Slider(1, 18, value=6, label="Guidance"),
    ],
    outputs=gr.Image(),
    title="Fine-tuning style diffusion Demo"
)

demo.launch()