| import os |
| from typing import Optional, Union |
|
|
| import numpy as np |
| import torch |
|
|
| from huggingface_hub import hf_hub_download, snapshot_download |
| from omegaconf import OmegaConf |
| from omegaconf.dictconfig import DictConfig |
|
|
| from .models.dit import get_dit |
| from .models.text_embedders import get_text_embedder |
| from .models.vae import build_vae |
| from .i2v_pipeline import Kandinsky5I2VPipeline |
| from .t2v_pipeline import Kandinsky5T2VPipeline |
| from .t2i_pipeline import Kandinsky5T2IPipeline |
| from .i2i_pipeline import Kandinsky5I2IPipeline |
| from .magcache_utils import set_magcache_params |
|
|
| from PIL import Image |
| from safetensors.torch import load_file |
|
|
| torch._dynamo.config.suppress_errors = True |
|
|
|
|
| HF_TOKEN = None |
|
|
|
|
| def get_hf_token(): |
| return HF_TOKEN |
|
|
|
|
| def set_hf_token(hf_token): |
| global HF_TOKEN |
| HF_TOKEN = hf_token |
|
|
|
|
| def get_T2V_pipeline( |
| device_map: Union[str, torch.device, dict], |
| cache_dir: str = "./weights/", |
| dit_path: str = None, |
| text_encoder_path: str = None, |
| text_encoder2_path: str = None, |
| vae_path: str = None, |
| conf_path: str = None, |
| magcache: bool = False, |
| quantized_qwen: bool = False, |
| text_token_padding: bool = False, |
| attention_engine: str = "auto", |
| ) -> Kandinsky5T2VPipeline: |
| if not isinstance(device_map, dict): |
| device_map = {"dit": device_map, "vae": device_map, "text_embedder": device_map} |
|
|
| os.makedirs(cache_dir, exist_ok=True) |
|
|
| if dit_path is None and conf_path is None: |
| dit_path = snapshot_download( |
| repo_id="kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s", |
| allow_patterns="model/*", |
| local_dir=cache_dir, |
| token=get_hf_token() |
| ) |
| dit_path = os.path.join(cache_dir, "model/kandinsky5lite_t2v_sft_5s.safetensors") |
|
|
| if vae_path is None and conf_path is None: |
| vae_path = snapshot_download( |
| repo_id="hunyuanvideo-community/HunyuanVideo", |
| allow_patterns="vae/*", |
| local_dir=cache_dir, |
| token=get_hf_token() |
| ) |
| vae_path = os.path.join(cache_dir, "vae/") |
|
|
| if text_encoder_path is None and conf_path is None: |
| text_encoder_path = snapshot_download( |
| repo_id="Qwen/Qwen2.5-VL-7B-Instruct", |
| local_dir=os.path.join(cache_dir, "text_encoder/"), |
| token=get_hf_token() |
| ) |
| text_encoder_path = os.path.join(cache_dir, "text_encoder/") |
|
|
| if text_encoder2_path is None and conf_path is None: |
| text_encoder2_path = snapshot_download( |
| repo_id="openai/clip-vit-large-patch14", |
| local_dir=os.path.join(cache_dir, "text_encoder2/"), |
| token=get_hf_token() |
| ) |
| text_encoder2_path = os.path.join(cache_dir, "text_encoder2/") |
|
|
| if conf_path is None: |
| conf = get_default_conf( |
| dit_path, vae_path, text_encoder_path, text_encoder2_path |
| ) |
| else: |
| conf = OmegaConf.load(conf_path) |
| conf.model.dit_params.attention_engine = attention_engine |
|
|
| conf.model.text_embedder.qwen.mode = "t2v" |
| text_embedder = get_text_embedder(conf.model.text_embedder, device=torch.device("cpu"), |
| quantized_qwen=quantized_qwen, text_token_padding=text_token_padding) |
|
|
| vae = build_vae(conf.model.vae) |
| vae = vae.eval() |
|
|
| dit = get_dit(conf.model.dit_params, text_token_padding=text_token_padding) |
|
|
| if magcache: |
| mag_ratios = conf.magcache.mag_ratios |
| num_steps = conf.model.num_steps |
| no_cfg = False |
| if conf.model.guidance_weight == 1.0: |
| no_cfg = True |
| set_magcache_params(dit, mag_ratios, num_steps, no_cfg, start_step=0) |
|
|
| state_dict = load_file(conf.model.checkpoint_path, device='cpu') |
| dit.load_state_dict(state_dict, assign=True) |
|
|
| return Kandinsky5T2VPipeline( |
| device_map=device_map, |
| dit=dit, |
| text_embedder=text_embedder, |
| vae=vae, |
| conf=conf, |
| ) |
|
|
| def get_I2V_pipeline( |
| device_map: Union[str, torch.device, dict], |
| cache_dir: str = "./weights/", |
| dit_path: str = None, |
| text_encoder_path: str = None, |
| text_encoder2_path: str = None, |
| vae_path: str = None, |
| conf_path: str = None, |
| magcache: bool = False, |
| quantized_qwen: bool = False, |
| text_token_padding: bool = False, |
| attention_engine: str = "auto", |
| ) -> Kandinsky5T2VPipeline: |
| if not isinstance(device_map, dict): |
| device_map = {"dit": device_map, "vae": device_map, "text_embedder": device_map} |
|
|
| os.makedirs(cache_dir, exist_ok=True) |
|
|
| if dit_path is None and conf_path is None: |
| dit_path = snapshot_download( |
| repo_id="kandinskylab/Kandinsky-5.0-I2V-Lite-5s", |
| allow_patterns="model/*", |
| local_dir=cache_dir, |
| token=get_hf_token() |
| ) |
| dit_path = os.path.join(cache_dir, "model/kandinsky5lite_i2v_5s.safetensors") |
|
|
| if vae_path is None and conf_path is None: |
| vae_path = snapshot_download( |
| repo_id="hunyuanvideo-community/HunyuanVideo", |
| allow_patterns="vae/*", |
| local_dir=cache_dir, |
| token=get_hf_token() |
| ) |
| vae_path = os.path.join(cache_dir, "vae/") |
|
|
| if text_encoder_path is None and conf_path is None: |
| text_encoder_path = snapshot_download( |
| repo_id="Qwen/Qwen2.5-VL-7B-Instruct", |
| local_dir=os.path.join(cache_dir, "text_encoder/"), |
| token=get_hf_token() |
| ) |
| text_encoder_path = os.path.join(cache_dir, "text_encoder/") |
|
|
| if text_encoder2_path is None and conf_path is None: |
| text_encoder2_path = snapshot_download( |
| repo_id="openai/clip-vit-large-patch14", |
| local_dir=os.path.join(cache_dir, "text_encoder2/"), |
| token=get_hf_token() |
| ) |
| text_encoder2_path = os.path.join(cache_dir, "text_encoder2/") |
|
|
| if conf_path is None: |
| conf = get_default_conf( |
| dit_path, vae_path, text_encoder_path, text_encoder2_path |
| ) |
| else: |
| conf = OmegaConf.load(conf_path) |
| conf.model.dit_params.attention_engine = attention_engine |
|
|
| conf.model.text_embedder.qwen.mode = "i2v" |
| text_embedder = get_text_embedder(conf.model.text_embedder, device=torch.device("cpu"), |
| quantized_qwen=quantized_qwen, text_token_padding=text_token_padding) |
| |
| vae = build_vae(conf.model.vae) |
| vae = vae.eval() |
|
|
| dit = get_dit(conf.model.dit_params, text_token_padding=text_token_padding) |
|
|
| if magcache: |
| mag_ratios = conf.magcache.mag_ratios |
| num_steps = conf.model.num_steps |
| no_cfg = False |
| if conf.model.guidance_weight == 1.0: |
| no_cfg = True |
| set_magcache_params(dit, mag_ratios, num_steps, no_cfg, start_step=0) |
|
|
| state_dict = load_file(conf.model.checkpoint_path, device='cpu') |
| dit.load_state_dict(state_dict, assign=True) |
|
|
| return Kandinsky5I2VPipeline( |
| device_map=device_map, |
| dit=dit, |
| text_embedder=text_embedder, |
| vae=vae, |
| conf=conf, |
| ) |
|
|
|
|
| def _get_TI2I_params( |
| instruct_type: bool, |
| model_name: str, |
| weights_name: str, |
| device_map: Union[str, torch.device, dict], |
| resolution: int = 1024, |
| cache_dir: str = "./weights/", |
| dit_path: str = None, |
| text_encoder_path: str = None, |
| text_encoder2_path: str = None, |
| vae_path: str = None, |
| conf_path: str = None, |
| magcache: bool = False, |
| quantized_qwen: bool = False, |
| text_token_padding: bool = False, |
| attention_engine: str = "auto", |
| ) -> Kandinsky5T2IPipeline: |
| assert resolution in [1024] |
|
|
| if not isinstance(device_map, dict): |
| device_map = {"dit": device_map, "vae": device_map, "text_embedder": device_map} |
|
|
| os.makedirs(cache_dir, exist_ok=True) |
|
|
| if dit_path is None and conf_path is None: |
| dit_path = snapshot_download( |
| repo_id=f"kandinskylab/{model_name}", |
| allow_patterns="model/*", |
| local_dir=cache_dir, |
| token=get_hf_token() |
| ) |
| dit_path = os.path.join(cache_dir, f"model/{weights_name}") |
|
|
| if vae_path is None and conf_path is None: |
| vae_path = snapshot_download( |
| repo_id="hunyuanvideo-community/HunyuanVideo", |
| allow_patterns="vae/*", |
| local_dir=cache_dir, |
| token=get_hf_token() |
| ) |
| vae_path = os.path.join(cache_dir, "vae/") |
|
|
| if text_encoder_path is None and conf_path is None: |
| text_encoder_path = snapshot_download( |
| repo_id="Qwen/Qwen2.5-VL-7B-Instruct", |
| local_dir=os.path.join(cache_dir, "text_encoder/"), |
| token=get_hf_token() |
| ) |
| text_encoder_path = os.path.join(cache_dir, "text_encoder/") |
|
|
| if text_encoder2_path is None and conf_path is None: |
| text_encoder2_path = snapshot_download( |
| repo_id="openai/clip-vit-large-patch14", |
| local_dir=os.path.join(cache_dir, "text_encoder2/"), |
| token=get_hf_token() |
| ) |
| text_encoder2_path = os.path.join(cache_dir, "text_encoder2/") |
|
|
| if conf_path is None: |
| conf = get_default_ti2i_conf( |
| dit_path, vae_path, text_encoder_path, text_encoder2_path, instruct_type=instruct_type, |
| ) |
| else: |
| conf = OmegaConf.load(conf_path) |
| conf.model.dit_params.attention_engine = attention_engine |
|
|
| conf.model.text_embedder.qwen.mode = "t2i" |
| text_embedder = get_text_embedder(conf.model.text_embedder, device=torch.device("cpu"), |
| quantized_qwen=quantized_qwen, text_token_padding=text_token_padding) |
|
|
| vae = build_vae(conf.model.vae) |
| vae = vae.eval() |
|
|
| dit = get_dit(conf.model.dit_params, text_token_padding=text_token_padding) |
|
|
| if magcache: |
| mag_ratios = conf.magcache.mag_ratios |
| num_steps = conf.model.num_steps |
| no_cfg = False |
| if conf.model.guidance_weight == 1.0: |
| no_cfg = True |
| set_magcache_params(dit, mag_ratios, num_steps, no_cfg, start_step=0) |
|
|
| state_dict = load_file(conf.model.checkpoint_path) |
| dit.load_state_dict(state_dict, assign=True) |
|
|
| return dict( |
| device_map=device_map, |
| dit=dit, |
| text_embedder=text_embedder, |
| vae=vae, |
| resolution=resolution, |
| conf=conf, |
| ) |
|
|
|
|
| def get_T2I_pipeline( |
| device_map: Union[str, torch.device, dict], |
| resolution: int = 1024, |
| cache_dir: str = "./weights/", |
| dit_path: str = None, |
| text_encoder_path: str = None, |
| text_encoder2_path: str = None, |
| vae_path: str = None, |
| conf_path: str = None, |
| magcache: bool = False, |
| quantized_qwen: bool = False, |
| text_token_padding: bool = False, |
| attention_engine: str = "auto", |
| ) -> Kandinsky5T2IPipeline: |
| kwargs = _get_TI2I_params( |
| instruct_type=None, |
| model_name='Kandinsky-5.0-T2I-Lite', |
| weights_name='kandinsky5lite_t2i.safetensors', |
| device_map=device_map, |
| resolution=resolution, |
| cache_dir=cache_dir, |
| dit_path=dit_path, |
| text_encoder_path=text_encoder_path, |
| text_encoder2_path=text_encoder2_path, |
| vae_path=vae_path, |
| conf_path=conf_path, |
| magcache=magcache, |
| quantized_qwen=quantized_qwen, |
| text_token_padding=text_token_padding, |
| attention_engine=attention_engine, |
| ) |
|
|
| return Kandinsky5T2IPipeline(**kwargs) |
|
|
|
|
| def get_I2I_pipeline( |
| device_map: Union[str, torch.device, dict], |
| resolution: int = 1024, |
| cache_dir: str = "./weights/", |
| dit_path: str = None, |
| text_encoder_path: str = None, |
| text_encoder2_path: str = None, |
| vae_path: str = None, |
| conf_path: str = None, |
| magcache: bool = False, |
| quantized_qwen: bool = False, |
| text_token_padding: bool = False, |
| attention_engine: str = "auto", |
| ) -> Kandinsky5T2IPipeline: |
| kwargs = _get_TI2I_params( |
| instruct_type='channel', |
| model_name='Kandinsky-5.0-I2I-Lite', |
| weights_name='kandinsky5lite_i2i.safetensors', |
| device_map=device_map, |
| resolution=resolution, |
| cache_dir=cache_dir, |
| dit_path=dit_path, |
| text_encoder_path=text_encoder_path, |
| text_encoder2_path=text_encoder2_path, |
| vae_path=vae_path, |
| conf_path=conf_path, |
| magcache=magcache, |
| quantized_qwen=quantized_qwen, |
| text_token_padding=text_token_padding, |
| attention_engine=attention_engine, |
| ) |
|
|
| return Kandinsky5I2IPipeline(**kwargs) |
|
|
|
|
| def get_default_conf( |
| dit_path, |
| vae_path, |
| text_encoder_path, |
| text_encoder2_path, |
| ) -> DictConfig: |
| dit_params = { |
| "in_visual_dim": 16, |
| "out_visual_dim": 16, |
| "time_dim": 512, |
| "patch_size": [1, 2, 2], |
| "model_dim": 1792, |
| "ff_dim": 7168, |
| "num_text_blocks": 2, |
| "num_visual_blocks": 32, |
| "axes_dims": [16, 24, 24], |
| "visual_cond": True, |
| "in_text_dim": 3584, |
| "in_text_dim2": 768, |
| } |
|
|
| attention = { |
| "type": "flash", |
| "causal": False, |
| "local": False, |
| "glob": False, |
| "window": 3, |
| } |
|
|
| vae = { |
| "checkpoint_path": vae_path, |
| "name": "hunyuan", |
| } |
|
|
| text_embedder = { |
| "qwen": { |
| "emb_size": 3584, |
| "checkpoint_path": text_encoder_path, |
| "max_length": 256, |
| }, |
| "clip": { |
| "checkpoint_path": text_encoder2_path, |
| "emb_size": 768, |
| "max_length": 77, |
| }, |
| } |
|
|
| conf = { |
| "model": { |
| "checkpoint_path": dit_path, |
| "vae": vae, |
| "text_embedder": text_embedder, |
| "dit_params": dit_params, |
| "attention": attention, |
| "num_steps": 50, |
| "guidance_weight": 5.0, |
| }, |
| "metrics": {"scale_factor": (1, 2, 2), "resolution": 512,}, |
| } |
|
|
| return DictConfig(conf) |
|
|
|
|
| def get_default_ti2i_conf( |
| dit_path, |
| vae_path, |
| text_encoder_path, |
| text_encoder2_path, |
| instruct_type=None, |
| ) -> DictConfig: |
| dit_params = { |
| "instruct_type": instruct_type, |
| "in_visual_dim": 16, |
| "out_visual_dim": 16, |
| "time_dim": 512, |
| "patch_size": [1, 2, 2], |
| "model_dim": 2560, |
| "ff_dim": 10240, |
| "num_text_blocks": 2, |
| "num_visual_blocks": 50, |
| "axes_dims": [32,48, 48], |
| } |
|
|
| attention = { |
| "type": "flash", |
| "causal": False, |
| "local": False, |
| "glob": False, |
| "window": 3, |
| } |
|
|
| vae = { |
| "checkpoint_path": vae_path, |
| "name": "hunyuan", |
| } |
|
|
| text_embedder = { |
| "qwen": { |
| "emb_size": 3584, |
| "checkpoint_path": text_encoder_path, |
| "max_length": 512, |
| }, |
| "clip": { |
| "checkpoint_path": text_encoder2_path, |
| "emb_size": 768, |
| "max_length": 77, |
| }, |
| } |
|
|
| conf = { |
| "model": { |
| "checkpoint_path": dit_path, |
| "vae": vae, |
| "text_embedder": text_embedder, |
| "dit_params": dit_params, |
| "attention": attention, |
| "num_steps": 50, |
| "guidance_weight": 3.5, |
| }, |
| "metrics": {"scale_factor": (1, 2, 2), "resolution": 1024}, |
| } |
|
|
| return DictConfig(conf) |
|
|
