1{
2  "op": "replace",
3  "filename": "diffusers-main/custom_train.py",
4  "start_line": 31,
5  "end_line": 58,
6  "content": "def build_model(device):\n    \"\"\"Focused Attention with Adaptive Output Calibration (FAOC).\n    \n    Novel contributions:\n    1. Optimized attention placement: 16x16 + 8x8 only (iDDPM pattern)\n       with larger head_dim=16 (from ADM). At 4x4 resolution, conv \n       receptive fields already cover the entire image, making attention\n       redundant. This saves capacity for where it matters.\n    2. Dropout regularization (0.1) for the 35k-step training regime.\n    3. Timestep-Adaptive Output Calibration (TAOC): A lightweight learned\n       module that applies noise-level-dependent per-channel scaling to\n       the output. Unlike internal scale-shift FiLM (which conditions \n       intermediate features), TAOC calibrates the final epsilon prediction\n       based on the global noise level, correcting systematic biases in\n       predictions across the diffusion schedule.\n    \"\"\"\n    channels = (128, 256, 256, 256)\n    if os.environ.get('BLOCK_OUT_CHANNELS'):\n        channels = tuple(int(x) for x in os.environ['BLOCK_OUT_CHANNELS'].split(','))\n    layers = int(os.environ.get('LAYERS_PER_BLOCK', 2))\n\n    unet = UNet2DModel(\n        sample_size=32, in_channels=3, out_channels=3,\n        block_out_channels=channels,\n        down_block_types=(\"DownBlock2D\", \"AttnDownBlock2D\", \"AttnDownBlock2D\", \"DownBlock2D\"),\n        up_block_types=(\"UpBlock2D\", \"AttnUpBlock2D\", \"AttnUpBlock2D\", \"UpBlock2D\"),\n        layers_per_block=layers, norm_num_groups=32, norm_eps=1e-6,\n        act_fn=\"silu\", time_embedding_type=\"positional\",\n        flip_sin_to_cos=True, freq_shift=0, downsample_padding=1,\n        attention_head_dim=16, resnet_time_scale_shift=\"scale_shift\",\n        dropout=0.1,\n    )\n\n    class TAOCWrapper(nn.Module):\n        \"\"\"Timestep-Adaptive Output Calibration wrapper.\"\"\"\n        def __init__(self, unet, time_dim=128):\n            super().__init__()\n            self.unet = unet\n            # Sinusoidal features -> per-channel scale+shift\n            self.calibrator = nn.Sequential(\n                nn.Linear(time_dim, time_dim), nn.SiLU(),\n                nn.Linear(time_dim, 6),  # 3 scale + 3 shift\n            )\n            # Zero-init so wrapper starts as identity\n            nn.init.zeros_(self.calibrator[-1].weight)\n            nn.init.zeros_(self.calibrator[-1].bias)\n            self.time_dim = time_dim\n        def _sinusoidal_embed(self, t):\n            half = self.time_dim // 2\n            freqs = torch.exp(-math.log(10000.0) * torch.arange(\n                half, device=t.device, dtype=torch.float32) / half)\n            args = t.float().unsqueeze(-1) * freqs.unsqueeze(0)\n            return torch.cat([args.cos(), args.sin()], dim=-1)\n        def forward(self, x, timestep, **kwargs):\n            result = self.unet(x, timestep, **kwargs)\n            t_emb = self._sinusoidal_embed(timestep)\n            params = self.calibrator(t_emb)  # [B, 6]\n            scale = 1.0 + params[:, :3].unsqueeze(-1).unsqueeze(-1)\n            shift = params[:, 3:].unsqueeze(-1).unsqueeze(-1)\n            result.sample = result.sample * scale + shift\n            return result\n\n    model = TAOCWrapper(unet)\n    return model.to(device)\n"
7}

