import torch
import gradio as gr
from diffusers import DiffusionPipeline

torch.set_num_threads(torch.get_num_threads())
torch.set_float32_matmul_precision("high")

pipe = DiffusionPipeline.from_pretrained(
    "stablediffusionapi/anything-v5",
    torch_dtype=torch.float32
)
pipe = pipe.to("cpu")

def generate(prompt, steps, seed):
    generator = torch.Generator(device="cpu").manual_seed(seed)

    for i, out in enumerate(pipe(
        prompt=prompt,
        num_inference_steps=steps,
        generator=generator,
        callback_steps=1,
        callback=lambda step, t, latents: None
    )):
        yield gr.Progress((i + 1) / steps), out.images[0]

with gr.Blocks() as demo:
    gr.Markdown("## 🌀 Anything-V5 CPU Anime Generator")

    with gr.Row():
        prompt = gr.Textbox(
            label="Prompt",
            value="Astronaut in a jungle, cold color palette, muted colors, detailed, anime style"
        )

    with gr.Row():
        steps = gr.Slider(10, 40, value=25, step=1, label="Steps")
        seed = gr.Number(value=42, precision=0, label="Seed")

    output = gr.Image(type="pil", label="Result")

    btn = gr.Button("Generate")

    btn.click(
        fn=generate,
        inputs=[prompt, steps, seed],
        outputs=output
    )

demo.launch()