import gradio as gr
import spaces
import torch
from torchvision.transforms import transforms
from huggingface_hub import hf_hub_download
from PIL import Image
from model import Unet, sample_ddim, T ,EMA

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

model = Unet().to(device)

model_path = hf_hub_download(repo_id='adityachaubey/Anime_Face_DDPM', filename='DDPM_weights.pth')
ema_path = hf_hub_download(repo_id='adityachaubey/Anime_Face_DDPM', filename='Ema_shadow.pth')

model_checkpoint = torch.load(model_path, map_location=device)
model.load_state_dict(model_checkpoint['model_state_dict'])

ema_checkpoint = torch.load(ema_path, map_location=device)

ema = EMA(model)
ema.register()
ema.shadow = ema_checkpoint['ema_shadow']  
ema.apply_shadow()

model.eval()
@spaces.GPU
def generate(ddim_steps=25):
    img = sample_ddim(model, T=T, img_size=64, batch_size=1, 
                      ddim_step=ddim_steps, device=device)
    img = (img.clamp(-1, 1) + 1) / 2      # denormalize to [0,1]
    img = img.squeeze(0).cpu()              # remove batch dim (1,3,64,64) → (3,64,64)
    img = transforms.ToPILImage()(img)      # tensor → PIL
    return img

demo = gr.Interface(
    fn=generate,
    inputs=gr.Slider(10, 100, value=25, step=5, label='DDIM Steps'),
    outputs=gr.Image(label='Generated Anime Face', width=256, height=256),
    title='Anime Face Generator',
    description='DDPM diffusion model trained on 63k anime faces for 40 epochs',
    cache_examples=False
)

demo.launch(share=True)