app.py

import gradio as gr
import torch
import torch.nn as nn
import math
from dataclasses import dataclass
from huggingface_hub import hf_hub_download
from PIL import Image
import pathlib
import platform

# ───────────────────────────────────────────────
# Correção para carregar no Linux modelos salvos no Windows
# ───────────────────────────────────────────────
if platform.system() == 'Linux':
    pathlib.WindowsPath = pathlib.PosixPath

# Otimizações para CPU no servidor do Hugging Face
torch.set_num_threads(4)
torch.backends.mkldnn.enabled = True

# ───────────────────────────────────────────────
# 1. Configuração e Arquitetura da U-Net
# ───────────────────────────────────────────────
@dataclass
class Config:
    image_size: int = 64
    in_channels: int = 3
    base_channels: int = 64
    channel_mults: tuple = (1, 2, 4)
    num_res_blocks: int = 1
    attention_resolutions: tuple = (16,)
    timesteps: int = 1000
    beta_start: float = 1e-4
    beta_end: float = 0.02
    dtype: torch.dtype = torch.float32

class SinusoidalEmbedding(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
    def forward(self, t: torch.Tensor) -> torch.Tensor:
        half = self.dim // 2
        dtype = t.dtype
        freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device, dtype=dtype) / (half - 1))
        args = t[:, None] * freqs[None]
        return torch.cat([args.sin(), args.cos()], dim=-1)

