1{
2  "op": "replace",
3  "filename": "diffusers-main/custom_train.py",
4  "start_line": 195,
5  "end_line": 227,
6  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Dual-path gated fusion: combines additive and multiplicative interactions.\n    \n    Instead of simple addition (FiLM) or no fusion (cross-attn/adanorm),\n    we use a content-dependent gating that allows the class embedding to\n    both additively shift AND multiplicatively modulate the time embedding.\n    This provides richer time-class interactions: at different timesteps,\n    different class dimensions become more/less relevant.\n    \n    Args:\n        time_emb:  [B, time_embed_dim] timestep embedding\n        class_emb: [B, time_embed_dim] class embedding\n\n    Returns: [B, time_embed_dim] embedding used in ResBlocks\n    \"\"\"\n    # Multiplicative gating: class determines which time dimensions are active\n    gate = torch.sigmoid(class_emb)\n    # Combined: additive class info + multiplicative modulation of time\n    return time_emb * (1.0 + 0.1 * gate) + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Adaptive LayerNorm-Zero conditioning applied after each block.\n    \n    Combined with the FiLM-style injection in prepare_conditioning,\n    this provides dual-path conditioning: class info enters both through\n    the time embedding (into every ResBlock) AND through per-block AdaLN-Zero\n    modulation. The AdaLN-Zero provides fine-grained, block-specific control\n    while FiLM provides deep integration throughout the ResBlock internals.\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        self.adaln = AdaLNBlock(channels, cond_dim)\n\n    def forward(self, h, class_emb):\n        return self.adaln(h, class_emb)\n"
7}

1{
2  "op": "replace",
3  "filename": "diffusers-main/custom_train.py",
4  "start_line": 195,
5  "end_line": 227,
6  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Dual-pathway conditioning: inject class info through BOTH the FiLM path\n    (adding to time_emb, which modulates ResBlock GroupNorm scale/shift) AND\n    through the cross-attention ClassConditioner after each block.\n\n    Motivation: FiLM provides global, channel-wise class conditioning through\n    the ResBlock's built-in adaptive GroupNorm. Cross-attention provides\n    spatially-varying, per-pixel class conditioning. These are complementary\n    signals — combining them gives the model two orthogonal ways to use class\n    information without any additional parameters.\n\n    Args:\n        time_emb:  [B, time_embed_dim] timestep embedding\n        class_emb: [B, time_embed_dim] class embedding\n\n    Returns: [B, time_embed_dim] embedding used in ResBlocks\n    \"\"\"\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Cross-attention pathway for spatial class conditioning.\n\n    Combined with the FiLM pathway in prepare_conditioning, this creates\n    a dual-pathway scheme: global (FiLM) + spatial (cross-attention).\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n\n    def forward(self, h, class_emb):\n        return self.cross_attn(h, class_emb)"
7}

