Big update, refactor model loading, breaks backwards compatibility

archiving old  version in legacy branch, probably not going to support SD3 version going forward

pyramid_dit/pyramid_dit_for_video_gen_pipeline.py

@@ -1,9 +1,6 @@
 import torch
-import os
-
 import torch.nn.functional as F
 
-from collections import OrderedDict
 from einops import rearrange
 from diffusers.utils.torch_utils import randn_tensor
 
@@ -12,18 +9,7 @@ from tqdm import tqdm
 
 from typing import List, Optional, Union
 from ..diffusion_schedulers import PyramidFlowMatchEulerDiscreteScheduler
-from ..video_vae.modeling_causal_vae import CausalVideoVAE
 
-
-from .mmdit_modules import (
-    PyramidDiffusionMMDiT,
-    SD3TextEncoderWithMask,
-)
-
-from .flux_modules import (
-    PyramidFluxTransformer,
-    FluxTextEncoderWithMask,
-)
 from comfy.utils import ProgressBar
 
 def compute_density_for_timestep_sampling(
@@ -40,74 +26,32 @@ def compute_density_for_timestep_sampling(
         u = torch.rand(size=(batch_size,), device="cpu")
     return u
 
-
-def build_pyramid_dit(
-    model_name : str,
-    model_path : str,
-    torch_dtype,
-    use_flash_attn : bool,
-    #use_mixed_training: bool,
-    interp_condition_pos: bool = True,
-    use_gradient_checkpointing: bool = False,
-    use_temporal_causal: bool = True,
-    gradient_checkpointing_ratio: float = 0.6,
-):
-    #model_dtype = torch.float32 if use_mixed_training else torch_dtype
-    if model_name == "pyramid_flux":
-        dit = PyramidFluxTransformer.from_pretrained(
-            model_path, torch_dtype=torch_dtype,
-            use_gradient_checkpointing=use_gradient_checkpointing, 
-            gradient_checkpointing_ratio=gradient_checkpointing_ratio,
-            use_flash_attn=use_flash_attn, use_temporal_causal=use_temporal_causal,
-            interp_condition_pos=interp_condition_pos, axes_dims_rope=[16, 24, 24],
-        )
-    elif model_name == "pyramid_mmdit":
-        dit = PyramidDiffusionMMDiT.from_pretrained(
-            model_path, torch_dtype=torch_dtype, use_gradient_checkpointing=use_gradient_checkpointing, 
-            gradient_checkpointing_ratio=gradient_checkpointing_ratio,
-            use_flash_attn=use_flash_attn, use_t5_mask=True, 
-            add_temp_pos_embed=True, temp_pos_embed_type='rope', 
-            use_temporal_causal=use_temporal_causal, interp_condition_pos=interp_condition_pos,
-        )
-    else:
-        raise NotImplementedError(f"Unsupported DiT architecture, please set the model_name to `pyramid_flux` or `pyramid_mmdit`")
-
-    return dit
-
-
-def build_text_encoder(
-    model_name : str,
-    model_path : str,
-    torch_dtype,
-    load_text_encoder: bool = True,
-):
-    # The text encoder
-    if load_text_encoder:
-        if model_name == "pyramid_flux":
-            text_encoder = FluxTextEncoderWithMask(model_path, torch_dtype=torch_dtype)
-        elif model_name == "pyramid_mmdit":
-            text_encoder = SD3TextEncoderWithMask(model_path, torch_dtype=torch_dtype)
-        else:
-            raise NotImplementedError(f"Unsupported Text Encoder architecture, please set the model_name to `pyramid_flux` or `pyramid_mmdit`")
-    else:
-        text_encoder = None
-
-    return text_encoder
-
-
 class PyramidDiTForVideoGeneration:
     """
         The pyramid dit for both image and video generation, The running class wrapper
         This class is mainly for fixed unit implementation: 1 + n + n + n
     """
-    def __init__(self, model_path, model_dtype, model_name, text_encoder_dtype, vae_dtype, use_gradient_checkpointing=False, return_log=True,
-        model_variant="diffusion_transformer_768p", timestep_shift=1.0, stage_range=[0, 1/3, 2/3, 1],
-        sample_ratios=[1, 1, 1], scheduler_gamma=1/3, use_flash_attn=False, 
-        load_text_encoder=True, load_vae=True, max_temporal_length=31, frame_per_unit=1, use_temporal_causal=True, 
-        corrupt_ratio=1/3, interp_condition_pos=True, stages=[1, 2, 4], fp8_fastmode=False, **kwargs,
-    ):
+    def __init__(
+            self,
+            transformer, 
+            model_dtype, 
+            model_name, 
+            return_log=True, 
+            timestep_shift=1.0, 
+            stage_range=[0, 1/3, 2/3, 1],
+            sample_ratios=[1, 1, 1], 
+            scheduler_gamma=1/3, 
+            max_temporal_length=31, 
+            frame_per_unit=1, 
+            use_temporal_causal=True, 
+            corrupt_ratio=1/3, 
+            interp_condition_pos=True, 
+            stages=[1, 2, 4], 
+            **kwargs,
+        ):
         super().__init__()
 
+        self.dit = transformer
         if model_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
             self.dtype = torch.bfloat16
         else:
@@ -117,43 +61,6 @@ class PyramidDiTForVideoGeneration:
         self.sample_ratios = sample_ratios
         self.corrupt_ratio = corrupt_ratio
         self.model_name = model_name
-
-       
-
-        dit_path = os.path.join(model_path, model_variant)
-
-        
-        # The dit
-        self.dit = build_pyramid_dit(
-            model_name, dit_path, self.dtype, 
-            use_flash_attn=use_flash_attn, 
-            interp_condition_pos=interp_condition_pos, use_gradient_checkpointing=use_gradient_checkpointing,
-            use_temporal_causal=use_temporal_causal,
-        )       
-
-        if model_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
-            for name, param in self.dit.named_parameters():
-                if name != "pos_embedding":
-                    param.data = param.data.to(model_dtype)
-
-        if model_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] and fp8_fastmode:
-            from ..fp8_optimization import convert_fp8_linear
-            convert_fp8_linear(self.dit, torch.bfloat16)
-
-        # The text encoder
-        self.text_encoder = build_text_encoder(
-            model_name, model_path, text_encoder_dtype, load_text_encoder=load_text_encoder,
-        )
-        self.load_text_encoder = load_text_encoder
-
-        # The base video vae decoder
-        if load_vae:
-            self.vae = CausalVideoVAE.from_pretrained(os.path.join(model_path, 'causal_video_vae'), torch_dtype=vae_dtype, interpolate=False)
-            # Freeze vae
-            for parameter in self.vae.parameters():
-                parameter.requires_grad = False
-        else:
-            self.vae = None
         
         # For the image latent
         self.vae_shift_factor = 0.1490
@@ -183,37 +90,7 @@ class PyramidDiTForVideoGeneration:
         
         self.cfg_rate = 0.1
         self.return_log = return_log
-        self.use_flash_attn = use_flash_attn
-
-    def load_checkpoint(self, checkpoint_path, model_key='model', **kwargs):
-        checkpoint = torch.load(checkpoint_path, map_location='cpu')
-        dit_checkpoint = OrderedDict()
-        for key in checkpoint:
-            if key.startswith('vae') or key.startswith('text_encoder'):
-                continue
-            if key.startswith('dit'):
-                new_key = key.split('.')
-                new_key = '.'.join(new_key[1:])
-                dit_checkpoint[new_key] = checkpoint[key]
-            else:
-                dit_checkpoint[key] = checkpoint[key]
-
-        load_result = self.dit.load_state_dict(dit_checkpoint, strict=True)
-        print(f"Load checkpoint from {checkpoint_path}, load result: {load_result}")
-
-    def load_vae_checkpoint(self, vae_checkpoint_path, model_key='model'):
-        checkpoint = torch.load(vae_checkpoint_path, map_location='cpu')
-        checkpoint = checkpoint[model_key]
-        loaded_checkpoint = OrderedDict()
-        
-        for key in checkpoint.keys():
-            if key.startswith('vae.'):
-                new_key = key.split('.')
-                new_key = '.'.join(new_key[1:])
-                loaded_checkpoint[new_key] = checkpoint[key]
-
-        load_result = self.vae.load_state_dict(loaded_checkpoint)
-        print(f"Load the VAE from {vae_checkpoint_path}, load result: {load_result}")
+        self.use_flash_attn = False
     
     @torch.no_grad()
     def get_pyramid_latent(self, x, stage_num):