Fix GroupNorm for arbitrary channel counts

liquid_diffusion/model.py

@@ -78,7 +78,11 @@ class AdaLN(nn.Module):
     """Adaptive Layer Normalization: out = norm(x) * (1 + scale(t)) + shift(t)"""
     def __init__(self, dim: int, cond_dim: int):
         super().__init__()
-        self.norm = nn.GroupNorm(num_groups=min(32, dim), num_channels=dim, affine=False)
+        # Find largest valid group count ≤ 32
+        num_groups = min(32, dim)
+        while dim % num_groups != 0:
+            num_groups -= 1
+        self.norm = nn.GroupNorm(num_groups=num_groups, num_channels=dim, affine=False)
         self.proj = nn.Sequential(nn.SiLU(), nn.Linear(cond_dim, dim * 2))
 
     def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
@@ -339,8 +343,11 @@ class LiquidDiffusionUNet(nn.Module):
             self.decoder_blocks.append(stage)
         
         # Output head (initialized to zero for stable start)
+        head_groups = min(32, channels[0])
+        while channels[0] % head_groups != 0:
+            head_groups -= 1
         self.head = nn.Sequential(
-            nn.GroupNorm(min(32, channels[0]), channels[0]),
+            nn.GroupNorm(head_groups, channels[0]),
             nn.SiLU(),
             nn.Conv2d(channels[0], in_channels, 3, padding=1),
         )
@@ -395,19 +402,19 @@ class LiquidDiffusionUNet(nn.Module):
 # =============================================================================
 
 def liquid_diffusion_tiny(**kwargs):
-    """~8M params, 256px, fits ~4GB VRAM."""
+    """~23M params, 256px, fits ~6GB VRAM."""
     return LiquidDiffusionUNet(
         channels=[64, 128, 256], blocks_per_stage=[2, 2, 4],
         t_dim=256, expand_ratio=2.0, kernel_size=7, **kwargs)
 
 def liquid_diffusion_small(**kwargs):
-    """~25M params, 256px, fits ~8GB VRAM."""
+    """~69M params, 256px, fits ~10GB VRAM."""
     return LiquidDiffusionUNet(
         channels=[96, 192, 384], blocks_per_stage=[2, 3, 6],
         t_dim=384, expand_ratio=2.0, kernel_size=7, **kwargs)
 
 def liquid_diffusion_base(**kwargs):
-    """~65M params, 512px, fits ~14GB VRAM."""
+    """~154M params, 512px, fits ~16GB VRAM."""
     return LiquidDiffusionUNet(
         channels=[128, 256, 512], blocks_per_stage=[2, 4, 8],
         t_dim=512, expand_ratio=2.0, kernel_size=7, **kwargs)