1{
2  "op": "replace",
3  "filename": "diffusers-main/custom_train.py",
4  "start_line": 31,
5  "end_line": 58,
6  "content": "def build_model(device):\n    \"\"\"Resolution-Adaptive Multi-Head Attention (RAMA) UNet.\n\n    Novel: per-level attention_head_dim that scales with spatial resolution.\n    At 16x16 (256 tokens): head_dim=4 -> many diverse attention patterns\n    At 8x8 (64 tokens): head_dim=8 -> balanced diversity and capacity\n    At 4x4 (16 tokens): head_dim=16 -> fewer but richer heads\n    Combined with 3-level attention, modern conditioning, and mild dropout.\n    \"\"\"\n    channels = (128, 256, 256, 256)\n    if os.environ.get('BLOCK_OUT_CHANNELS'):\n        channels = tuple(int(x) for x in os.environ['BLOCK_OUT_CHANNELS'].split(','))\n    layers = int(os.environ.get('LAYERS_PER_BLOCK', 2))\n    # Graduated head dims: no_attn at 32x32, then 4->8->16 as resolution drops\n    head_dims = (8, 4, 8, 16)\n    return UNet2DModel(\n        sample_size=32,\n        in_channels=3,\n        out_channels=3,\n        block_out_channels=channels,\n        down_block_types=(\"DownBlock2D\", \"AttnDownBlock2D\", \"AttnDownBlock2D\", \"AttnDownBlock2D\"),\n        up_block_types=(\"AttnUpBlock2D\", \"AttnUpBlock2D\", \"AttnUpBlock2D\", \"UpBlock2D\"),\n        layers_per_block=layers,\n        norm_num_groups=32,\n        norm_eps=1e-6,\n        act_fn=\"silu\",\n        time_embedding_type=\"positional\",\n        flip_sin_to_cos=True,\n        freq_shift=0,\n        downsample_padding=1,\n        attention_head_dim=head_dims,\n        resnet_time_scale_shift=\"scale_shift\",\n        dropout=0.05,\n    ).to(device)\n"
7}

1{
2  "content": "def build_model(device):\n    \"\"\"Hybrid Asymmetric Attention UNet for CIFAR-10 diffusion.\n\n    Novel contributions:\n    1) Asymmetric encoder-decoder attention: encoder uses attention only at\n       deepest levels (8x8, 4x4) for efficient feature extraction, while\n       decoder uses attention at 3 levels (16x16, 8x8, 4x4) for coherent\n       generation. This breaks the standard symmetric design.\n    2) Large attention heads (dim=16) for richer per-head representations\n       at low resolutions where token count is small.\n    3) Mild dropout (0.05) for regularization under short training.\n    \"\"\"\n    channels = (128, 256, 256, 256)\n    if os.environ.get('BLOCK_OUT_CHANNELS'):\n        channels = tuple(int(x) for x in os.environ['BLOCK_OUT_CHANNELS'].split(','))\n    layers = int(os.environ.get('LAYERS_PER_BLOCK', 2))\n    return UNet2DModel(\n        sample_size=32,\n        in_channels=3,\n        out_channels=3,\n        block_out_channels=channels,\n        down_block_types=(\"DownBlock2D\", \"DownBlock2D\", \"AttnDownBlock2D\", \"AttnDownBlock2D\"),\n        up_block_types=(\"AttnUpBlock2D\", \"AttnUpBlock2D\", \"AttnUpBlock2D\", \"UpBlock2D\"),\n        layers_per_block=layers,\n        norm_num_groups=32,\n        norm_eps=1e-6,\n        act_fn=\"silu\",\n        time_embedding_type=\"positional\",\n        flip_sin_to_cos=True,\n        freq_shift=0,\n        downsample_padding=1,\n        attention_head_dim=16,\n        resnet_time_scale_shift=\"scale_shift\",\n        dropout=0.05,\n    ).to(device)\n",
3  "end_line": 64,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 31
7}

1{
2  "op": "replace",
3  "filename": "diffusers-main/custom_train.py",
4  "start_line": 31,
5  "end_line": 58,
6  "content": "def build_model(device):\n    \"\"\"Novel: Top-heavy attention with resolution-scaled head dimensions.\n    \n    Key innovations over baselines:\n    1. Attention at 32/16/8 but NOT 4x4 — at 4x4 (16 tokens), self-attention\n       is redundant since convolutions already have global receptive field.\n       This saves parameters vs full-attn while keeping 32x32 global coherence.\n    2. Resolution-aware head_dim: smaller head_dim (more heads) at high res\n       captures diverse spatial patterns; larger head_dim (fewer heads) at low\n       res provides richer per-head feature representations.\n    \"\"\"\n    channels = (128, 256, 256, 256)\n    if os.environ.get('BLOCK_OUT_CHANNELS'):\n        channels = tuple(int(x) for x in os.environ['BLOCK_OUT_CHANNELS'].split(','))\n    layers = int(os.environ.get('LAYERS_PER_BLOCK', 2))\n\n    # Resolution-scaled head dims: 8 @ 32x32, 16 @ 16x16, 32 @ 8x8\n    head_dims = []\n    for i, c in enumerate(channels):\n        hd = min(8 * (2 ** i), c)\n        while c % hd != 0 and hd > 1:\n            hd -= 1\n        head_dims.append(hd)\n\n    return UNet2DModel(\n        sample_size=32, in_channels=3, out_channels=3,\n        block_out_channels=channels, layers_per_block=layers,\n        down_block_types=(\"AttnDownBlock2D\", \"AttnDownBlock2D\", \"AttnDownBlock2D\", \"DownBlock2D\"),\n        up_block_types=(\"UpBlock2D\", \"AttnUpBlock2D\", \"AttnUpBlock2D\", \"AttnUpBlock2D\"),\n        norm_num_groups=32, norm_eps=1e-6, act_fn=\"silu\",\n        time_embedding_type=\"positional\", flip_sin_to_cos=True, freq_shift=0,\n        downsample_padding=1, attention_head_dim=tuple(head_dims),\n        resnet_time_scale_shift=\"scale_shift\",\n    ).to(device)\n"
7}

