rename

Author: a-r-r-o-w

scripts/convert_hunyuan_video_to_diffusers.py

@@ -3,8 +3,9 @@
 
 import torch
 from accelerate import init_empty_weights
+from transformers import AutoModel, AutoTokenizer, CLIPTextModel, CLIPTokenizer
 
-from diffusers import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel
+from diffusers import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel, HunyuanVideoPipeline
 
 
 def remap_norm_scale_shift_(key, state_dict):
@@ -76,6 +77,8 @@ def remap_single_transformer_blocks_(key, state_dict):
     # "guidance_in.mlp.2": "time_text_embed.guidance_embedder.linear_2",
     # "vector_in.in_layer": "time_text_embed.text_embedder.linear_1",
     # "vector_in.out_layer": "time_text_embed.text_embedder.linear_2",
+    "txt_in.t_embedder": "txt_in.time_embed",
+    "txt_in.c_embedder": "txt_in.context_embed",
     "double_blocks": "transformer_blocks",
     "individual_token_refiner.blocks": "token_refiner.refiner_blocks",
     "img_attn_q_norm": "attn.norm_q",
@@ -179,6 +182,8 @@ def get_args():
         "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint"
     )
     parser.add_argument("--vae_ckpt_path", type=str, default=None, help="Path to original VAE checkpoint")
+    parser.add_argument("--text_encoder_path", type=str, default=None, help="Path to original llama checkpoint")
+    parser.add_argument("--text_encoder_2_path", type=str, default=None, help="Path to original clip checkpoint")
     parser.add_argument("--save_pipeline", action="store_true")
     parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved")
     parser.add_argument("--dtype", default="bf16", help="Torch dtype to save the transformer in.")
@@ -200,6 +205,8 @@ def get_args():
 
     if args.save_pipeline:
         assert args.transformer_ckpt_path is not None and args.vae_ckpt_path is not None
+        assert args.text_encoder_path is not None
+        assert args.text_encoder_2_path is not None
 
     if args.transformer_ckpt_path is not None:
         transformer = convert_transformer(args.transformer_ckpt_path)
@@ -211,3 +218,19 @@ def get_args():
         vae = convert_vae(args.vae_ckpt_path)
         if not args.save_pipeline:
             vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
+
+    if args.save_pipeline:
+        text_encoder = AutoModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.float16)
+        tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_path, padding_side="right")
+        text_encoder_2 = CLIPTextModel.from_pretrained(args.text_encoder_2_path, torch_dtype=torch.float16)
+        tokenizer_2 = CLIPTokenizer.from_pretrained(args.text_encoder_2_path)
+
+        pipe = HunyuanVideoPipeline(
+            transformer=transformer,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            text_encoder_2=text_encoder_2,
+            tokenizer_2=tokenizer_2,
+        )
+        pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")

src/diffusers/models/transformers/transformer_hunyuan_video.py

@@ -168,31 +168,6 @@ def forward(self, caption):
         return hidden_states
 
 
-def timestep_embedding(t, dim, max_period=10000):
-    """
-    Create sinusoidal timestep embeddings.
-
-    Args:
-        t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional.
-        dim (int): the dimension of the output.
-        max_period (int): controls the minimum frequency of the embeddings.
-
-    Returns:
-        embedding (torch.Tensor): An (N, D) Tensor of positional embeddings.
-
-    .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
-    """
-    half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-        device=t.device
-    )
-    args = t[:, None].float() * freqs[None]
-    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-    if dim % 2:
-        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-    return embedding
-
-
 class TimestepEmbedder(nn.Module):
     """
     Embeds scalar timesteps into vector representations.
@@ -219,7 +194,6 @@ def __init__(
         )
 
     def forward(self, t):
-        # t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).type(self.mlp[0].weight.dtype)
         t_freq = get_timestep_embedding(t, self.frequency_embedding_size, flip_sin_to_cos=True, max_period=self.max_period, downscale_freq_shift=0).type(self.mlp[0].weight.dtype)
         t_emb = self.mlp(t_freq)
         return t_emb
@@ -340,10 +314,8 @@ def __init__(
         hidden_size = num_attention_heads * attention_head_dim
 
         self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True)
-        # self.time_embed = TimestepEmbedder(hidden_size, nn.SiLU)
-        # self.context_embed = TextProjection(in_channels, hidden_size, nn.SiLU)
-        self.t_embedder = TimestepEmbedder(hidden_size, nn.SiLU)
-        self.c_embedder = TextProjection(in_channels, hidden_size, nn.SiLU)
+        self.time_embed = TimestepEmbedder(hidden_size, nn.SiLU)
+        self.context_embed = TextProjection(in_channels, hidden_size, nn.SiLU)
 
         self.token_refiner = IndividualTokenRefiner(
             num_attention_heads=num_attention_heads,
@@ -361,8 +333,7 @@ def forward(
         attention_mask: Optional[torch.LongTensor] = None,
     ) -> torch.Tensor:
         original_dtype = hidden_states.dtype
-        # temb = self.time_embed(timestep)
-        temb = self.t_embedder(timestep)
+        temb = self.time_embed(timestep)
 
         if attention_mask is None:
             pooled_projections = hidden_states.mean(dim=1)
@@ -371,8 +342,7 @@ def forward(
             pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1)
             pooled_projections = pooled_projections.to(original_dtype)
 
-        # pooled_projections = self.context_embed(pooled_projections)
-        pooled_projections = self.c_embedder(pooled_projections)
+        pooled_projections = self.context_embed(pooled_projections)
         emb = temb + pooled_projections
 
         hidden_states = self.input_embedder(hidden_states)