parse subfolder when needed.

This commit is contained in:
sayakpaul
2025-01-09 12:54:57 +05:30
parent f311f16d98
commit acffc2d9fc
3 changed files with 6 additions and 3 deletions
+2 -1
View File
@@ -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}
@@ -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}
+2 -1
View File
@@ -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}