mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-17 00:23:35 +00:00
parse subfolder when needed.
This commit is contained in:
@@ -42,8 +42,9 @@ def load_diffusion_models(
|
|||||||
cache_dir: Optional[str] = None,
|
cache_dir: Optional[str] = None,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
):
|
):
|
||||||
|
subfolder = kwargs.get("subfolder", None)
|
||||||
transformer = CogVideoXTransformer3DModel.from_pretrained(
|
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")
|
scheduler = CogVideoXDDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
|
||||||
return {"transformer": transformer, "scheduler": scheduler}
|
return {"transformer": transformer, "scheduler": scheduler}
|
||||||
|
|||||||
@@ -63,8 +63,9 @@ def load_diffusion_models(
|
|||||||
cache_dir: Optional[str] = None,
|
cache_dir: Optional[str] = None,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
) -> Dict[str, Union[nn.Module, FlowMatchEulerDiscreteScheduler]]:
|
) -> Dict[str, Union[nn.Module, FlowMatchEulerDiscreteScheduler]]:
|
||||||
|
subfolder = kwargs.get("subfolder", None)
|
||||||
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
|
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)
|
scheduler = FlowMatchEulerDiscreteScheduler(shift=shift)
|
||||||
return {"transformer": transformer, "scheduler": scheduler}
|
return {"transformer": transformer, "scheduler": scheduler}
|
||||||
|
|||||||
@@ -45,8 +45,9 @@ def load_diffusion_models(
|
|||||||
cache_dir: Optional[str] = None,
|
cache_dir: Optional[str] = None,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
) -> Dict[str, nn.Module]:
|
) -> Dict[str, nn.Module]:
|
||||||
|
subfolder = kwargs.get("subfolder", None)
|
||||||
transformer = LTXVideoTransformer3DModel.from_pretrained(
|
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()
|
scheduler = FlowMatchEulerDiscreteScheduler()
|
||||||
return {"transformer": transformer, "scheduler": scheduler}
|
return {"transformer": transformer, "scheduler": scheduler}
|
||||||
|
|||||||
Reference in New Issue
Block a user