| import os |
| import json |
| import torch |
| from accelerate import init_empty_weights |
| from omegaconf import OmegaConf |
| from mmgp import offload |
| from shared.utils import files_locator as fl |
| from shared.utils.utils import convert_image_to_tensor, convert_tensor_to_image |
| from PIL import Image |
|
|
| from .kandinsky.models.dit import get_dit |
| from .kandinsky.models.text_embedders import get_text_embedder |
| from models.hyvideo.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D |
| from .kandinsky.pipeline import Kandinsky5Pipeline |
|
|
|
|
| def _resolve_repo_path(path): |
| if path is None: |
| return None |
| if os.path.isabs(path): |
| return path |
| base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) |
| return os.path.join(base_dir, path) |
|
|
|
|
| def _infer_task(base_model_type): |
| if not base_model_type: |
| return "t2v" |
| base = base_model_type.lower() |
| if "i2v" in base: |
| return "i2v" |
| if "t2v" in base: |
| return "t2v" |
| if "i2i" in base: |
| return "i2i" |
| if "t2i" in base: |
| return "t2i" |
| return "t2v" |
|
|
|
|
| def _get_config_path(base_model_type): |
| if not base_model_type: |
| raise ValueError("Missing base_model_type for Kandinsky config resolution.") |
| return os.path.join("models", "kandinsky5", "configs", f"{base_model_type}.yaml") |
|
|
|
|
| def _select_qwen_checkpoint(path_hint=None, folder_hint=None): |
| if path_hint is not None: |
| return path_hint |
| if folder_hint: |
| return fl.locate_folder(folder_hint) |
| candidates = [ |
| os.path.join("Qwen2.5-VL-7B-Instruct", "Qwen2.5-VL-7B-Instruct_bf16.safetensors"), |
| os.path.join("Qwen2.5-VL-7B-Instruct", "Qwen2.5-VL-7B-Instruct_quanto_bf16_int8.safetensors"), |
| ] |
| for candidate in candidates: |
| resolved = fl.locate_file(candidate, error_if_none=False) |
| if resolved is not None: |
| return resolved |
| return fl.locate_folder("Qwen2.5-VL-7B-Instruct") |
|
|
|
|
| def _preprocess_dit_state_dict(state_dict): |
| prefixes = ["model.diffusion_model.", "diffusion_model.", "model."] |
| for prefix in prefixes: |
| if any(key.startswith(prefix) for key in state_dict): |
| return {key[len(prefix):]: value for key, value in state_dict.items() if key.startswith(prefix)} |
| return state_dict |
|
|
|
|
| def _apply_vae_tile_config(vae, tile_config): |
| if tile_config is None: |
| if hasattr(vae, "disable_tiling"): |
| vae.disable_tiling() |
| return |
| if hasattr(vae, "apply_tile_config"): |
| vae.apply_tile_config(tile_config) |
| return |
| for key in ( |
| "tile_sample_min_tsize", |
| "tile_latent_min_tsize", |
| "tile_sample_min_size", |
| "tile_latent_min_size", |
| "tile_overlap_factor", |
| ): |
| if key in tile_config and hasattr(vae, key): |
| setattr(vae, key, tile_config[key]) |
| if hasattr(vae, "enable_tiling"): |
| vae.enable_tiling() |
|
|
|
|
| class model_factory: |
| def __init__( |
| self, |
| checkpoint_dir, |
| model_filename=None, |
| model_type=None, |
| model_def=None, |
| base_model_type=None, |
| text_encoder_filename=None, |
| quantizeTransformer=False, |
| save_quantized=False, |
| dtype=torch.bfloat16, |
| VAE_dtype=torch.float32, |
| mixed_precision_transformer=False, |
| **kwargs, |
| ): |
| self.device = torch.device("cuda") |
| load_device = torch.device("cpu") |
| self.dtype = dtype |
| self.VAE_dtype = VAE_dtype |
| self.model_def = model_def or {} |
| self.base_model_type = base_model_type or self.model_def.get("architecture") or model_type |
| self.mode = _infer_task(self.base_model_type) |
|
|
| config_path = _resolve_repo_path(_get_config_path(self.base_model_type)) |
| conf = OmegaConf.load(config_path) |
| overrides = self.model_def.get("k5_config_overrides") |
| if overrides: |
| conf = OmegaConf.merge(conf, OmegaConf.create(overrides)) |
|
|
| if isinstance(model_filename, (list, tuple)): |
| model_filename = model_filename[0] |
| conf.model.checkpoint_path = model_filename |
|
|
| text_encoder_folder = self.model_def.get("text_encoder_folder") |
| qwen_path = _select_qwen_checkpoint(text_encoder_filename, text_encoder_folder) |
| conf.model.text_embedder.qwen.checkpoint_path = qwen_path |
| conf.model.text_embedder.clip.checkpoint_path = fl.locate_folder("clip_vit_large_patch14") |
|
|
| vae_filename = fl.locate_file("hunyuan_video_VAE_fp32.safetensors") |
| vae_config = fl.locate_file("hunyuan_video_VAE_config.json") |
| conf.model.vae.checkpoint_path = vae_filename |
| conf.model.vae.name = "hunyuan" |
| conf.model.vae.config_path = vae_config |
|
|
| with init_empty_weights(): |
| dit = get_dit(conf.model.dit_params) |
| offload.load_model_data( |
| dit, |
| model_filename, |
| writable_tensors=False, |
| preprocess_sd=_preprocess_dit_state_dict, |
| ) |
| offload.change_dtype(dit, dtype, True) |
| dit.eval().requires_grad_(False) |
| if save_quantized: |
| from wgp import save_quantized_model |
| save_quantized_model(dit, model_type, model_filename, dtype, None) |
|
|
| quantized_qwen = isinstance(qwen_path, str) and "quanto" in qwen_path.lower() |
| text_embedder = get_text_embedder( |
| conf.model.text_embedder, device=load_device, quantized_qwen=quantized_qwen |
| ) |
|
|
| vae = offload.fast_load_transformers_model( |
| vae_filename, |
| writable_tensors=True, |
| modelClass=AutoencoderKLCausal3D, |
| defaultConfigPath=vae_config, |
| default_dtype=VAE_dtype, |
| ) |
| vae = vae.to(dtype=VAE_dtype, device=load_device).eval() |
| vae._model_dtype = VAE_dtype |
|
|
| device_map = {"dit": self.device, "vae": self.device, "text_embedder": self.device} |
| self.pipeline = Kandinsky5Pipeline( |
| mode=self.mode, |
| device_map=device_map, |
| dit=dit, |
| text_embedder=text_embedder, |
| vae=vae, |
| conf=conf, |
| ) |
| self._interrupt = False |
|
|
| self.transformer = dit |
| self.text_embedder = text_embedder |
| self.vae = vae |
|
|
| def generate( |
| self, |
| seed: int | None = None, |
| input_prompt: str = "", |
| n_prompt: str = "", |
| sampling_steps: int = 30, |
| guide_scale: float = 5.0, |
| frame_num: int = 81, |
| width: int = 768, |
| height: int = 512, |
| image_start=None, |
| image_mode: int = 0, |
| shift: float = 10.0, |
| callback=None, |
| progress: bool | None = None, |
| joint_pass: bool = False, |
| VAE_tile_size: dict | None = None, |
| **kwargs, |
| ): |
| if seed is not None and seed < 0: |
| seed = None |
| if progress is None: |
| progress = callback is None |
|
|
| negative_caption = n_prompt or "" |
| scheduler_scale = shift |
|
|
| if VAE_tile_size is not None: |
| _apply_vae_tile_config(self.vae, VAE_tile_size) |
| self.vae._use_vae_tiling = True |
| else: |
| _apply_vae_tile_config(self.vae, None) |
| self.vae._use_vae_tiling = False |
|
|
| image_pil = None |
| if image_start is not None: |
| if torch.is_tensor(image_start): |
| image_pil = convert_tensor_to_image(image_start) |
| elif isinstance(image_start, Image.Image): |
| image_pil = image_start |
| else: |
| raise ValueError(f"Unsupported image_start type: {type(image_start)}") |
|
|
| time_length = 0 if image_mode > 0 else 1 |
| outputs = self.pipeline( |
| text=input_prompt, |
| image=image_pil, |
| time_length=time_length, |
| width=width, |
| height=height, |
| frame_num=frame_num, |
| seed=seed, |
| num_steps=sampling_steps, |
| guidance_weight=guide_scale, |
| scheduler_scale=scheduler_scale, |
| negative_caption=negative_caption, |
| expand_prompts=False, |
| save_path=None, |
| progress=progress, |
| callback=callback, |
| joint_pass=joint_pass, |
| ) |
| if outputs is None: |
| return None |
|
|
| if isinstance(outputs, list): |
| frames = [convert_image_to_tensor(frame) for frame in outputs] |
| frames = torch.stack(frames, dim=0) |
| video = frames.permute(1, 0, 2, 3) |
| else: |
| if outputs.dim() == 5: |
| video = outputs[0].float().div_(127.5).sub_(1.0) |
| else: |
| video = outputs.float().div_(127.5).sub_(1.0) |
| return video |
|
|
| def get_loras_transformer(self, *args, **kwargs): |
| return [], [] |
|
|
| @property |
| def _interrupt(self): |
| if hasattr(self, "pipeline"): |
| return self.pipeline._interrupt |
| return False |
|
|
| @_interrupt.setter |
| def _interrupt(self, value): |
| if hasattr(self, "pipeline"): |
| self.pipeline._interrupt = value |
|
|
