1{
2  "start_line": 195,
3  "op": "replace",
4  "end_line": 227,
5  "filename": "custom_train.py",
6  "content": "def prepare_conditioning(time_emb, class_emb):\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    def __init__(self, channels, cond_dim, num_heads=4, num_tokens=4):\n        super().__init__()\n        self.num_heads = num_heads\n        self.head_dim = channels // num_heads\n        self.num_tokens = num_tokens\n        self.norm = nn.GroupNorm(32, channels)\n        \n        token_dim = cond_dim // num_tokens\n        self.q_proj = nn.Linear(channels, channels)\n        self.k_proj = nn.Linear(token_dim, channels)\n        self.v_proj = nn.Linear(token_dim, channels)\n        self.out_proj = zero_module(nn.Linear(channels, channels))\n\n    def forward(self, h, class_emb):\n        B, C, H, W = h.shape\n        hx = self.norm(h).view(B, C, -1).transpose(1, 2)\n        \n        ctx = class_emb.view(B, self.num_tokens, -1)\n        \n        q = self.q_proj(hx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        k = self.k_proj(ctx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        v = self.v_proj(ctx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        \n        attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.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 = self.out_proj(out)\n        return h + out.transpose(1, 2).view(B, C, H, W)"
7}

1{
2  "filename": "diffusers-main/custom_train.py",
3  "op": "replace",
4  "start_line": 195,
5  "end_line": 228,
6  "content": "def prepare_conditioning(time_emb, class_emb):\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    def __init__(self, channels, cond_dim, num_heads=4, num_tokens=4):\n        super().__init__()\n        self.num_heads = num_heads\n        self.head_dim = channels // num_heads\n        self.num_tokens = num_tokens\n        self.norm = nn.GroupNorm(32, channels)\n        \n        token_dim = cond_dim // num_tokens\n        self.q_proj = nn.Linear(channels, channels)\n        self.k_proj = nn.Linear(token_dim, channels)\n        self.v_proj = nn.Linear(token_dim, channels)\n        self.out_proj = zero_module(nn.Linear(channels, channels))\n\n    def forward(self, h, class_emb):\n        B, C, H, W = h.shape\n        hx = self.norm(h).view(B, C, -1).transpose(1, 2)\n        \n        ctx = class_emb.view(B, self.num_tokens, -1)\n        \n        q = self.q_proj(hx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        k = self.k_proj(ctx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        v = self.v_proj(ctx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        \n        attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.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 = self.out_proj(out)\n        return h + out.transpose(1, 2).view(B, C, H, W)"
7}

1{
2  "content": "def prepare_conditioning(time_emb, class_emb):\n    return time_emb + class_emb\n\n\nclass ClassConditioner(nn.Module):\n    def __init__(self, channels, cond_dim, num_heads=4, num_tokens=4):\n        super().__init__()\n        self.num_heads = num_heads\n        self.head_dim = channels // num_heads\n        self.num_tokens = num_tokens\n        self.norm = nn.GroupNorm(32, channels)\n        \n        token_dim = cond_dim // num_tokens\n        self.q_proj = nn.Linear(channels, channels)\n        self.k_proj = nn.Linear(token_dim, channels)\n        self.v_proj = nn.Linear(token_dim, channels)\n        self.out_proj = zero_module(nn.Linear(channels, channels))\n\n    def forward(self, h, class_emb):\n        B, C, H, W = h.shape\n        hx = self.norm(h).view(B, C, -1).transpose(1, 2)\n        \n        ctx = class_emb.view(B, self.num_tokens, -1)\n        \n        q = self.q_proj(hx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        k = self.k_proj(ctx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        v = self.v_proj(ctx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        \n        attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.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 = self.out_proj(out)\n        return h + out.transpose(1, 2).view(B, C, H, W)",
3  "op": "replace",
4  "filename": "diffusers-main/custom_train.py",
5  "start_line": 195,
6  "end_line": 227
7}

1{
2  "filename": "diffusers-main/custom_train.py",
3  "end_line": 228,
4  "content": "def prepare_conditioning(time_emb, class_emb):\n    # Keep time embedding pure to prevent interference\n    return time_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Feature-Adaptive FiLM (FA-FiLM).\n    \n    Dynamically predicts scale and shift for the features based on both\n    the global class embedding and the current spatial feature state.\n    \"\"\"\n    def __init__(self, channels, cond_dim):\n        super().__init__()\n        self.pool = nn.AdaptiveAvgPool2d(1)\n        self.mlp = nn.Sequential(\n            nn.Linear(channels + cond_dim, channels),\n            nn.LayerNorm(channels),\n            nn.SiLU(),\n            nn.Linear(channels, 2 * channels)\n        )\n        # Zero-initialize for identity mapping at start\n        nn.init.zeros_(self.mlp[-1].weight)\n        nn.init.zeros_(self.mlp[-1].bias)\n\n    def forward(self, h, class_emb):\n        B, C, _, _ = h.shape\n        # Extract current feature state\n        pooled = self.pool(h).view(B, C)\n        # Combine state with global class condition\n        ctx = torch.cat([pooled, class_emb], dim=1)\n        \n        # Predict scale and shift\n        params = self.mlp(ctx).view(B, 2, C, 1, 1)\n        scale, shift = params.chunk(2, dim=1)\n        scale = scale.squeeze(1)\n        shift = shift.squeeze(1)\n        \n        return h * (1 + scale) + shift",
5  "start_line": 195,
6  "op": "replace"
7}

1{
2  "start_line": 195,
3  "end_line": 232,
4  "filename": "diffusers-main/custom_train.py",
5  "content": "def prepare_conditioning(time_emb, class_emb):\n    # Keep time_emb pure to avoid timestep interference\n    return time_emb\n\n\nclass ClassConditioner(nn.Module):\n    \"\"\"Dual-Injection Conditioning: Global FiLM + Multi-Token Spatial Attention.\"\"\"\n    def __init__(self, channels, cond_dim, num_heads=4, num_tokens=4):\n        super().__init__()\n        self.num_heads = num_heads\n        self.head_dim = channels // num_heads\n        self.num_tokens = num_tokens\n        \n        # 1. Global FiLM\n        self.film = nn.Sequential(\n            nn.SiLU(),\n            nn.Linear(cond_dim, 2 * channels)\n        )\n        nn.init.zeros_(self.film[-1].weight)\n        nn.init.zeros_(self.film[-1].bias)\n        \n        # 2. Multi-Token Cross Attention\n        self.norm = nn.GroupNorm(32, channels)\n        token_dim = cond_dim // num_tokens\n        self.q_proj = nn.Linear(channels, channels)\n        self.k_proj = nn.Linear(token_dim, channels)\n        self.v_proj = nn.Linear(token_dim, channels)\n        self.out_proj = zero_module(nn.Linear(channels, channels))\n\n    def forward(self, h, class_emb):\n        B, C, H, W = h.shape\n        \n        # Global modulation\n        scale, shift = self.film(class_emb).unsqueeze(-1).unsqueeze(-1).chunk(2, dim=1)\n        h_film = h * (1 + scale) + shift\n        \n        # Spatial Multi-Token routing\n        hx = self.norm(h_film).view(B, C, -1).transpose(1, 2)\n        ctx = class_emb.view(B, self.num_tokens, -1)\n        \n        q = self.q_proj(hx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        k = self.k_proj(ctx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        v = self.v_proj(ctx).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        \n        attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.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 = self.out_proj(out)\n        \n        return h_film + out.transpose(1, 2).view(B, C, H, W)",
6  "op": "replace"
7}