Update README.md

README.md

@@ -79,12 +79,134 @@ Limitations:
 
 ### 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 train import StudentUNet, DDPMScheduler, Config
+
+
+# ----------------------------
+# Sampling (com controle de steps)
+# ----------------------------
+@torch.no_grad()
+def generate_samples(model, scheduler, n=4, steps=50, device="cpu", dtype=torch.float32):
+    model.eval()
+
+    x = torch.randn(n, 3, cfg.image_size, cfg.image_size, device=device, dtype=dtype)
+
+    # 🔥 Usa menos steps → muito mais rápido
+    step_size = scheduler.T // steps
+    timesteps = list(range(0, scheduler.T, step_size))
+    timesteps = list(reversed(timesteps))
+
+    for t_val in timesteps:
+        t = torch.full((n,), 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)
+            prev_t = max(t_val - step_size, 0)
+            ab_prev = scheduler.alpha_bar[prev_t].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)
+
+    model.train()
+    return x.clamp(-1, 1)
+
+
+# ----------------------------
+# 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]
+
+    grid = Image.new("RGB", (n * w, w))
+
+    for i, s in enumerate(samples):
+        grid.paste(Image.fromarray(s), (i * w, 0))
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    grid.save(path)
+
+
+# ----------------------------
+# 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
 python generate.py \
   --checkpoint model.pt \
   --n_images 8 \
   --steps 50 \
   --seed 42
+
 📁 Output
 
 Generated images are saved as a horizontal grid: