Spaces:
Running on Zero
Running on Zero
| """Inference-facing model registry for StreamDiffusionV2.""" | |
| from .wan.wan_wrapper import ( | |
| CausalWanDiffusionWrapper, | |
| WanDiffusionWrapper, | |
| WanTextEncoder, | |
| WanVAEWrapper, | |
| ) | |
| DIFFUSION_NAME_TO_CLASS = { | |
| "wan": WanDiffusionWrapper, | |
| "causal_wan": CausalWanDiffusionWrapper, | |
| } | |
| TEXT_ENCODER_NAME_TO_CLASS = { | |
| "wan": WanTextEncoder, | |
| "causal_wan": WanTextEncoder, | |
| } | |
| VAE_NAME_TO_CLASS = { | |
| "wan": WanVAEWrapper, | |
| "causal_wan": WanVAEWrapper, | |
| } | |
| def get_diffusion_wrapper(model_name): | |
| return DIFFUSION_NAME_TO_CLASS[model_name] | |
| def get_text_encoder_wrapper(model_name): | |
| return TEXT_ENCODER_NAME_TO_CLASS[model_name] | |
| def get_vae_wrapper(model_name): | |
| return VAE_NAME_TO_CLASS[model_name] | |