class ResBlock(nn.Module):
    def __init__(self, in_ch: int, out_ch: int, time_dim: int):
        super().__init__()
        self.norm1 = nn.GroupNorm(8, in_ch)
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.norm2 = nn.GroupNorm(8, out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.time_proj = nn.Linear(time_dim, out_ch)
        self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
        self.act = nn.SiLU()
    def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
        h = self.act(self.norm1(x))
        h = self.conv1(h)
        h = h + self.time_proj(self.act(t_emb))[:, :, None, None]
        h = self.act(self.norm2(h))
        h = self.conv2(h)
        return h + self.skip(x)

class AttentionBlock(nn.Module):
    def __init__(self, channels: int, heads: int = 4):
        super().__init__()
        self.norm = nn.GroupNorm(8, channels)
        self.attn = nn.MultiheadAttention(channels, heads, batch_first=True)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        h = self.norm(x).view(B, C, H * W).transpose(1, 2)
        h, _ = self.attn(h, h, h)
        return x + h.transpose(1, 2).view(B, C, H, W)

class DownBlock(nn.Module):
    def __init__(self, in_ch, out_ch, time_dim, use_attn=False, n_res=1):
        super().__init__()
        self.res = nn.ModuleList([ResBlock(in_ch if i == 0 else out_ch, out_ch, time_dim) for i in range(n_res)])
        self.attn = AttentionBlock(out_ch) if use_attn else nn.Identity()
        self.down = nn.Conv2d(out_ch, out_ch, 3, stride=2, padding=1)
    def forward(self, x, t_emb):
        for r in self.res:
            x = r(x, t_emb)
        x = self.attn(x)
        skip = x
        x = self.down(x)
        return x, skip

class UpBlock(nn.Module):
    def __init__(self, in_ch, out_ch, time_dim, use_attn=False, n_res=1):
        super().__init__()
        self.up = nn.ConvTranspose2d(in_ch, in_ch, 2, stride=2)
        self.res = nn.ModuleList([ResBlock(in_ch + out_ch if i == 0 else out_ch, out_ch, time_dim) for i in range(n_res)])
        self.attn = AttentionBlock(out_ch) if use_attn else nn.Identity()
    def forward(self, x, skip, t_emb):
        x = self.up(x)
        x = torch.cat([x, skip], dim=1)
        for r in self.res:
            x = r(x, t_emb)
        x = self.attn(x)
        return x

class StudentUNet(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        ch, mults, time_dim, n_res, attn_res = cfg.base_channels, cfg.channel_mults, cfg.base_channels * 4, cfg.num_res_blocks, cfg.attention_resolutions
        self.time_embed = nn.Sequential(SinusoidalEmbedding(ch), nn.Linear(ch, time_dim), nn.SiLU(), nn.Linear(time_dim, time_dim))
        self.input_conv = nn.Conv2d(cfg.in_channels, ch, 3, padding=1)
        self.downs = nn.ModuleList()
        in_ch, res = ch, cfg.image_size
        self.skip_channels = []
        for mult in mults:
            out_ch = ch * mult
            self.downs.append(DownBlock(in_ch, out_ch, time_dim, res in attn_res, n_res))
            self.skip_channels.append(out_ch)
            in_ch = out_ch; res //= 2
        self.mid_res1 = ResBlock(in_ch, in_ch, time_dim)
        self.mid_attn = AttentionBlock(in_ch)
        self.mid_res2 = ResBlock(in_ch, in_ch, time_dim)
        self.ups = nn.ModuleList()
        for mult in reversed(mults):
            out_ch = ch * mult
            self.ups.append(UpBlock(in_ch, out_ch, time_dim, res in attn_res, n_res))
            in_ch = out_ch; res *= 2
        self.out_norm = nn.GroupNorm(8, in_ch)
        self.out_conv = nn.Conv2d(in_ch, cfg.in_channels, 3, padding=1)
        self.act = nn.SiLU()

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        target_dtype = self.input_conv.weight.dtype
        x = x.to(target_dtype)
        t_emb = self.time_embed(t.to(target_dtype))
        h = self.input_conv(x)
        skips = []
        for down in self.downs:
            h, skip = down(h, t_emb)
            skips.append(skip)
        h = self.mid_res1(h, t_emb)
        h = self.mid_attn(h)
        h = self.mid_res2(h, t_emb)
        for up, skip in zip(self.ups, reversed(skips)):
            h = up(h, skip, t_emb)
        h = self.act(self.out_norm(h))
        return self.out_conv(h)

class DDPMScheduler:
    def __init__(self, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.T = timesteps
        betas = torch.linspace(beta_start, beta_end, timesteps)
        alphas = 1.0 - betas
        self.alpha_bar = torch.cumprod(alphas, dim=0)

    def predict_x0(self, xt: torch.Tensor, noise_pred: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        ab = self.alpha_bar[t].view(-1, 1, 1, 1).to(xt.dtype)
        return (xt - (1 - ab).sqrt() * noise_pred) / ab.sqrt()


# ───────────────────────────────────────────────
# 2. Carregando o Modelo do seu Repositório
# ───────────────────────────────────────────────
REPO_ID = "AxionLab-Co/PokePixels1-9M"
FILENAME = "model.pt"

print("Baixando e carregando o modelo...")
model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

cfg = Config()
model = StudentUNet(cfg)
ckpt = torch.load(model_path, map_location="cpu", weights_only=False)

if "model_state" in ckpt:
    model.load_state_dict(ckpt["model_state"])
else:
    model.load_state_dict(ckpt)
    
model.eval()
scheduler = DDPMScheduler(cfg.timesteps, cfg.beta_start, cfg.beta_end)


# ───────────────────────────────────────────────
# 3. Lógica de Geração com Barra de Progresso Gradio
# ───────────────────────────────────────────────
@torch.no_grad()
def generate_fakemons(num_images, progress=gr.Progress()):
    device = "cpu"
    x = torch.randn(num_images, 3, cfg.image_size, cfg.image_size, device=device)
    
    for t_val in progress.tqdm(reversed(range(scheduler.T)), total=scheduler.T, desc="Removendo ruído (DDPM)"):
        t = torch.full((num_images,), t_val, device=device, dtype=torch.long)
        noise_pred = model(x, t)
        
        if t_val > 0:
            ab = scheduler.alpha_bar[t_val].to(x.dtype)
            ab_prev = scheduler.alpha_bar[t_val - 1].to(x.dtype)
            beta_t = 1.0 - (ab / ab_prev)
            alpha_t = 1.0 - beta_t
            
            mean = (1.0 / alpha_t.sqrt()) * (x - (beta_t / (1.0 - ab).sqrt()) * noise_pred)
            sigma = beta_t.sqrt()
            x = mean + sigma * torch.randn_like(x)
        else:
            x = scheduler.predict_x0(x, noise_pred, t)

    x = x.clamp(-1, 1)
    x = (x + 1) / 2
    x = (x * 255).to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy()
    
    images = [Image.fromarray(img) for img in x]
    return images


# ───────────────────────────────────────────────
# 4. Interface Web (Gradio)
# ───────────────────────────────────────────────
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# ⚡ PokePixels 9M (Unconditional)")
    gr.Markdown("Fakemon generator created from scratch and trained entirely on a CPU (Ryzen 5 5600G) by AxionLab-Co. As an unconditional Diffusion model, it creates creatures based on pure noise using 1000 steps (DDPM).")
    
    with gr.Row():
        with gr.Column(scale=1):
            num_slider = gr.Slider(minimum=1, maximum=4, step=1, value=2, label="Fakemon Quantity", info="More images take more time on HuggingFace CPU.")
            gen_btn = gr.Button("Generate Fakemons! 🚀", variant="primary")
            gr.Markdown("*Note: Hugging Face's free server runs on CPU. Generating images may take 30 to 60 seconds.*")
            
        with gr.Column(scale=2):
            gallery = gr.Gallery(label="Generated fakemons:", show_label=True, elem_id="gallery", columns=[2], rows=[2], object_fit="contain", height="auto")

    gen_btn.click(fn=generate_fakemons, inputs=num_slider, outputs=gallery)

demo.launch()