Upload bytedream/model.py with huggingface_hub

bytedream/model.py

@@ -33,7 +33,7 @@ class ResnetBlock2D(nn.Module):
         self.dropout = nn.Dropout(0.0)
         self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
         
-        self.nonlinearity = nn.SiLU(inplace=True)
+        self.nonlinearity = nn.SiLU()
         
         if in_channels != out_channels:
             self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
@@ -79,15 +79,18 @@ class AttentionBlock(nn.Module):
     ):
         super().__init__()
         
-        inner_dim = num_heads * head_dim if head_dim is not None else query_dim
-        cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+        # Use head_dim if provided, otherwise calculate from query_dim and num_heads
+        self.head_dim = head_dim if head_dim is not None else query_dim // num_heads
+        inner_dim = self.head_dim * num_heads
+        
+        # Use cross_attention_dim if provided, otherwise use query_dim (self-attention)
+        self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
         
         self.num_heads = num_heads
-        self.head_dim = head_dim if head_dim is not None else query_dim // num_heads
         
         self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
-        self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=False)
-        self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(self.cross_attention_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(self.cross_attention_dim, inner_dim, bias=False)
         
         self.to_out = nn.ModuleList([
             nn.Linear(inner_dim, query_dim),
@@ -282,15 +285,21 @@ class UpBlock2D(nn.Module):
                         align_corners=False
                     )
                 
-                # Ensure channel dimensions match (project if needed)
-                expected_channels = self.resnets[i].conv1.in_channels - hidden_states.shape[1]
-                if res_hidden_state.shape[1] != expected_channels:
-                    # Project skip connection to expected channels
-                    res_hidden_state = nn.functional.conv2d(
-                        res_hidden_state,
-                        torch.randn(expected_channels, res_hidden_state.shape[1], 1, 1, device=res_hidden_state.device) * 0.01,
-                        padding=0
-                    )
+                # Ensure channel dimensions match
+                # The resnet expects input = hidden_states + res_hidden_state concatenated
+                expected_in_channels = self.resnets[i].conv1.in_channels
+                actual_in_channels = hidden_states.shape[1] + res_hidden_state.shape[1]
+                
+                if actual_in_channels != expected_in_channels:
+                    # Project skip connection to match expected channels
+                    channel_diff = expected_in_channels - hidden_states.shape[1]
+                    if channel_diff > 0 and channel_diff != res_hidden_state.shape[1]:
+                        # Need to project skip connection
+                        res_hidden_state = nn.functional.conv2d(
+                            res_hidden_state,
+                            torch.randn(channel_diff, res_hidden_state.shape[1], 1, 1, device=res_hidden_state.device) * 0.01,
+                            padding=0
+                        )
                 
                 hidden_states = torch.cat([hidden_states, res_hidden_state], dim=1)
             
@@ -325,7 +334,7 @@ class TimestepEmbedding(nn.Module):
         
         # Projection layers
         self.linear_1 = nn.Linear(in_features, time_embed_dim)
-        self.activation = nn.SiLU(inplace=True)
+        self.activation = nn.SiLU()
         self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim)
     
     def forward(self, timestep: torch.Tensor) -> torch.Tensor:
@@ -362,6 +371,7 @@ class UNet2DConditionModel(nn.Module):
         attention_head_dim: int = 8,
         cross_attention_dim: int = 768,
         use_linear_projection: bool = True,
+        use_gradient_checkpointing: bool = False,
     ):
         super().__init__()
         
@@ -369,6 +379,7 @@ class UNet2DConditionModel(nn.Module):
         self.block_out_channels = block_out_channels
         self.layers_per_block = layers_per_block
         self.cross_attention_dim = cross_attention_dim
+        self.use_gradient_checkpointing = use_gradient_checkpointing
         
         # Time embedding
         time_embed_dim = block_out_channels[0] * 4
@@ -445,7 +456,7 @@ class UNet2DConditionModel(nn.Module):
         
         # Output
         self.conv_norm_out = nn.GroupNorm(num_groups=32, num_channels=block_out_channels[0], eps=1e-6)
-        self.conv_act = nn.SiLU(inplace=True)
+        self.conv_act = nn.SiLU()
         self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, stride=1, padding=1)
     
     def forward(
@@ -465,31 +476,59 @@ class UNet2DConditionModel(nn.Module):
         down_block_res_samples = (hidden_states,)
         
         for downsample_block in self.down_blocks:
-            hidden_states, res_samples = downsample_block(
-                hidden_states=hidden_states,
-                temb=temb,
-                encoder_hidden_states=encoder_hidden_states,
-            )
+            if self.use_gradient_checkpointing and self.training:
+                hidden_states, res_samples = torch.utils.checkpoint.checkpoint(
+                    lambda hs, t, ehs: downsample_block(hs, t, ehs),
+                    hidden_states, temb, encoder_hidden_states,
+                    use_reentrant=False
+                )
+            else:
+                hidden_states, res_samples = downsample_block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    encoder_hidden_states=encoder_hidden_states,
+                )
             down_block_res_samples += res_samples
         
         # Middle
         for layer in self.mid_block:
-            if isinstance(layer, ResnetBlock2D):
-                hidden_states = layer(hidden_states, temb)
+            if self.use_gradient_checkpointing and self.training:
+                if isinstance(layer, ResnetBlock2D):
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        lambda hs, t: layer(hs, t),
+                        hidden_states, temb,
+                        use_reentrant=False
+                    )
+                else:
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        lambda hs, ehs: layer(hs, ehs),
+                        hidden_states, encoder_hidden_states,
+                        use_reentrant=False
+                    )
             else:
