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 fusion of time and class embeddings.\n    \n    The class embedding modulates the time embedding via a sigmoid gate,\n    creating a synergistic combined embedding that carries both temporal\n    (noise level) and semantic (class) information to ResBlocks.\n    \"\"\"\n    gate = torch.sigmoid(class_emb)\n    return time_emb * (1 - gate * 0.5) + class_emb * gate * 0.5\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Cross-attention with complementary FiLM modulation.\n    \n    Combines spatial cross-attention (which learns WHERE to attend to class info)\n    with FiLM-style channel modulation (which learns WHICH channels to emphasize).\n    The FiLM path is zero-initialized and provides a complementary signal to\n    the cross-attention, improving conditioning expressiveness within budget.\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n        # Lightweight FiLM scale modulation (zero-initialized)\n        self.film_proj = nn.Linear(cond_dim, channels)\n        nn.init.zeros_(self.film_proj.weight)\n        nn.init.zeros_(self.film_proj.bias)\n        \n    def forward(self, h, class_emb):\n        # Cross-attention path: spatial features attend to class embedding\n        attn_out = self.cross_attn(h, class_emb)  # = h + residual_attn\n        # FiLM path: channel-wise scale modulation\n        scale = self.film_proj(class_emb).unsqueeze(-1).unsqueeze(-1)\n        # Combine: cross-attn residual + FiLM-modulated skip connection\n        # = residual_attn + h * (1 + scale)\n        return attn_out - h + h * (1 + scale)"
7}

1{
2  "op": "replace",
3  "filename": "diffusers-main/custom_train.py",
4  "start_line": 195,
5  "end_line": 229,
6  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Dual-pathway: inject class info into time embedding.\n    \n    The class embedding is added to the time embedding, creating a combined\n    embedding that carries class information into every ResBlock's FiLM layers.\n    This works synergistically with the spatial cross-attention in ClassConditioner.\n    \"\"\"\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Cross-attention with class-dependent residual gating.\n    \n    Combines spatial cross-attention (WHERE to attend) with a learned\n    per-channel gating mechanism that controls how much cross-attention\n    signal passes through each feature channel. The gate is predicted from\n    the class embedding, enabling class-specific conditioning strength.\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n        # Class-dependent residual gate (zero-initialized for stability)\n        self.gate_proj = nn.Sequential(nn.SiLU(), nn.Linear(cond_dim, channels))\n        nn.init.zeros_(self.gate_proj[-1].weight)\n        nn.init.zeros_(self.gate_proj[-1].bias)\n        \n    def forward(self, h, class_emb):\n        # Spatial cross-attention: features attend to class embedding\n        attn_out = self.cross_attn(h, class_emb)  # = h + residual_attn\n        # Extract the attention residual\n        residual = attn_out - h\n        # Class-dependent gate: sigmoid ensures 0..1 range\n        gate = torch.sigmoid(self.gate_proj(class_emb)).unsqueeze(-1).unsqueeze(-1)\n        # Gated residual: class controls how much attention signal flows\n        return h + residual * gate"
7}

1{
2  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Keep time embedding unchanged; class conditioning via ClassConditioner only.\"\"\"\n    return time_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Cross-attention + FiLM: complementary additive and multiplicative conditioning.\n    \n    Cross-attention provides a class-dependent additive signal to spatial features.\n    FiLM provides a class-dependent multiplicative signal to channels.\n    Together they give richer class conditioning than either alone.\n    FiLM is zero-initialized (scale=1, shift=0) for stable training start.\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        # Spatial conditioning via cross-attention\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n        # Channel conditioning via FiLM (scale + shift from class embedding)\n        self.film_proj = nn.Linear(cond_dim, 2 * channels)\n        # Initialize to identity: scale=1, shift=0\n        nn.init.zeros_(self.film_proj.weight)\n        nn.init.ones_(self.film_proj.bias[:channels])\n        nn.init.zeros_(self.film_proj.bias[channels:])\n        \n    def forward(self, h, class_emb):\n        # Cross-attention: class-dependent additive signal\n        attn_out = self.cross_attn(h, class_emb)\n        # FiLM: class-dependent multiplicative signal\n        scale, shift = self.film_proj(class_emb).chunk(2, dim=-1)\n        scale = scale.unsqueeze(-1).unsqueeze(-1)\n        shift = shift.unsqueeze(-1).unsqueeze(-1)\n        return attn_out * scale + shift",
3  "end_line": 229,
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-pathway: class embedding is added to time embedding for FiLM ResBlocks.\"\"\"\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Cross-attention + FiLM: complementary additive and multiplicative conditioning.\n    \n    Cross-attention provides a class-dependent additive signal to spatial features.\n    FiLM provides a class-dependent multiplicative signal to channels.\n    Together they give richer class conditioning than either alone.\n    FiLM is zero-initialized (scale=1, shift=0) for stable training start.\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        # Spatial conditioning via cross-attention\n        self.cross_attn = CrossAttentionLayer(channels, cond_dim, num_heads=4)\n        # Channel conditioning via FiLM (scale + shift from class embedding)\n        self.film_proj = nn.Linear(cond_dim, 2 * channels)\n        # Initialize to identity: scale=1, shift=0\n        nn.init.zeros_(self.film_proj.weight)\n        nn.init.ones_(self.film_proj.bias[:channels])\n        nn.init.zeros_(self.film_proj.bias[channels:])\n        \n    def forward(self, h, class_emb):\n        # Cross-attention: class-dependent additive signal\n        attn_out = self.cross_attn(h, class_emb)\n        # FiLM: class-dependent multiplicative signal\n        scale, shift = self.film_proj(class_emb).chunk(2, dim=-1)\n        scale = scale.unsqueeze(-1).unsqueeze(-1)\n        shift = shift.unsqueeze(-1).unsqueeze(-1)\n        return attn_out * scale + shift",
3  "end_line": 229,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 195
7}

1{
2  "end_line": 226,
3  "filename": "diffusers-main/custom_train.py",
4  "op": "replace",
5  "start_line": 195,
6  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Dual-pathway: inject class embedding into time embedding for ResBlock FiLM layers.\"\"\"\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Cross-attention for spatial class conditioning.\n    \n    Class embedding serves as key/value in cross-attention, allowing each\n    spatial position to attend to class-specific information. This provides\n    a complementary conditioning pathway to the FiLM injection via time_emb.\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  "content": "def prepare_conditioning(time_emb, class_emb):\n    \"\"\"Keep time embedding unchanged; class conditioning entirely via ClassConditioner.\"\"\"\n    return time_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Class-projected cross-attention: project class to channel space before attention.\n    \n    Instead of using the raw 512-dim class embedding as cross-attention context,\n    we first project it to match the feature channel dimension. This creates a\n    class-specific representation in the same space as spatial features, making\n    the K/V projections more effective and natural. The cross-attention then\n    operates on features and class projections in the same channel space.\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        # Project class embedding to feature channel space\n        self.class_proj = nn.Linear(cond_dim, channels)\n        # Cross-attention with matched dimensions\n        self.cross_attn = CrossAttentionLayer(channels, channels, num_heads=4)\n        \n    def forward(self, h, class_emb):\n        context = self.class_proj(class_emb)\n        return self.cross_attn(h, context)",
3  "end_line": 212,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 195
7}