Upload inference.py

inference.py

@@ -56,6 +56,8 @@ class ShellDConfig:
     beta_end: float = 0.02
     model_name: str = "ShellD"
     dropout: float = 0.0               # match train.py; 0 = no dropout (inference)
+    kl_weight: float = 0.1             # VAE KL weight (not used in inference)
+    cfg_dropout_prob: float = 0.15     # CFG dropout prob (not used in inference)
 
     @classmethod
     def from_json(cls, path: str) -> "ShellDConfig":
@@ -142,7 +144,7 @@ class Decoder(nn.Module):
                 nn.Sequential(
                     ResidualBlock(ch, out_ch),
                     ResidualBlock(out_ch, out_ch),
-                    nn.Upsample(scale_factor=2, mode="nearest"),
+                    nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
                 )
             )
             ch = out_ch
@@ -151,7 +153,7 @@ class Decoder(nn.Module):
         self.out_conv = nn.Sequential(
             ResidualBlock(ch, ch),
             nn.Conv2d(ch, 3, 3, padding=1),
-            nn.Sigmoid(),
+            nn.Tanh(),
         )
 
     def forward(self, z: torch.Tensor) -> torch.Tensor:
@@ -329,6 +331,8 @@ class ShellDModel(nn.Module):
         self.vae = VAE(cfg)
         self.dit = DiT(cfg)
         self.text_encoder = None  # loaded separately
+        # Null text embedding for classifier-free guidance (loaded from checkpoint)
+        self.null_text_embed: Optional[torch.Tensor] = None
 
     def encode_text(self, prompts: List[str], device: torch.device) -> torch.Tensor:
         assert self.text_encoder is not None, "Text encoder not loaded"
@@ -396,9 +400,13 @@ class DiffusionSchedule:
         # Start from random noise
         z = torch.randn(B, cfg.latent_dim, H, W, device=device)
 
-        # For classifier-free guidance, we need an unconditional embedding (zeros)
+        # For classifier-free guidance, we need the learned null embedding
         if cfg_scale != 1.0:
-            uncond_emb = torch.zeros_like(text_emb)
+            if model.null_text_embed is not None:
+                uncond_emb = model.null_text_embed.to(device).expand(B, -1, -1)
+            else:
+                # Fallback: zeros (works if model was trained with zero-dropped text)
+                uncond_emb = torch.zeros_like(text_emb)
 
         # Time step resampling for faster inference (DDPM-style, evenly spaced)
         step_indices = torch.linspace(0, cfg.num_timesteps - 1, num_steps, device=device, dtype=torch.long)
@@ -535,6 +543,13 @@ class ShellDInference:
         self.model.vae.load_state_dict(vae_sd)
         self.model.dit.load_state_dict(dit_sd)
 
+        # Restore null text embedding for CFG
+        if "null_text_embed" in sd:
+            self.model.null_text_embed = sd["null_text_embed"].to(self.device)
+            print("Null text embedding loaded for CFG.")
+        else:
+            print("No null_text_embed in checkpoint — CFG will use zeros (may degrade quality).")
+
         if txt_sd and self.model.text_encoder is not None:
             self.model.text_encoder.load_state_dict(txt_sd)
             print("Text encoder weights loaded from safetensors.")
@@ -579,12 +594,12 @@ class ShellDInference:
             seed=seed,
         )
 
-        # Decode latent to image
-        img_tensor = self.model.vae.decode(z)  # [1, 3, 256, 256], values in [0, 1]
+        # Decode latent to image (VAE decoder outputs [-1, 1] via Tanh)
+        img_tensor = self.model.vae.decode(z)  # [1, 3, 256, 256], values in [-1, 1]
 
-        # Convert to PIL
+        # Convert to PIL: [-1,1] → [0,255]
         img_np = img_tensor[0].permute(1, 2, 0).cpu().numpy()  # [256, 256, 3]
-        img_np = (img_np * 255).clip(0, 255).astype(np.uint8)
+        img_np = ((img_np + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
         img = Image.fromarray(img_np)
 
         if output_size is not None:
@@ -594,9 +609,9 @@ class ShellDInference:
 
     def _decode_latent(self, z: torch.Tensor) -> Image.Image:
         """Decode a latent tensor to a PIL Image (helper for streaming)."""
-        img_tensor = self.model.vae.decode(z)  # [B, 3, H, W] in [0,1]
+        img_tensor = self.model.vae.decode(z)  # [B, 3, H, W] in [-1,1]
         img_np = img_tensor[0].permute(1, 2, 0).cpu().numpy()
-        img_np = (img_np * 255).clip(0, 255).astype(np.uint8)
+        img_np = ((img_np + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
         return Image.fromarray(img_np)
 
     @torch.no_grad()
@@ -625,7 +640,13 @@ class ShellDInference:
 
         # --- Encode prompt ---
         text_emb = self.model.encode_text([prompt], device)  # [1, 1, 384]
-        uncond_emb = torch.zeros_like(text_emb) if cfg_scale != 1.0 else None
+        if cfg_scale != 1.0:
+            if self.model.null_text_embed is not None:
+                uncond_emb = self.model.null_text_embed.to(device).expand(1, -1, -1)
+            else:
+                uncond_emb = torch.zeros_like(text_emb)
+        else:
+            uncond_emb = None
 
         # --- Start from pure noise ---
         H = W = cfg.image_size // (2 ** cfg.ae_num_blocks)
@@ -710,12 +731,12 @@ class ShellDInference:
         )
 
         # Decode
-        img_tensor = self.model.vae.decode(z)  # [B, 3, 256, 256]
+        img_tensor = self.model.vae.decode(z)  # [B, 3, 256, 256] in [-1,1]
 
         images = []
         for i in range(B):
             img_np = img_tensor[i].permute(1, 2, 0).cpu().numpy()
-            img_np = (img_np * 255).clip(0, 255).astype(np.uint8)
+            img_np = ((img_np + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
             images.append(Image.fromarray(img_np))
 
         return images