import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
import spaces

if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"
    
model_id = "nikkijiang/sd-pokemon-generator" 


pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32,safety_checker=None).to(device)

@spaces.GPU
def generate(prompt, negative_prompt):
    image = pipe(
        prompt='an official style pokemon :'+prompt,
        negative_prompt=negative_prompt,
        height=256,
        width=256,
        num_inference_steps=100,
        guidance_scale=7.5
    ).images[0]
    return image

gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="description", placeholder="e.g., lovely cat, cute, shining"),
        gr.Textbox(label="negative_prompt", placeholder="e.g., simple"),
    ],
    outputs=gr.Image(type="pil"),
    title="Custom Pokémon Generator (beta)",
    description="Enter a prompt to generate your own Pokémon-style character.",
).launch()