drizzymedia/StreamDiffusionV2-Realtime

1"""Inference-facing model registry for StreamDiffusionV2."""2 3from .wan.wan_wrapper import (4    CausalWanDiffusionWrapper,5    WanDiffusionWrapper,6    WanTextEncoder,7    WanVAEWrapper,8)9 10 11DIFFUSION_NAME_TO_CLASS = {12    "wan": WanDiffusionWrapper,13    "causal_wan": CausalWanDiffusionWrapper,14}15 16 17TEXT_ENCODER_NAME_TO_CLASS = {18    "wan": WanTextEncoder,19    "causal_wan": WanTextEncoder,20}21 22 23VAE_NAME_TO_CLASS = {24    "wan": WanVAEWrapper,25    "causal_wan": WanVAEWrapper,26}27 28 29def get_diffusion_wrapper(model_name):30    return DIFFUSION_NAME_TO_CLASS[model_name]31 32 33def get_text_encoder_wrapper(model_name):34    return TEXT_ENCODER_NAME_TO_CLASS[model_name]35 36 37def get_vae_wrapper(model_name):38    return VAE_NAME_TO_CLASS[model_name]39