From acffc2d9fcefd23244cd6e897b9c7f8db04c093d Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 9 Jan 2025 12:54:57 +0530 Subject: [PATCH] parse subfolder when needed. --- finetrainers/cogvideox/cogvideox_lora.py | 3 ++- finetrainers/hunyuan_video/hunyuan_video_lora.py | 3 ++- finetrainers/ltx_video/ltx_video_lora.py | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/finetrainers/cogvideox/cogvideox_lora.py b/finetrainers/cogvideox/cogvideox_lora.py index c3b754a..9d3a77c 100644 --- a/finetrainers/cogvideox/cogvideox_lora.py +++ b/finetrainers/cogvideox/cogvideox_lora.py @@ -42,8 +42,9 @@ def load_diffusion_models( cache_dir: Optional[str] = None, **kwargs, ): + subfolder = kwargs.get("subfolder", None) transformer = CogVideoXTransformer3DModel.from_pretrained( - model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir + model_id, subfolder=subfolder, torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir ) scheduler = CogVideoXDDIMScheduler.from_pretrained(model_id, subfolder="scheduler") return {"transformer": transformer, "scheduler": scheduler} diff --git a/finetrainers/hunyuan_video/hunyuan_video_lora.py b/finetrainers/hunyuan_video/hunyuan_video_lora.py index 9bfea53..450101d 100644 --- a/finetrainers/hunyuan_video/hunyuan_video_lora.py +++ b/finetrainers/hunyuan_video/hunyuan_video_lora.py @@ -63,8 +63,9 @@ def load_diffusion_models( cache_dir: Optional[str] = None, **kwargs, ) -> Dict[str, Union[nn.Module, FlowMatchEulerDiscreteScheduler]]: + subfolder = kwargs.get("subfolder", None) transformer = HunyuanVideoTransformer3DModel.from_pretrained( - model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir + model_id, subfolder=subfolder, torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir ) scheduler = FlowMatchEulerDiscreteScheduler(shift=shift) return {"transformer": transformer, "scheduler": scheduler} diff --git a/finetrainers/ltx_video/ltx_video_lora.py b/finetrainers/ltx_video/ltx_video_lora.py index 0e1af9b..7a80543 100644 --- a/finetrainers/ltx_video/ltx_video_lora.py +++ b/finetrainers/ltx_video/ltx_video_lora.py @@ -45,8 +45,9 @@ def load_diffusion_models( cache_dir: Optional[str] = None, **kwargs, ) -> Dict[str, nn.Module]: + subfolder = kwargs.get("subfolder", None) transformer = LTXVideoTransformer3DModel.from_pretrained( - model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir + model_id, subfolder=subfolder, torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir ) scheduler = FlowMatchEulerDiscreteScheduler() return {"transformer": transformer, "scheduler": scheduler}