models/kandinsky5/kandinsky/utils.py

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)