Create app.py

app.py

@@ -0,0 +1,226 @@
+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
+
+# 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 (Cópia do seu treino)
+# ───────────────────────────────────────────────
+@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"  # ← Substitua se o nome exato do arquivo no seu repo for outro (ex: checkpoint_epoch100.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")
+
+# Trata se você salvou apenas o state_dict ou o dicionário inteiro de treino
+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)
+    
+    # Passa pelos 1000 passos e atualiza a barra na tela do usuário
+    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)
+
+    # Converte o Tensor para uma lista de imagens PIL para o Gradio
+    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 1.9M (Unconditional)")
+    gr.Markdown("Gerador de 'Fakemons' criado do zero e treinado inteiramente em uma CPU (Ryzen 5 5600G) por **AxionLab-Co**. Como é um modelo incondicional de Difusão, ele inventa criaturas baseadas em ruído puro usando 1000 passos (DDPM).")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            num_slider = gr.Slider(minimum=1, maximum=4, step=1, value=2, label="Quantidade de Fakemons", info="Mais imagens demoram mais tempo na CPU.")
+            gen_btn = gr.Button("Gerar Fakemons! 🚀", variant="primary")
+            gr.Markdown("*Nota: O servidor gratuito do Hugging Face roda em CPU. Gerar imagens pode levar de 30 a 60 segundos.*")
+            
+        with gr.Column(scale=2):
+            gallery = gr.Gallery(label="Fakemons Gerados", 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()