From b045037850d6824b13da6f6f30bb88fdde64d95f Mon Sep 17 00:00:00 2001 From: Aryan Date: Sat, 21 Dec 2024 15:19:03 +0100 Subject: [PATCH] make style --- .../hunyuan_video/hunyuan_video_lora.py | 23 +++++++++++++++---- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/finetrainers/hunyuan_video/hunyuan_video_lora.py b/finetrainers/hunyuan_video/hunyuan_video_lora.py index 5fefb9c..e510102 100644 --- a/finetrainers/hunyuan_video/hunyuan_video_lora.py +++ b/finetrainers/hunyuan_video/hunyuan_video_lora.py @@ -25,10 +25,21 @@ def load_condition_models( **kwargs, ) -> Dict[str, nn.Module]: tokenizer = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, cache_dir=cache_dir) - text_encoder = LlamaModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=text_encoder_dtype, revision=revision, cache_dir=cache_dir) - tokenizer_2 = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer_2", revision=revision, cache_dir=cache_dir) - text_encoder_2 = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder_2", torch_dtype=text_encoder_2_dtype, revision=revision, cache_dir=cache_dir) - return {"tokenizer": tokenizer, "text_encoder": text_encoder, "tokenizer_2": tokenizer_2, "text_encoder_2": text_encoder_2} + text_encoder = LlamaModel.from_pretrained( + model_id, subfolder="text_encoder", torch_dtype=text_encoder_dtype, revision=revision, cache_dir=cache_dir + ) + tokenizer_2 = CLIPTokenizer.from_pretrained( + model_id, subfolder="tokenizer_2", revision=revision, cache_dir=cache_dir + ) + text_encoder_2 = CLIPTextModel.from_pretrained( + model_id, subfolder="text_encoder_2", torch_dtype=text_encoder_2_dtype, revision=revision, cache_dir=cache_dir + ) + return { + "tokenizer": tokenizer, + "text_encoder": text_encoder, + "tokenizer_2": tokenizer_2, + "text_encoder_2": text_encoder_2, + } def load_latent_models( @@ -38,7 +49,9 @@ def load_latent_models( cache_dir: Optional[str] = None, **kwargs, ) -> Dict[str, nn.Module]: - vae = AutoencoderKLHunyuanVideo.from_pretrained(model_id, subfolder="vae", torch_dtype=vae_dtype, revision=revision, cache_dir=cache_dir) + vae = AutoencoderKLHunyuanVideo.from_pretrained( + model_id, subfolder="vae", torch_dtype=vae_dtype, revision=revision, cache_dir=cache_dir + ) return {"vae": vae}