fix denoising stage init

2025-06-06 11:44:36 -07:00
1 changed files with 11 additions and 8 deletions
@@ -47,14 +47,17 @@ class DenoisingStage(PipelineStage):
        super().__init__()
        self.transformer = transformer
        self.scheduler = scheduler
-        attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
+        # when used for validation, transformer is None as it is taking from the
-        self.attn_backend = get_attn_backend(
+        # training loop
-            head_size=attn_head_size,
+        if transformer is not None:
-            dtype=torch.float16,  # TODO(will): hack
+            attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
-            supported_attention_backends=(_Backend.SLIDING_TILE_ATTN,
+            self.attn_backend = get_attn_backend(
-                                          _Backend.FLASH_ATTN,
+                head_size=attn_head_size,
-                                          _Backend.TORCH_SDPA)  # hack
+                dtype=torch.float16,  # TODO(will): hack
-        )
+                supported_attention_backends=(_Backend.SLIDING_TILE_ATTN,
                                              _Backend.FLASH_ATTN,
                                              _Backend.TORCH_SDPA)  # hack
            )
    def forward(
        self,