1{
2  "op": "replace",
3  "filename": "diffusers-main/custom_train.py",
4  "start_line": 195,
5  "end_line": 227,
6  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Dual-path conditioning: FiLM (additive to time_emb) + cross-attention.\n\n    Adding class_emb to time_emb injects class info through the ResBlock\n    adaptive GroupNorm (scale/shift), providing multiplicative channel\n    modulation inside each ResBlock. This complements the cross-attention\n    ClassConditioner which adds a residual bias after each block.\n    \"\"\"\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Multi-token cross-attention with class-conditioned token expansion.\n\n    Standard cross-attention with 1 KV token degenerates (softmax=1 always),\n    making Q irrelevant. We expand the class embedding into K learnable\n    context tokens, making the attention non-degenerate: different spatial\n    positions can attend to different class aspects.\n\n    Combined with FiLM (via prepare_conditioning), this gives dual-path\n    class conditioning: FiLM for global channel modulation inside ResBlocks,\n    multi-token cross-attention for spatially-adaptive conditioning after.\n    \"\"\"\n    def __init__(self, channels, cond_dim, num_tokens=4):\n        super().__init__()\n        self.num_tokens = num_tokens\n        # Project class embedding to multiple context tokens\n        self.token_proj = nn.Linear(cond_dim, num_tokens * cond_dim)\n        # Cross-attention with multi-token context\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n\n    def forward(self, h, class_emb):\n        # Expand class_emb into multiple tokens: [B, cond_dim] -> [B, num_tokens, cond_dim]\n        B = class_emb.shape[0]\n        multi_ctx = self.token_proj(class_emb).view(B, self.num_tokens, -1)\n        # Use multi-token cross attention\n        return self._multi_token_cross_attn(h, multi_ctx)\n\n    def _multi_token_cross_attn(self, x, context):\n        \"\"\"Cross-attention with multiple KV tokens.\"\"\"\n        B, C, H, W = x.shape\n        ca = self.cross_attn\n        h = ca.norm(x).view(B, C, -1).transpose(1, 2)  # [B, HW, C]\n        # context: [B, num_tokens, cond_dim]\n        q = ca.q_proj(h).view(B, -1, ca.num_heads, ca.head_dim).transpose(1, 2)\n        k = ca.k_proj(context).view(B, -1, ca.num_heads, ca.head_dim).transpose(1, 2)\n        v = ca.v_proj(context).view(B, -1, ca.num_heads, ca.head_dim).transpose(1, 2)\n        attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(ca.head_dim)\n        attn = F.softmax(attn, dim=-1)\n        out = torch.matmul(attn, v)\n        out = out.transpose(1, 2).reshape(B, H * W, C)\n        out = ca.out_proj(out)\n        return x + out.transpose(1, 2).view(B, C, H, W)\n"
7}

1{
2  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Dual-path conditioning: FiLM (additive to time_emb) + cross-attention.\n\n    Adding class_emb to time_emb injects class info through the ResBlock\n    adaptive GroupNorm (scale/shift), providing multiplicative channel\n    modulation inside each ResBlock. This complements the cross-attention\n    ClassConditioner which adds a residual bias after each block.\n    \"\"\"\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Multi-token cross-attention with class-conditioned token expansion.\n\n    Standard cross-attention with 1 KV token degenerates (softmax=1 always),\n    making Q irrelevant. We expand the class embedding into K learnable\n    context tokens via a lightweight bottleneck projection, making the\n    attention non-degenerate: different spatial positions can attend to\n    different class aspects.\n\n    Combined with FiLM (via prepare_conditioning), this gives dual-path\n    class conditioning: FiLM for global channel modulation inside ResBlocks,\n    multi-token cross-attention for spatially-adaptive conditioning after.\n    \"\"\"\n    def __init__(self, channels, cond_dim, num_tokens=4):\n        super().__init__()\n        self.num_tokens = num_tokens\n        self.cond_dim = cond_dim\n        # Lightweight bottleneck to generate multiple tokens\n        # bottleneck_dim keeps param count low\n        bottleneck_dim = cond_dim // 4\n        self.token_proj = nn.Sequential(\n            nn.SiLU(),\n            nn.Linear(cond_dim, bottleneck_dim),\n            nn.SiLU(),\n            nn.Linear(bottleneck_dim, num_tokens * cond_dim),\n        )\n        # Cross-attention with multi-token context\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n\n    def forward(self, h, class_emb):\n        # Expand class_emb into multiple tokens: [B, cond_dim] -> [B, num_tokens, cond_dim]\n        B = class_emb.shape[0]\n        multi_ctx = self.token_proj(class_emb).view(B, self.num_tokens, self.cond_dim)\n        # Use multi-token cross attention\n        return self._multi_token_cross_attn(h, multi_ctx)\n\n    def _multi_token_cross_attn(self, x, context):\n        \"\"\"Cross-attention with multiple KV tokens.\"\"\"\n        B, C, H, W = x.shape\n        ca = self.cross_attn\n        h = ca.norm(x).view(B, C, -1).transpose(1, 2)  # [B, HW, C]\n        # context: [B, num_tokens, cond_dim]\n        q = ca.q_proj(h).view(B, -1, ca.num_heads, ca.head_dim).transpose(1, 2)\n        k = ca.k_proj(context).view(B, -1, ca.num_heads, ca.head_dim).transpose(1, 2)\n        v = ca.v_proj(context).view(B, -1, ca.num_heads, ca.head_dim).transpose(1, 2)\n        attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(ca.head_dim)\n        attn = F.softmax(attn, dim=-1)\n        out = torch.matmul(attn, v)\n        out = out.transpose(1, 2).reshape(B, H * W, C)\n        out = ca.out_proj(out)\n        return x + out.transpose(1, 2).view(B, C, H, W)\n",
3  "end_line": 247,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 195
7}

