File size: 2,417 Bytes
1156de8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """
Classifier-Free Guidance implementation using the diffusers library.
"""
import torch.nn as nn
from diffusers import UNet2DModel
class CustomClassConditionedUnet(UNet2DModel):
"""UNet2DModel adapted for multi-hot classification vectors"""
def __init__(
self,
sample_size,
in_channels=1,
out_channels=1,
layers_per_block=2,
block_out_channels=(128, 128, 256, 256, 512, 512),
down_block_types=(
"DownBlock2D",
"DownBlock2D",
"DownBlock2D",
"DownBlock2D",
"AttnDownBlock2D",
"DownBlock2D",
),
up_block_types=(
"UpBlock2D",
"AttnUpBlock2D",
"UpBlock2D",
"UpBlock2D",
"UpBlock2D",
"UpBlock2D",
),
multihot_dim=14,
**kwargs,
):
# Remove conflicting parameters if they exist
kwargs.pop("class_embed_type", None)
kwargs.pop("num_class_embeds", None)
# Initialize the base model without class conditioning
super().__init__(
sample_size=sample_size,
in_channels=in_channels,
out_channels=out_channels,
layers_per_block=layers_per_block,
block_out_channels=block_out_channels,
down_block_types=down_block_types,
up_block_types=up_block_types,
**kwargs,
)
# Compute the time embedding dimension
time_embed_dim = block_out_channels[0] * 4
# Replace the class embedding with a linear layer for multihot vectors
self.class_embedding = nn.Linear(multihot_dim, time_embed_dim)
# Save the multihot dimension
self.config.multihot_dim = multihot_dim
def forward(self, sample, timestep, class_labels=None, return_dict=True):
"""
Forward pass that accepts multi-hot vectors for class_labels
Args:
sample: Image tensor [batch_size, channels, height, width]
timestep: Time steps [batch_size] or scalar
class_labels: Multi-hot vector [batch_size, multihot_dim]
return_dict: Whether to return a dictionary or just the sample
Returns:
Model prediction (noise or clean image depending on configuration)
"""
return super().forward(sample, timestep, class_labels, return_dict)
|