from dataclasses import dataclass
import diffusers
import transformers


@dataclass
class Model():
    name: str
    url: str = ''
    repo: str = None
    repo_cls: classmethod = None
    repo_revision: str = None
    dit: str = None
    dit_cls: classmethod = None
    dit_folder: str = 'transformer'
    dit_revision: str = None
    te: str = None
    te_cls: classmethod = None
    te_folder: str = 'text_encoder'
    te_hijack: bool = True
    te_revision: str = None
    image_hijack: bool = True
    vae_hijack: bool = True
    vae_remote: bool = False

    def __str__(self):
        return f'name="{self.name}" url="{self.url}" repo="{self.repo}" repo_cls="{self.repo_cls}" dit="{self.dit}" dit_cls="{self.dit_cls}" dit_folder="{self.dit_folder}" te="{self.te}" te_cls="{self.te_cls}" te_folder="{self.te_folder}" te_hijack={self.te_hijack} vae_hijack={self.vae_hijack} vae_remote={self.vae_remote}'


models = {
    'None': [],
    'Hunyuan Video': [
        Model(name='None'),
        Model(name='Hunyuan Video T2V',
              url='https://huggingface.co/tencent/HunyuanVideo',
              vae_remote=True,
              repo='hunyuanvideo-community/HunyuanVideo',
              repo_cls=diffusers.HunyuanVideoPipeline,
              te_cls=transformers.LlamaModel,
              dit_cls=diffusers.HunyuanVideoTransformer3DModel),
        Model(name='Hunyuan Video I2V', # https://github.com/huggingface/diffusers/pull/10983
              url='https://huggingface.co/tencent/HunyuanVideo-I2V',
              vae_remote=True,
              repo='hunyuanvideo-community/HunyuanVideo-I2V',
              repo_cls=diffusers.HunyuanVideoImageToVideoPipeline,
              te_cls=transformers.LlavaForConditionalGeneration,
              dit_cls=diffusers.HunyuanVideoTransformer3DModel),
        Model(name='SkyReels Hunyuan T2V', # https://github.com/huggingface/diffusers/pull/10837
              url='https://huggingface.co/Skywork/SkyReels-V1-Hunyuan-T2V',
              vae_remote=True,
              repo='hunyuanvideo-community/HunyuanVideo',
              repo_cls=diffusers.HunyuanVideoPipeline,
              te_cls=transformers.LlamaModel,
              dit='Skywork/SkyReels-V1-Hunyuan-T2V',
              dit_folder=None,
              dit_cls=diffusers.HunyuanVideoTransformer3DModel),
        Model(name='SkyReels Hunyuan I2V', # https://github.com/huggingface/diffusers/pull/10837
              url='https://huggingface.co/Skywork/SkyReels-V1-Hunyuan-I2V',
              vae_remote=True,
              repo='hunyuanvideo-community/HunyuanVideo',
              repo_cls=diffusers.HunyuanSkyreelsImageToVideoPipeline,
              te_cls=transformers.LlamaModel,
              dit='Skywork/SkyReels-V1-Hunyuan-I2V',
              dit_folder=None,
              dit_cls=diffusers.HunyuanVideoTransformer3DModel),
        Model(name='Fast Hunyuan T2V', # https://github.com/hao-ai-lab/FastVideo/blob/8a77cf22c9b9e7f931f42bc4b35d21fd91d24e45/fastvideo/models/hunyuan/inference.py#L213
              url='https://huggingface.co/FastVideo/FastHunyuan',
              vae_remote=True,
              repo='hunyuanvideo-community/HunyuanVideo',
              repo_cls=diffusers.HunyuanVideoPipeline,
              te_cls=transformers.LlamaModel,
              dit='FastVideo/FastHunyuan-diffusers',
              dit_cls=diffusers.HunyuanVideoTransformer3DModel),
    ],
    'LTX Video': [
        Model(name='None'),
        Model(name='LTXVideo 0.9.8 13B',
              url='https://huggingface.co/Lightricks/LTX-Video-0.9.8-13B-distilled',
              repo='Lightricks/LTX-Video-0.9.8-13B-distilled',
              repo_cls=diffusers.LTXConditionPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.7 13B',
              url='https://huggingface.co/Lightricks/LTX-Video-0.9.7-dev',
              repo='a-r-r-o-w/LTX-Video-0.9.7-diffusers',
              repo_cls=diffusers.LTXConditionPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.6 2B T2V',
              url='https://huggingface.co/Lightricks/LTX-Video',
              repo='Lightricks/LTX-Video',
              repo_cls=diffusers.LTXConditionPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.6 2B I2V',
              url='https://huggingface.co/Lightricks/LTX-Video',
              repo='Lightricks/LTX-Video',
              repo_cls=diffusers.LTXConditionPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.6 2B T2V Distilled',
              url='https://huggingface.co/Lightricks/LTX-Video-2B-0.9.6-Distilled-04-25',
              repo='Lightricks/LTX-Video-2B-0.9.6-Distilled-04-25',
              repo_cls=diffusers.LTXConditionPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.6 2B I2V Distilled',
              url='https://huggingface.co/Lightricks/LTX-Video-2B-0.9.6-Distilled-04-25',
              repo='Lightricks/LTX-Video-2B-0.9.6-Distilled-04-25',
              repo_cls=diffusers.LTXConditionPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.5 T2V', # https://github.com/huggingface/diffusers/pull/10968
              url='https://huggingface.co/Lightricks/LTX-Video-0.9.5',
              repo='Lightricks/LTX-Video-0.9.5',
              repo_cls=diffusers.LTXConditionPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.5 I2V',
              url='https://huggingface.co/Lightricks/LTX-Video-0.9.5',
              repo='Lightricks/LTX-Video-0.9.5',
              repo_cls=diffusers.LTXConditionPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.1 T2V',
              url='https://huggingface.co/a-r-r-o-w/LTX-Video-0.9.1-diffusers',
              repo='a-r-r-o-w/LTX-Video-0.9.1-diffusers',
              repo_cls=diffusers.LTXPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.1 I2V',
              url='https://huggingface.co/a-r-r-o-w/LTX-Video-0.9.1-diffusers',
              repo='a-r-r-o-w/LTX-Video-0.9.1-diffusers',
              repo_cls=diffusers.LTXImageToVideoPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.0 T2V',
              url='https://huggingface.co/a-r-r-o-w/LTX-Video-diffusers',
              repo='a-r-r-o-w/LTX-Video-diffusers',
              repo_cls=diffusers.LTXPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
        Model(name='LTXVideo 0.9.0 I2V',
              url='https://huggingface.co/a-r-r-o-w/LTX-Video-diffusers',
              repo='a-r-r-o-w/LTX-Video-diffusers',
              repo_cls=diffusers.LTXImageToVideoPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LTXVideoTransformer3DModel),
    ],
    'WAN Video': [
        Model(name='None'),
        Model(name='WAN 2.2 5B T2V',
              url='https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers',
              repo='Wan-AI/Wan2.2-TI2V-5B-Diffusers',
              repo_cls=diffusers.WanPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanTransformer3DModel),
        Model(name='WAN 2.2 5B I2V',
              url='https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers',
              repo='Wan-AI/Wan2.2-TI2V-5B-Diffusers',
              repo_cls=diffusers.WanImageToVideoPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanTransformer3DModel),
        Model(name='WAN 2.2 A14B T2V',
              url='https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B-Diffusers',
              repo='Wan-AI/Wan2.2-T2V-A14B-Diffusers',
              repo_cls=diffusers.WanPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanTransformer3DModel,
              dit_folder=("transformer", "transformer_2")),
        Model(name='WAN 2.2 A14B I2V',
              url='https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers',
              repo='Wan-AI/Wan2.2-I2V-A14B-Diffusers',
              repo_cls=diffusers.WanImageToVideoPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanTransformer3DModel,
              dit_folder=("transformer", "transformer_2")),
        Model(name='WAN 2.2 14B VACE',
              url='https://huggingface.co/linoyts/Wan2.2-VACE-Fun-14B-diffusers',
              repo='linoyts/Wan2.2-VACE-Fun-14B-diffusers',
              repo_cls=diffusers.WanVACEPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanVACETransformer3DModel,
              dit_folder=("transformer", "transformer_2")),
        Model(name='WAN 2.1 1.3B T2V',
              url='https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers',
              repo='Wan-AI/Wan2.1-T2V-1.3B-Diffusers',
              repo_cls=diffusers.WanPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanTransformer3DModel),
        Model(name='WAN 2.1 14B T2V',
              url='https://huggingface.co/Wan-AI/Wan2.1-T2V-14B-Diffusers',
              repo='Wan-AI/Wan2.1-T2V-14B-Diffusers',
              repo_cls=diffusers.WanPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanTransformer3DModel),
        Model(name='WAN 2.1 14B I2V 480p',
              url='https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P-Diffusers',
              repo='Wan-AI/Wan2.1-I2V-14B-480P-Diffusers',
              repo_cls=diffusers.WanImageToVideoPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanTransformer3DModel),
        Model(name='WAN 2.1 14B I2V 720p',
              url='https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P-Diffusers',
              repo='Wan-AI/Wan2.1-I2V-14B-720P-Diffusers',
              repo_cls=diffusers.WanImageToVideoPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanTransformer3DModel),
        Model(name='WAN 2.1 14B FLF2V 720p',
              url='https://huggingface.co/Wan-AI/Wan2.1-FLF2V-14B-720P',
              repo='Wan-AI/Wan2.1-FLF2V-14B-720P-diffusers',
              repo_cls=diffusers.WanImageToVideoPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanTransformer3DModel),
        Model(name='WAN 2.1 VACE 1.3B',
              url='https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers',
              repo='Wan-AI/Wan2.1-VACE-1.3B-diffusers',
              repo_cls=diffusers.WanVACEPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanVACETransformer3DModel),
        Model(name='WAN 2.1 VACE 14B',
              url='https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers',
              repo='Wan-AI/Wan2.1-VACE-14B-diffusers',
              repo_cls=diffusers.WanVACEPipeline,
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.WanVACETransformer3DModel),
    ],
    'SkyReels V2': [
        Model(name='None'),
        Model(name='SkyReels-V2 T2V-DF 1.3B-540P',
              url='https://huggingface.co/Skywork/SkyReels-V2-DF-1.3B-540P-Diffusers',
              repo='Skywork/SkyReels-V2-DF-1.3B-540P-Diffusers',
              repo_cls=diffusers.SkyReelsV2DiffusionForcingPipeline,
              repo_revision='refs/pr/1',
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.SkyReelsV2Transformer3DModel),
        Model(name='SkyReels-V2 T2V-DF 14B-720P',
              url='https://huggingface.co/Skywork/SkyReels-V2-DF-14B-720P-Diffusers',
              repo='Skywork/SkyReels-V2-DF-14B-720P-Diffusers',
              repo_cls=diffusers.SkyReelsV2DiffusionForcingPipeline,
              repo_revision='refs/pr/1',
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.SkyReelsV2Transformer3DModel),
        Model(name='SkyReels-V2 I2V-DF 14B-720P',
              url='https://huggingface.co/Skywork/SkyReels-V2-DF-14B-720P-Diffusers',
              repo='Skywork/SkyReels-V2-DF-14B-720P-Diffusers',
              repo_cls=diffusers.SkyReelsV2DiffusionForcingImageToVideoPipeline,
              repo_revision='refs/pr/1',
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.SkyReelsV2Transformer3DModel),
        Model(name='SkyReels-V2 T2V 14B-720P',
              url='https://huggingface.co/Skywork/SkyReels-V2-T2V-14B-720P-Diffusers',
              repo='Skywork/SkyReels-V2-T2V-14B-720P-Diffusers',
              repo_cls=diffusers.SkyReelsV2Pipeline,
              repo_revision='refs/pr/1',
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.SkyReelsV2Transformer3DModel),
        Model(name='SkyReels-V2 I2V 14B-720P',
              url='https://huggingface.co/Skywork/SkyReels-V2-I2V-14B-720P-Diffusers',
              repo='Skywork/SkyReels-V2-I2V-14B-720P-Diffusers',
              repo_cls=diffusers.SkyReelsV2ImageToVideoPipeline,
              repo_revision='refs/pr/1',
              te_cls=transformers.UMT5EncoderModel,
              dit_cls=diffusers.SkyReelsV2Transformer3DModel),
    ],
    'Mochi Video': [
        Model(name='None'),
        Model(name='Mochi 1 T2V',
              url='https://huggingface.co/genmo/mochi-1-preview',
              repo='genmo/mochi-1-preview',
              repo_cls=diffusers.MochiPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.MochiTransformer3DModel),
    ],
    'Latte Video': [
        Model(name='None'),
        Model(name='Latte 1 T2V',
              url='https://huggingface.co/maxin-cn/Latte-1',
              repo='maxin-cn/Latte-1',
              repo_cls=diffusers.LattePipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.LatteTransformer3DModel),
    ],
    'Allegro Video': [
        Model(name='None'),
        Model(name='Allegro T2V',
              url='https://huggingface.co/rhymes-ai/Allegro',
              repo='rhymes-ai/Allegro',
              repo_cls=diffusers.AllegroPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.AllegroTransformer3DModel),
    ],
    'Cog Video': [
        Model(name='None'),
        Model(name='CogVideoX 1.0 2B T2V',
              url='https://huggingface.co/THUDM/CogVideoX-2b',
              repo='THUDM/CogVideoX-2b',
              repo_cls=diffusers.CogVideoXPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.CogVideoXTransformer3DModel),
        Model(name='CogVideoX 1.0 5B T2V',
              url='https://huggingface.co/THUDM/CogVideoX-5b',
              repo='THUDM/CogVideoX-5b',
              repo_cls=diffusers.CogVideoXPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.CogVideoXTransformer3DModel),
        Model(name='CogVideoX 1.0 5B I2V',
              url='https://huggingface.co/THUDM/CogVideoX-5b-I2V',
              repo='THUDM/CogVideoX-5b-I2V',
              repo_cls=diffusers.CogVideoXImageToVideoPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.CogVideoXTransformer3DModel),
        Model(name='CogVideoX 1.5 5B T2V',
              url='https://huggingface.co/THUDM/CogVideoX1.5-5B',
              repo='THUDM/CogVideoX1.5-5B',
              repo_cls=diffusers.CogVideoXPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.CogVideoXTransformer3DModel),
        Model(name='CogVideoX 1.5 5B I2V',
              url='https://huggingface.co/THUDM/CogVideoX1.5-5B-I2V',
              repo='THUDM/CogVideoX1.5-5B-I2V',
              repo_cls=diffusers.CogVideoXImageToVideoPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.CogVideoXTransformer3DModel),
        Model(name='Index Anisora 1.0 5B I2V',
              url='https://huggingface.co/Disty0/Index-anisora-5B-diffusers',
              repo='Disty0/Index-anisora-5B-diffusers',
              repo_cls=diffusers.CogVideoXImageToVideoPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.CogVideoXTransformer3DModel),
        Model(name='Index Anisora 1.0 5B RL I2V',
              url='https://huggingface.co/Disty0/Index-anisora-5B_RL-diffusers',
              repo='Disty0/Index-anisora-5B_RL-diffusers',
              repo_cls=diffusers.CogVideoXImageToVideoPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.CogVideoXTransformer3DModel),
    ],
    'nVidia Cosmos': [
        Model(name='nvidia Cosmos Predict2 2B I2V',
              url='https://huggingface.co/nvidia/Cosmos-Predict2-2B-Text2Image',
              repo='nvidia/Cosmos-Predict2-2B-Video2World',
              repo_cls=diffusers.Cosmos2VideoToWorldPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.CosmosTransformer3DModel),
        Model(name='nvidia Cosmos Predict2 2B I2V',
              url='https://huggingface.co/nvidia/Cosmos-Predict2-2B-Text2Image',
              repo='nvidia/Cosmos-Predict2-2B-Video2World',
              repo_cls=diffusers.Cosmos2VideoToWorldPipeline,
              te_cls=transformers.T5EncoderModel,
              dit_cls=diffusers.CosmosTransformer3DModel),
    ],
    'Kandinsky': [
        Model(name='Kandinsky 5.0 Lite SFT T2V',
              url='https://huggingface.co/ai-forever/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers',
              repo='ai-forever/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers',
              repo_cls=diffusers.Kandinsky5T2VPipeline,
              te_cls=transformers.Qwen2_5_VLForConditionalGeneration,
              dit_cls=diffusers.Kandinsky5Transformer3DModel),
        Model(name='Kandinsky 5.0 Lite CFG-distilled T2V',
              url='https://huggingface.co/ai-forever/Kandinsky-5.0-T2V-Lite-nocfg-5s-Diffusers',
              repo='ai-forever/Kandinsky-5.0-T2V-Lite-nocfg-5s-Diffusers',
              repo_cls=diffusers.Kandinsky5T2VPipeline,
              te_cls=transformers.Qwen2_5_VLForConditionalGeneration,
              dit_cls=diffusers.Kandinsky5Transformer3DModel),
        Model(name='Kandinsky 5.0 Lite Steps-distilled T2V',
              url='https://huggingface.co/ai-forever/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers',
              repo='ai-forever/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers',
              repo_cls=diffusers.Kandinsky5T2VPipeline,
              te_cls=transformers.Qwen2_5_VLForConditionalGeneration,
              dit_cls=diffusers.Kandinsky5Transformer3DModel),
    ],
}