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    \"\"\"Prepare the combined embedding used in ResBlocks.\n\n    TODO: Implement your conditioning preparation here.\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    # Combined additive and multiplicative conditioning\n    # Simple additive conditioning like concat-film\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Conditioning module applied after each UNet block.\n\n    TODO: Implement your conditioning method here.\n\n    Args (forward):\n        h:         [B, C, H, W] feature map\n        class_emb: [B, time_embed_dim] class embedding\n\n    Available utilities:\n        CrossAttentionLayer(channels, context_dim, num_heads)\n        AdaLNBlock(channels, cond_dim)\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        # Cross-attention path\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n        # Adaptive normalization path  \n        self.adaln = AdaLNBlock(channels, cond_dim)\n        # Gating network: produces per-channel weight\n        self.gate_proj = nn.Sequential(\n            nn.SiLU(),\n            nn.Linear(cond_dim, channels)\n        )\n        # Initialize gate output to 0.0 (sigmoid(0.0)=0.5 for equal mixing)\n        nn.init.constant_(self.gate_proj[-1].weight, 0.0)\n        nn.init.constant_(self.gate_proj[-1].bias, 0.0)\n        \n    def forward(self, h, class_emb):\n        attn_out = self.cross_attn(h, class_emb)  # = h + attn_res\n        norm_out = self.adaln(h, class_emb)       # = h + norm_res\n        # Compute gating weights per channel\n        gate = torch.sigmoid(self.gate_proj(class_emb))  # [B, C]\n        gate = gate.view(gate.shape[0], gate.shape[1], 1, 1)  # [B, C, 1, 1]\n        # Combine residuals\n        combined_res = gate * (attn_out - h) + (1 - gate) * (norm_out - h)\n        return h + combined_res"
7}

1{
2  "n": 1
3}

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    \"\"\"Prepare the combined embedding used in ResBlocks.\n\n    TODO: Implement your conditioning preparation here.\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    # Simple additive conditioning like concat-film\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Conditioning module applied after each UNet block.\n\n    TODO: Implement your conditioning method here.\n\n    Args (forward):\n        h:         [B, C, H, W] feature map\n        class_emb: [B, time_embed_dim] class embedding\n\n    Available utilities:\n        CrossAttentionLayer(channels, context_dim, num_heads)\n        AdaLNBlock(channels, cond_dim)\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        # Cross-attention layer\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n        # Learn per-channel gating weight from class embedding\n        self.gate_proj = nn.Linear(cond_dim, channels)\n        # Initialize zero gate (start with no conditioning, let training learn)\n        nn.init.zeros_(self.gate_proj.weight)\n        nn.init.zeros_(self.gate_proj.bias)\n        \n    def forward(self, h, class_emb):\n        attn_out = self.cross_attn(h, class_emb)  # = h + attn_res\n        # Compute per-channel gate\n        gate = torch.sigmoid(self.gate_proj(class_emb))  # [B, C]\n        gate = gate.view(gate.shape[0], gate.shape[1], 1, 1)\n        # Blend between original h and attn_out\n        # h + gate * (attn_out - h) = (1-gate)*h + gate*attn_out\n        return h + gate * (attn_out - h)"
7}

1{
2  "n": 1
3}

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    \"\"\"Prepare the combined embedding used in ResBlocks.\n\n    TODO: Implement your conditioning preparation here.\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    # Simple additive conditioning like concat-film\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Conditioning module applied after each UNet block.\n\n    TODO: Implement your conditioning method here.\n\n    Args (forward):\n        h:         [B, C, H, W] feature map\n        class_emb: [B, time_embed_dim] class embedding\n\n    Available utilities:\n        CrossAttentionLayer(channels, context_dim, num_heads)\n        AdaLNBlock(channels, cond_dim)\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        # Adaptive LayerNorm-Zero conditioning\n        self.adaln = AdaLNBlock(channels, cond_dim)\n        # Self-gating: compute per-channel gate from feature map\n        self.self_gate = nn.Sequential(\n            nn.AdaptiveAvgPool2d(1),\n            nn.Conv2d(channels, channels, kernel_size=1),\n            nn.Sigmoid()\n        )\n        # Initialize conv weights zero so initial gate ~0.5\n        nn.init.zeros_(self.self_gate[1].weight)\n        nn.init.zeros_(self.self_gate[1].bias)\n        \n    def forward(self, h, class_emb):\n        # Adaptive normalization output\n        adaln_out = self.adaln(h, class_emb)  # = h + gate * (norm - h)\n        # Compute self-gate from original features\n        gate = self.self_gate(h)  # [B, C, 1, 1]\n        # Blend between original h and conditioned output\n        return h + gate * (adaln_out - h)"
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    \"\"\"Gated cross-attention: additive conditioning in time embedding.\"\"\"\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Gated cross-attention with channel-wise gating.\"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n        # Gate network: produce per-channel gating scalar\n        self.gate_net = nn.Sequential(\n            nn.SiLU(),\n            nn.Linear(cond_dim, channels)\n        )\n        # Initialize gate output to ~1.0 (bias) so initially cross-attn contributes fully\n        nn.init.zeros_(self.gate_net[-1].weight)\n        nn.init.ones_(self.gate_net[-1].bias)\n\n    def forward(self, h, class_emb):\n        cross_out = self.cross_attn(h, class_emb)\n        # Compute gate per channel, shape [B, C]\n        gate = torch.sigmoid(self.gate_net(class_emb))\n        # Reshape to [B, C, 1, 1] for broadcasting\n        gate = gate.view(gate.shape[0], gate.shape[1], 1, 1)\n        # Blend: h + gate * cross_out\n        return h + gate * cross_out"
7}