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()