Update app.py

app.py

@@ -5,13 +5,21 @@ 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 (Cópia do seu treino)
+# 1. Configuração e Arquitetura da U-Net
 # ───────────────────────────────────────────────
 @dataclass
 class Config:
@@ -153,7 +161,7 @@ class DDPMScheduler:
 # 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)
+FILENAME = "model.pt"
 
 print("Baixando e carregando o modelo...")
 model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
@@ -162,7 +170,6 @@ cfg = Config()
 model = StudentUNet(cfg)
 ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
 
-# 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:
@@ -180,7 +187,6 @@ 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)
@@ -197,7 +203,6 @@ def generate_fakemons(num_images, progress=gr.Progress()):
         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()
@@ -205,6 +210,7 @@ def generate_fakemons(num_images, progress=gr.Progress()):
     images = [Image.fromarray(img) for img in x]
     return images
 
+
 # ───────────────────────────────────────────────
 # 4. Interface Web (Gradio)
 # ───────────────────────────────────────────────