1{
2  "content": "def build_model(device):\n    \"\"\"Novel: Top-heavy attention — attend at 32/16/8, skip 4x4.\n    \n    Key insight: At 4x4 resolution (16 spatial tokens), self-attention is\n    redundant — convolutions already span the full spatial extent. Meanwhile,\n    32x32 attention provides crucial global coherence that helps both small\n    and medium models. No baseline uses this specific 32/16/8 placement.\n    Combined with dropout=0.1 for attention regularization.\n    \"\"\"\n    channels = (128, 256, 256, 256)\n    if os.environ.get('BLOCK_OUT_CHANNELS'):\n        channels = tuple(int(x) for x in os.environ['BLOCK_OUT_CHANNELS'].split(','))\n    layers = int(os.environ.get('LAYERS_PER_BLOCK', 2))\n\n    return UNet2DModel(\n        sample_size=32, in_channels=3, out_channels=3,\n        block_out_channels=channels, layers_per_block=layers,\n        down_block_types=(\"AttnDownBlock2D\", \"AttnDownBlock2D\",\n                          \"AttnDownBlock2D\", \"DownBlock2D\"),\n        up_block_types=(\"UpBlock2D\", \"AttnUpBlock2D\",\n                        \"AttnUpBlock2D\", \"AttnUpBlock2D\"),\n        norm_num_groups=32, norm_eps=1e-6, act_fn=\"silu\",\n        time_embedding_type=\"positional\", flip_sin_to_cos=True, freq_shift=0,\n        downsample_padding=1, attention_head_dim=8,\n        resnet_time_scale_shift=\"scale_shift\",\n        dropout=0.1,\n    ).to(device)\n",
3  "end_line": 64,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 31
7}

1{
2  "content": "def build_model(device):\n    \"\"\"Novel: Top-heavy attention — attend at 32/16/8, skip 4x4.\n    \n    Key insight: At 4x4 resolution (16 spatial tokens), self-attention is\n    redundant — convolutions already span the full spatial extent. Meanwhile,\n    32x32 attention provides crucial global coherence that helps both small\n    and medium models. No baseline uses this specific 32/16/8 placement.\n    \"\"\"\n    channels = (128, 256, 256, 256)\n    if os.environ.get('BLOCK_OUT_CHANNELS'):\n        channels = tuple(int(x) for x in os.environ['BLOCK_OUT_CHANNELS'].split(','))\n    layers = int(os.environ.get('LAYERS_PER_BLOCK', 2))\n\n    return UNet2DModel(\n        sample_size=32, in_channels=3, out_channels=3,\n        block_out_channels=channels, layers_per_block=layers,\n        down_block_types=(\"AttnDownBlock2D\", \"AttnDownBlock2D\",\n                          \"AttnDownBlock2D\", \"DownBlock2D\"),\n        up_block_types=(\"UpBlock2D\", \"AttnUpBlock2D\",\n                        \"AttnUpBlock2D\", \"AttnUpBlock2D\"),\n        norm_num_groups=32, norm_eps=1e-6, act_fn=\"silu\",\n        time_embedding_type=\"positional\", flip_sin_to_cos=True, freq_shift=0,\n        downsample_padding=1, attention_head_dim=8,\n        resnet_time_scale_shift=\"scale_shift\",\n    ).to(device)\n",
3  "end_line": 57,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 31
7}