mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-12 21:49:57 +00:00
parse subfolder when needed.
This commit is contained in:
@@ -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}
|
||||
|
||||
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user