Upload scheduler.py with huggingface_hub

scheduler.py

@@ -0,0 +1,123 @@
+import torch
+import torch.nn.functional as F
+import math
+
+
+"""
+  This scheduler has 3 main responsibilities:
+
+  1. Setup (init) - Pre-compute noise schedule
+  2. Training (q_sample) - Add noise to images
+  3. Generation (p_sample_text + sample_text) - Remove noise
+  step-by-step
+
+"""
+
+
+class SimpleDDPMScheduler:
+    def __init__(self, num_timesteps=1000, beta_start=0.0001, beta_end=0.02):
+        self.num_timesteps = num_timesteps
+
+        # Linear beta schedule - can replace with cosine
+        self.betas = torch.linspace(beta_start, beta_end, num_timesteps)
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(
+            self.alphas, dim=0
+        )  # cumulative product - lets us jump to any timestep immediately.
+        self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)
+
+        # Calculations for diffusion q(x_t | x_{t-1}) and others (pre-compute for efficiency)
+        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
+
+        # Calculations for posterior q(x_{t-1} | x_t, x_0)
+        # This tells us how much randomness is appropriate at this step.
+        # Removing this would lead to mode-seeking behavior (and poor sample quality).
+        self.posterior_variance = (
+            self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+
+    def q_sample(self, x_start, t, noise=None):
+        """Add noise to the clean images according to the noise schedule.
+
+        So we can have examples at any timestep in the forward process."""
+        # Generate original noise
+        if noise is None:
+            noise = torch.randn_like(x_start)
+
+        sqrt_alphas_cumprod_t = extract(self.sqrt_alphas_cumprod, t, x_start.shape)
+        sqrt_one_minus_alphas_cumprod_t = extract(
+            self.sqrt_one_minus_alphas_cumprod, t, x_start.shape
+        )
+
+        return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
+
+    def p_sample_text(self, model, x, t, text_embeddings, guidance_scale=1.0):
+        """Sample x_{t-1} from x_t using the model with text conditioning and CFG.
+
+        Args:
+            model: The diffusion model
+            x: Current noisy image
+            t: Current timestep
+            text_embeddings: Text embeddings for conditioning
+            guidance_scale: Classifier-free guidance scale (1.0 = no guidance, higher = stronger)
+        """
+        # Get model prediction with text conditioning
+        predicted_noise = model(x, t, text_embeddings)
+
+        # Apply classifier-free guidance if scale > 1.0
+        if guidance_scale > 1.0:
+            # Also get unconditional prediction (with zero text embeddings)
+            uncond_embeddings = torch.zeros_like(text_embeddings)
+            uncond_noise = model(x, t, uncond_embeddings)
+
+            # Amplify the difference between conditional and unconditional
+            predicted_noise = uncond_noise + guidance_scale * (predicted_noise - uncond_noise)
+
+        # Get coefficients
+        betas_t = extract(self.betas, t, x.shape)
+        sqrt_one_minus_alphas_cumprod_t = extract(
+            self.sqrt_one_minus_alphas_cumprod, t, x.shape
+        )
+        sqrt_recip_alphas_t = extract(1.0 / torch.sqrt(self.alphas), t, x.shape)
+
+        # Compute x_{t-1}
+        model_mean = sqrt_recip_alphas_t * (
+            x - betas_t * predicted_noise / sqrt_one_minus_alphas_cumprod_t
+        )
+
+        if t[0] == 0:
+            return model_mean
+        else:
+            posterior_variance_t = extract(self.posterior_variance, t, x.shape)
+            noise = torch.randn_like(x)
+            return model_mean + torch.sqrt(posterior_variance_t) * noise
+
+    def sample_text(self, model, shape, text_embeddings, device="cuda", guidance_scale=1.0):
+        """Generate samples using DDPM sampling with text conditioning and CFG.
+
+        Args:
+            model: The diffusion model
+            shape: Output shape (B, C, H, W)
+            text_embeddings: Text embeddings for conditioning
+            device: Device to use
+            guidance_scale: Classifier-free guidance scale (1.0 = no guidance, 3.0-7.0 typical)
+        """
+        b = shape[0]
+        img = torch.randn(shape, device=device)
+
+        for i in reversed(range(0, self.num_timesteps)):
+            t = torch.full((b,), i, device=device, dtype=torch.long)
+            img = self.p_sample_text(model, img, t, text_embeddings, guidance_scale)
+
+            # Clamp to prevent explosion
+            img = torch.clamp(img, -2.0, 2.0)
+
+        return img
+
+
+def extract(a, t, x_shape):
+    """Extract coefficients from a based on t and reshape to match x_shape."""
+    batch_size = t.shape[0]
+    out = a.gather(-1, t.cpu())
+    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)