1{
2  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Dual-path conditioning: FiLM (additive to time_emb) + cross-attention.\n\n    Adding class_emb to time_emb injects class info through the ResBlock\n    adaptive GroupNorm (scale/shift), providing multiplicative channel\n    modulation inside each ResBlock. This complements the cross-attention\n    ClassConditioner which adds a learned residual bias after each block.\n\n    Motivation: FiLM and cross-attention are complementary conditioning paths.\n    FiLM modulates feature statistics globally via GroupNorm affine params,\n    while the post-block conditioner adds spatially-uniform but channel-wise\n    adaptive residuals. Having both paths provides redundant but\n    complementary class signal, improving conditioning robustness.\n    \"\"\"\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Gated channel modulation: projects class embedding to scale, shift,\n    and gate parameters applied after GroupNorm, with sigmoid gating for\n    stable training.\n\n    Unlike AdaLN-Zero which uses unbounded linear gating, we use sigmoid\n    gating which is bounded in [0,1], providing better gradient flow and\n    more stable interpolation between the identity (gate=0.5 at init with\n    appropriate bias) and the modulated features.\n\n    Combined with FiLM (via prepare_conditioning), this provides dual-path\n    class conditioning: FiLM for in-ResBlock modulation, gated channel\n    modulation for post-block refinement.\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        self.norm = nn.GroupNorm(32, channels)\n        # Project class embedding to scale, shift, gate\n        self.proj = nn.Sequential(\n            nn.SiLU(),\n            nn.Linear(cond_dim, 3 * channels),\n        )\n        # Zero-init for stable start (identity mapping)\n        nn.init.zeros_(self.proj[-1].weight)\n        nn.init.zeros_(self.proj[-1].bias)\n        # Set gate bias so sigmoid(0) = 0.5 initially -> start near identity\n        # Actually with zero init, gate=sigmoid(0)=0.5, scale=0, shift=0\n        # Result: x + 0.5 * (norm(x) * 1 + 0 - x) = 0.5*x + 0.5*norm(x)\n        # Better: init gate bias to -3 so sigmoid(-3)≈0.05, nearly identity\n        with torch.no_grad():\n            self.proj[-1].bias.data[2*channels:] = -3.0\n\n    def forward(self, h, class_emb):\n        params = self.proj(class_emb)\n        scale, shift, gate_logit = params.unsqueeze(-1).unsqueeze(-1).chunk(3, dim=1)\n        gate = torch.sigmoid(gate_logit)\n        normed = self.norm(h.float()).to(h.dtype)\n        modulated = normed * (1 + scale) + shift\n        return h + gate * (modulated - h)\n",
3  "end_line": 256,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 195
7}

1{
2  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Dual-path conditioning: FiLM (additive to time_emb) + cross-attention.\n\n    Adding class_emb to time_emb injects class info through the ResBlock\n    adaptive GroupNorm (scale/shift), providing multiplicative channel\n    modulation inside each ResBlock. This complements the cross-attention\n    ClassConditioner which adds a residual bias after each block.\n\n    Motivation: FiLM and cross-attention are complementary conditioning paths.\n    FiLM modulates feature statistics globally via GroupNorm affine params,\n    while cross-attention adds spatially-uniform but channel-wise adaptive\n    residuals. Having both paths provides redundant but complementary class\n    signal, improving conditioning robustness.\n    \"\"\"\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Cross-attention with FiLM dual-path conditioning.\n\n    Uses standard cross-attention (class embedding as key/value) combined\n    with FiLM conditioning through the time embedding path. The dual-path\n    approach provides complementary conditioning: FiLM for in-ResBlock\n    channel modulation, cross-attention for post-block refinement.\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n\n    def forward(self, h, class_emb):\n        return self.cross_attn(h, class_emb)\n",
3  "end_line": 250,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 195
7}