import gradio as gr
from huggingface_hub import hf_hub_download
from diffusers import StableDiffusionPipeline
import torch
import os

print("Downloading model...")
model_path = hf_hub_download(
    repo_id="calcuis/pony",
    filename="blackmagic-q4_k_m.gguf",
    token=os.environ.get("HF_TOKEN")
)
print(f"Model ready at: {model_path}")

print("Loading pipeline...")
pipe = StableDiffusionPipeline.from_single_file(
    model_path,
    torch_dtype=torch.float32,
)
pipe.to("cpu")
print("Pipeline ready!")

def generate(prompt, negative_prompt="bad quality, blurry", steps=5):
    print(f"Generating: {prompt}")
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=int(steps),
        height=512,
        width=512,
    ).images[0]
    print("Done!")
    return image

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", value="a cute cat sitting on a bench"),
        gr.Textbox(label="Negative Prompt", value="bad quality, blurry"),
        gr.Slider(minimum=1, maximum=30, value=5, step=1, label="Steps")
    ],
    outputs=gr.Image(label="Generated Image", type="pil"),
    title="Pony Image Generator API"
)

demo.launch()