Update README.md

README.md

@@ -75,44 +75,59 @@ Limitations:
 
 ---
 
+## THE INITIAL IDEA WAS A STUDENT U-NET FROM A TEACHER U-NET, BUT THIS WAS DISCONTINUED BECAUSE THE TEACHER WAS INITIALIZATED WITH RANDOM WEIGHTS, THAT WOULD KILL THE STUDENT LEARNING
+
 ## 🚀 Usage
 
 ### Generate images
 
-## THE INITIAL IDEA WAS A STUDENT U-NET FROM A TEACHER U-NET, BUT THIS WAS DISCONTINUED BECAUSE THE TEACHER WAS INITIALIZATED WITH RANDOM WEIGHTS, THAT WOULD KILL THE STUDENT LEARNING
 
 
 ```python
 
-import argparse
 import torch
 from pathlib import Path
+from PIL import Image
 
+# ===== CONFIG =====
+CHECKPOINT = "model.pt"
+N_IMAGES = 8
+STEPS = 50
+SEED = 42
+OUT = "generated.png"
+
+# ===== IMPORT MODEL =====
 from train import StudentUNet, DDPMScheduler, Config
 
+# ===== LOAD =====
+torch.manual_seed(SEED)
 
-# ----------------------------
-# Sampling (com controle de steps)
-# ----------------------------
-@torch.no_grad()
-def generate_samples(model, scheduler, n=4, steps=50, device="cpu", dtype=torch.float32):
-    model.eval()
+ckpt = torch.load(CHECKPOINT, map_location="cpu")
+cfg = ckpt.get("config", Config())
+
+model = StudentUNet(cfg)
+model.load_state_dict(ckpt["model_state"])
+model.eval()
 
-    x = torch.randn(n, 3, cfg.image_size, cfg.image_size, device=device, dtype=dtype)
+scheduler = DDPMScheduler(cfg.timesteps, cfg.beta_start, cfg.beta_end)
+
+# ===== SAMPLING =====
+@torch.no_grad()
+def sample(model, scheduler, n, steps):
+    x = torch.randn(n, 3, cfg.image_size, cfg.image_size)
 
-    # 🔥 Usa menos steps → muito mais rápido
     step_size = scheduler.T // steps
-    timesteps = list(range(0, scheduler.T, step_size))
-    timesteps = list(reversed(timesteps))
+    timesteps = list(range(0, scheduler.T, step_size))[::-1]
 
     for t_val in timesteps:
-        t = torch.full((n,), t_val, device=device, dtype=torch.long)
+        t = torch.full((n,), t_val, dtype=torch.long)
+
         noise_pred = model(x, t)
 
         if t_val > 0:
-            ab = scheduler.alpha_bar[t_val].to(x.dtype)
+            ab = scheduler.alpha_bar[t_val]
             prev_t = max(t_val - step_size, 0)
-            ab_prev = scheduler.alpha_bar[prev_t].to(x.dtype)
+            ab_prev = scheduler.alpha_bar[prev_t]
 
             beta_t = 1.0 - (ab / ab_prev)
             alpha_t = 1.0 - beta_t
@@ -121,83 +136,28 @@ def generate_samples(model, scheduler, n=4, steps=50, device="cpu", dtype=torch.
                 x - (beta_t / (1.0 - ab).sqrt()) * noise_pred
             )
 
-            sigma = beta_t.sqrt()
-            x = mean + sigma * torch.randn_like(x)
+            x = mean + beta_t.sqrt() * torch.randn_like(x)
         else:
             x = scheduler.predict_x0(x, noise_pred, t)
 
-    model.train()
     return x.clamp(-1, 1)
 
+samples = sample(model, scheduler, N_IMAGES, STEPS)
 
-# ----------------------------
-# Save
-# ----------------------------
-def save_samples(samples, path: Path):
-    samples = (samples + 1) / 2
-    samples = (samples * 255).byte().permute(0, 2, 3, 1).cpu().numpy()
-
-    from PIL import Image
-
-    n = len(samples)
-    w = samples.shape[1]
+# ===== SAVE =====
+samples = (samples + 1) / 2
+samples = (samples * 255).byte().permute(0, 2, 3, 1).numpy()
 
-    grid = Image.new("RGB", (n * w, w))
+grid = Image.new("RGB", (cfg.image_size * N_IMAGES, cfg.image_size))
 
-    for i, s in enumerate(samples):
-        grid.paste(Image.fromarray(s), (i * w, 0))
+for i, img in enumerate(samples):
+    grid.paste(Image.fromarray(img), (i * cfg.image_size, 0))
 
-    path.parent.mkdir(parents=True, exist_ok=True)
-    grid.save(path)
+grid.save(OUT)
 
+print(f"✅ Saved to {OUT}")
 
-# ----------------------------
-# Main
-# ----------------------------
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--checkpoint", type=str, required=True)
-    parser.add_argument("--n_images", type=int, default=8)
-    parser.add_argument("--steps", type=int, default=50)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--out", type=str, default="outputs/generated.png")
 
-    args = parser.parse_args()
-
-    # Seed
-    torch.manual_seed(args.seed)
-
-    # Load checkpoint
-    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
-    global cfg
-    cfg = ckpt.get("config", Config())
-
-    # Model
-    model = StudentUNet(cfg)
-    model.load_state_dict(ckpt["model_state"])
-    model.eval()
-
-    # Scheduler
-    scheduler = DDPMScheduler(cfg.timesteps, cfg.beta_start, cfg.beta_end)
-
-    print(f"\n🚀 Generating {args.n_images} images")
-    print(f"⚙️ Steps: {args.steps} | Seed: {args.seed}")
-
-    samples = generate_samples(
-        model,
-        scheduler,
-        n=args.n_images,
-        steps=args.steps,
-        dtype=cfg.dtype
-    )
-
-    save_samples(samples, Path(args.out))
-
-    print(f"✅ Saved to: {args.out}")
-
-
-if __name__ == "__main__":
-    main()
 ```
 
 ```bash
@@ -206,6 +166,7 @@ python generate.py \
   --n_images 8 \
   --steps 50 \
   --seed 42
+```
 
 📁 Output