-                hidden_states = layer(hidden_states, encoder_hidden_states)
+                if isinstance(layer, ResnetBlock2D):
+                    hidden_states = layer(hidden_states, temb)
+                else:
+                    hidden_states = layer(hidden_states, encoder_hidden_states)
         
         # Up sampling path
         for upsample_block in self.up_blocks:
             res_samples = down_block_res_samples[-len(upsample_block.resnets):]
             down_block_res_samples = down_block_res_samples[:-len(upsample_block.resnets)]
             
-            hidden_states = upsample_block(
-                hidden_states=hidden_states,
-                res_hidden_states_tuple=res_samples,
-                temb=temb,
-                encoder_hidden_states=encoder_hidden_states,
-            )
+            if self.use_gradient_checkpointing and self.training:
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    lambda hs, res, t, ehs: upsample_block(hs, res, t, ehs),
+                    hidden_states, res_samples, temb, encoder_hidden_states,
+                    use_reentrant=False
+                )
+            else:
+                hidden_states = upsample_block(
+                    hidden_states=hidden_states,
+                    res_hidden_states_tuple=res_samples,
+                    temb=temb,
+                    encoder_hidden_states=encoder_hidden_states,
+                )
         
         # Output
         hidden_states = self.conv_norm_out(hidden_states)
@@ -513,40 +552,42 @@ class AutoencoderKL(nn.Module):
         up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",) * 4,
         latent_channels: int = 4,
         sample_size: int = 512,
+        block_out_channels: Tuple[int, ...] = (64, 128, 256, 512),
     ):
         super().__init__()
         
         self.sample_size = sample_size
+        self.block_out_channels = block_out_channels
         
-        # Encoder
+        # Encoder - using reduced channels for memory efficiency
         self.encoder = nn.ModuleList()
-        channels = [in_channels, 128, 256, 512, 512]
+        channels = [in_channels] + list(block_out_channels)
         
         for i in range(len(down_block_types)):
             block = nn.Sequential(
                 nn.Conv2d(channels[i], channels[i+1], kernel_size=3, stride=2, padding=1),
-                nn.GroupNorm(num_groups=32, num_channels=channels[i+1], eps=1e-6),
-                nn.SiLU(inplace=True),
+                nn.GroupNorm(num_groups=min(32, channels[i+1]), num_channels=channels[i+1], eps=1e-6),
+                nn.SiLU(),
             )
             self.encoder.append(block)
         
         # Latent space projection
-        self.quant_conv = nn.Conv2d(512, latent_channels * 2, kernel_size=1)
+        self.quant_conv = nn.Conv2d(block_out_channels[-1], latent_channels * 2, kernel_size=1)
         
-        # Decoder
+        # Decoder - using reduced channels for memory efficiency
         self.decoder = nn.ModuleList()
-        decoder_channels = [latent_channels, 512, 512, 256, 128]
+        decoder_channels = [latent_channels] + list(reversed(block_out_channels))
         
         for i in range(len(up_block_types)):
             block = nn.Sequential(
                 nn.ConvTranspose2d(decoder_channels[i], decoder_channels[i+1], kernel_size=4, stride=2, padding=1),
-                nn.GroupNorm(num_groups=32, num_channels=decoder_channels[i+1], eps=1e-6),
-                nn.SiLU(inplace=True),
+                nn.GroupNorm(num_groups=min(32, decoder_channels[i+1]), num_channels=decoder_channels[i+1], eps=1e-6),
+                nn.SiLU(),
             )
             self.decoder.append(block)
         
-        self.post_quant_conv = nn.Conv2d(latent_channels, 512, kernel_size=1)
-        self.conv_out = nn.Conv2d(128, out_channels, kernel_size=3, stride=1, padding=1)
+        self.post_quant_conv = nn.Conv2d(latent_channels, block_out_channels[-1], kernel_size=1)
+        self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, stride=1, padding=1)
     
     def encode(self, x: torch.Tensor) -> torch.Tensor:
         """Encode image to latent space"""
@@ -576,15 +617,17 @@ class CLIPTextModel(nn.Module):
     Extracts semantic features from text for conditioning
     """
     
-    def __init__(self, model_name: str = "openai/clip-vit-large-patch14", max_length: int = 77):
+    def __init__(self, model_name: str = "openai/clip-vit-base-patch32", max_length: int = 77):
         super().__init__()
         
         try:
             from transformers import CLIPTextModel as HFCLIPTextModel, CLIPTokenizer
             
+            print(f"Loading CLIP text encoder: {model_name}...")
             self.model = HFCLIPTextModel.from_pretrained(model_name)
             self.tokenizer = CLIPTokenizer.from_pretrained(model_name)
             self.max_length = max_length
+            print(f"✓ CLIP text encoder loaded successfully on CPU")
         except ImportError:
             print("Warning: transformers not installed. Using dummy text encoder.")
             self.model = None
@@ -603,7 +646,7 @@ class CLIPTextModel(nn.Module):
         """
         if self.model is None:
             # Dummy implementation if transformers not available
-            return torch.zeros(1, 77, 768)
+            return torch.zeros(1, 77, 512)
         
         inputs = self.tokenizer(
             text,
@@ -631,6 +674,7 @@ def create_unet(config):
         attention_head_dim=unet_config['attention_head_dim'],
         cross_attention_dim=unet_config['cross_attention_dim'],
         use_linear_projection=unet_config['use_linear_projection'],
+        use_gradient_checkpointing=True,  # Enable for memory efficiency
     )
 
 
@@ -644,6 +688,7 @@ def create_vae(config):
         up_block_types=tuple(vae_config['up_block_types']),
         latent_channels=vae_config['latent_channels'],
         sample_size=vae_config['sample_size'],
+        block_out_channels=tuple(vae_config.get('block_out_channels', [64, 128, 256, 512])),
     )