diff --git a/.gitignore b/.gitignore index d95eb33..b42c7cd 100644 --- a/.gitignore +++ b/.gitignore @@ -14,7 +14,7 @@ *.pth *.ckpt *.safetensors -*.json +#*.json # *.txt *.backup *.pkl diff --git a/README.md b/README.md index 0ba6bd6..ceeb7a4 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,17 @@ WanGP supports the Wan (and derived models), Hunyuan Video and LTV Video models **Follow DeepBeepMeep on Twitter/X to get the Latest News**: https://x.com/deepbeepmeep ## 🔥 Latest Updates : +### August 6 2025: WanGP v7.7 - Picky, picky + +This release comes with two new models : +- Qwen Image: a Commercial grade Image generator capable to inject full sentences in the generated Image while still offering incredible visuals +- Wan 2.2 TextImage to Video 5B: the last Wan 2.2 needed if you want to complete your Wan 2.2 collection (loras for this folder can be stored in "\loras\5B" ) + +There is catch though, they are very picky if you want to get good generations: first they both need lots of steps (50 ?) to show what they have to offer. Then for Qwen Image I had to hardcode the supported resolutions, because if you try anything else, you will get garbage. Likiwise Wan 2.2 5B will remind you of Wan 1.0 if you don't ask for at least 720p. + +Please note that the VAE decoding of Wan 2.2 TextImage is not tiled yet and it may produce VRAM consumption peaks (this doens't mix well with the 720p requirement). + + ### August 4 2025: WanGP v7.6 - Remuxed With this new version you won't have any excuse if there is no sound in your video. diff --git a/configs/qwen_image_20B.json b/configs/qwen_image_20B.json new file mode 100644 index 0000000..4bff1e5 --- /dev/null +++ b/configs/qwen_image_20B.json @@ -0,0 +1,18 @@ +{ + "_class_name": "QwenImageTransformer2DModel", + "_diffusers_version": "0.34.0.dev0", + "attention_head_dim": 128, + "axes_dims_rope": [ + 16, + 56, + 56 + ], + "guidance_embeds": false, + "in_channels": 64, + "joint_attention_dim": 3584, + "num_attention_heads": 24, + "num_layers": 60, + "out_channels": 16, + "patch_size": 2, + "pooled_projection_dim": 768 +} diff --git a/configs/ti2v_2_2.json b/configs/ti2v_2_2.json new file mode 100644 index 0000000..d58edcc --- /dev/null +++ b/configs/ti2v_2_2.json @@ -0,0 +1,14 @@ +{ + "_class_name": "WanModel", + "_diffusers_version": "0.33.0", + "dim": 3072, + "eps": 1e-06, + "ffn_dim": 14336, + "freq_dim": 256, + "in_dim": 48, + "model_type": "ti2v2_2", + "num_heads": 24, + "num_layers": 30, + "out_dim": 48, + "text_len": 512 +} diff --git a/defaults/ltxv_13B.json b/defaults/ltxv_13B.json index dc61e31..639442e 100644 --- a/defaults/ltxv_13B.json +++ b/defaults/ltxv_13B.json @@ -13,7 +13,7 @@ "https://huggingface.co/DeepBeepMeep/LTX_Video/resolve/main/ltxv-097-ic-lora-depth-control-diffusers.safetensors", "https://huggingface.co/DeepBeepMeep/LTX_Video/resolve/main/ltxv-097-ic-lora-canny-control-diffusers.safetensors" ], - "LTXV_config": "ltx_video/configs/ltxv-13b-0.9.8-dev.yaml" + "LTXV_config": "models/ltx_video/configs/ltxv-13b-0.9.8-dev.yaml" }, "num_inference_steps": 30 } diff --git a/defaults/ltxv_distilled.json b/defaults/ltxv_distilled.json index 8973b11..c570057 100644 --- a/defaults/ltxv_distilled.json +++ b/defaults/ltxv_distilled.json @@ -9,7 +9,7 @@ "https://huggingface.co/DeepBeepMeep/LTX_Video/resolve/main/ltxv_0.9.8_13B_distilled_quanto_bf16_int8.safetensors" ], "preload_URLs" : "ltxv_13B", - "LTXV_config": "ltx_video/configs/ltxv-13b-0.9.8-distilled.yaml" + "LTXV_config": "models/ltx_video/configs/ltxv-13b-0.9.8-distilled.yaml" }, "num_inference_steps": 6 } diff --git a/defaults/qwen_image_20B.json b/defaults/qwen_image_20B.json new file mode 100644 index 0000000..7abe026 --- /dev/null +++ b/defaults/qwen_image_20B.json @@ -0,0 +1,21 @@ +{ + "model": { + "name": "Qwen Image 20B", + "architecture": "qwen_image_20B", + "description": "Qwen Image is generative model that will very high quality images. It is one of the few models capable to generate in the image very long texts.", + "URLs": [ + "https://huggingface.co/DeepBeepMeep/Qwen/resolve/main/qwen_image_20B_bf16.safetensors", + "https://huggingface.co/DeepBeepMeep/Qwen/resolve/main/qwen_image_20B_quanto_bf16_int8.safetensors" + ], + "resolutions": [ ["1328x1328 (1:1)", "1328x1328"], + ["1664x928 (16:9)", "1664x928"], + ["928x1664 (9:16)", "928x1664"], + ["1472x1140 (4:3)", "1472x1140"], + ["1140x1472 (3:4)", "1140x1472"] + ], + "image_outputs": true + }, + "prompt": "draw a hat", + "resolution": "1280x720", + "batch_size": 1 +} \ No newline at end of file diff --git a/flux/__init__.py b/flux/__init__.py deleted file mode 100644 index dddc6a3..0000000 --- a/flux/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -try: - from ._version import ( - version as __version__, # type: ignore - version_tuple, - ) -except ImportError: - __version__ = "unknown (no version information available)" - version_tuple = (0, 0, "unknown", "noinfo") - -from pathlib import Path - -PACKAGE = __package__.replace("_", "-") -PACKAGE_ROOT = Path(__file__).parent diff --git a/loras_qwen/Readme.txt b/loras_qwen/Readme.txt new file mode 100644 index 0000000..14a70a8 --- /dev/null +++ b/loras_qwen/Readme.txt @@ -0,0 +1 @@ +LTX Video loras \ No newline at end of file diff --git a/hyvideo/__init__.py b/models/__init__.py similarity index 100% rename from hyvideo/__init__.py rename to models/__init__.py diff --git a/models/flux/__init__.py b/models/flux/__init__.py new file mode 100644 index 0000000..d0a07ae --- /dev/null +++ b/models/flux/__init__.py @@ -0,0 +1,2 @@ +from .flux_main import model_factory +from . import flux_handler diff --git a/flux/__main__.py b/models/flux/__main__.py similarity index 100% rename from flux/__main__.py rename to models/flux/__main__.py diff --git a/flux/_version.py b/models/flux/_version.py similarity index 100% rename from flux/_version.py rename to models/flux/_version.py diff --git a/models/flux/flux_handler.py b/models/flux/flux_handler.py new file mode 100644 index 0000000..a42f9a5 --- /dev/null +++ b/models/flux/flux_handler.py @@ -0,0 +1,103 @@ +import torch + +def get_ltxv_text_encoder_filename(text_encoder_quantization): + text_encoder_filename = "ckpts/T5_xxl_1.1/T5_xxl_1.1_enc_bf16.safetensors" + if text_encoder_quantization =="int8": + text_encoder_filename = text_encoder_filename.replace("bf16", "quanto_bf16_int8") + return text_encoder_filename + +class family_handler(): + @staticmethod + def query_model_def(base_model_type, model_def): + flux_model = model_def.get("flux-model", "flux-dev") + flux_schnell = flux_model == "flux-schnell" + model_def_output = { + "image_outputs" : True, + "no_negative_prompt" : True, + } + if flux_schnell: + model_def_output["no_guidance"] = True + else: + model_def_output["embedded_guidance"] = True + + + return model_def_output + + @staticmethod + def query_supported_types(): + return ["flux"] + + @staticmethod + def query_family_maps(): + return {}, {} + + @staticmethod + def get_rgb_factors(model_type): + from shared.RGB_factors import get_rgb_factors + latent_rgb_factors, latent_rgb_factors_bias = get_rgb_factors("flux") + return latent_rgb_factors, latent_rgb_factors_bias + + + @staticmethod + def query_model_family(): + return "flux" + + @staticmethod + def query_family_infos(): + return {"flux":(30, "Flux 1")} + + @staticmethod + def query_model_files(computeList, base_model_type, model_filename, text_encoder_quantization): + text_encoder_filename = get_ltxv_text_encoder_filename(text_encoder_quantization) + return [ + { + "repoId" : "DeepBeepMeep/Flux", + "sourceFolderList" : [""], + "fileList" : [ ["flux_vae.safetensors"] ] + }, + { + "repoId" : "DeepBeepMeep/LTX_Video", + "sourceFolderList" : ["T5_xxl_1.1"], + "fileList" : [ ["added_tokens.json", "special_tokens_map.json", "spiece.model", "tokenizer_config.json"] + computeList(text_encoder_filename) ] + }, + { + "repoId" : "DeepBeepMeep/HunyuanVideo", + "sourceFolderList" : [ "clip_vit_large_patch14", ], + "fileList" :[ + ["config.json", "merges.txt", "model.safetensors", "preprocessor_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json", "vocab.json"], + ] + } + ] + + @staticmethod + def load_model(model_filename, model_type, base_model_type, model_def, quantizeTransformer = False, text_encoder_quantization = None, dtype = torch.bfloat16, VAE_dtype = torch.float32, mixed_precision_transformer = False, save_quantized = False): + from .flux_main import model_factory + + flux_model = model_factory( + checkpoint_dir="ckpts", + model_filename=model_filename, + model_type = model_type, + model_def = model_def, + base_model_type=base_model_type, + text_encoder_filename= get_ltxv_text_encoder_filename(text_encoder_quantization), + quantizeTransformer = quantizeTransformer, + dtype = dtype, + VAE_dtype = VAE_dtype, + mixed_precision_transformer = mixed_precision_transformer, + save_quantized = save_quantized + ) + + pipe = { "transformer": flux_model.model, "vae" : flux_model.vae, "text_encoder" : flux_model.clip, "text_encoder_2" : flux_model.t5} + + return flux_model, pipe + + @staticmethod + def update_default_settings(base_model_type, model_def, ui_defaults): + ui_defaults.update({ + "embedded_guidance": 2.5, + }) + if model_def.get("reference_image", False): + ui_defaults.update({ + "video_prompt_type": "KI", + }) + diff --git a/flux/flux_main.py b/models/flux/flux_main.py similarity index 90% rename from flux/flux_main.py rename to models/flux/flux_main.py index 303765a..16659fc 100644 --- a/flux/flux_main.py +++ b/models/flux/flux_main.py @@ -5,10 +5,10 @@ from dataclasses import dataclass from glob import iglob from mmgp import offload as offload import torch -from wan.utils.utils import calculate_new_dimensions -from flux.sampling import denoise, get_schedule, prepare_kontext, unpack -from flux.modules.layers import get_linear_split_map -from flux.util import ( +from shared.utils.utils import calculate_new_dimensions +from .sampling import denoise, get_schedule, prepare_kontext, unpack +from .modules.layers import get_linear_split_map +from .util import ( aspect_ratio_to_height_width, load_ae, load_clip, @@ -146,13 +146,3 @@ class model_factory: x = x.transpose(0, 1) return x -def query_model_def(model_type, model_def): - flux_model = model_def.get("flux-model", "flux-dev") - flux_schnell = flux_model == "flux-schnell" - model_def_output = { - "image_outputs" : True, - } - if flux_schnell: - model_def_output["no_guidance"] = True - - return model_def_output \ No newline at end of file diff --git a/flux/math.py b/models/flux/math.py similarity index 97% rename from flux/math.py rename to models/flux/math.py index 9e8aa59..a249f19 100644 --- a/flux/math.py +++ b/models/flux/math.py @@ -1,7 +1,7 @@ import torch from einops import rearrange from torch import Tensor -from wan.modules.attention import pay_attention +from shared.attention import pay_attention def attention(qkv_list, pe: Tensor) -> Tensor: diff --git a/flux/model.py b/models/flux/model.py similarity index 98% rename from flux/model.py rename to models/flux/model.py index d6c1b6c..d84ceb3 100644 --- a/flux/model.py +++ b/models/flux/model.py @@ -3,7 +3,7 @@ from dataclasses import dataclass import torch from torch import Tensor, nn -from flux.modules.layers import ( +from .modules.layers import ( DoubleStreamBlock, EmbedND, LastLayer, @@ -11,7 +11,7 @@ from flux.modules.layers import ( SingleStreamBlock, timestep_embedding, ) -from flux.modules.lora import LinearLora, replace_linear_with_lora +from .modules.lora import LinearLora, replace_linear_with_lora @dataclass diff --git a/flux/modules/autoencoder.py b/models/flux/modules/autoencoder.py similarity index 100% rename from flux/modules/autoencoder.py rename to models/flux/modules/autoencoder.py diff --git a/flux/modules/conditioner.py b/models/flux/modules/conditioner.py similarity index 100% rename from flux/modules/conditioner.py rename to models/flux/modules/conditioner.py diff --git a/flux/modules/image_embedders.py b/models/flux/modules/image_embedders.py similarity index 98% rename from flux/modules/image_embedders.py rename to models/flux/modules/image_embedders.py index aa26d9b..011f840 100644 --- a/flux/modules/image_embedders.py +++ b/models/flux/modules/image_embedders.py @@ -7,7 +7,7 @@ from safetensors.torch import load_file as load_sft from torch import nn from transformers import AutoModelForDepthEstimation, AutoProcessor, SiglipImageProcessor, SiglipVisionModel -from flux.util import print_load_warning +from ..util import print_load_warning class DepthImageEncoder: diff --git a/flux/modules/layers copy.py b/models/flux/modules/layers copy.py similarity index 100% rename from flux/modules/layers copy.py rename to models/flux/modules/layers copy.py diff --git a/flux/modules/layers.py b/models/flux/modules/layers.py similarity index 99% rename from flux/modules/layers.py rename to models/flux/modules/layers.py index d0273b7..b937143 100644 --- a/flux/modules/layers.py +++ b/models/flux/modules/layers.py @@ -5,7 +5,7 @@ import torch from einops import rearrange from torch import Tensor, nn -from flux.math import attention, rope +from ..math import attention, rope def get_linear_split_map(): hidden_size = 3072 diff --git a/flux/modules/lora.py b/models/flux/modules/lora.py similarity index 100% rename from flux/modules/lora.py rename to models/flux/modules/lora.py diff --git a/flux/sampling.py b/models/flux/sampling.py similarity index 99% rename from flux/sampling.py rename to models/flux/sampling.py index 5a15c5e..23cfcf3 100644 --- a/flux/sampling.py +++ b/models/flux/sampling.py @@ -343,7 +343,7 @@ def denoise( updated_num_steps= len(timesteps) -1 if callback != None: - from wan.utils.loras_mutipliers import update_loras_slists + from shared.utils.loras_mutipliers import update_loras_slists update_loras_slists(model, loras_slists, updated_num_steps) callback(-1, None, True, override_num_inference_steps = updated_num_steps) from mmgp import offload diff --git a/flux/util.py b/models/flux/util.py similarity index 99% rename from flux/util.py rename to models/flux/util.py index a9815af..a729dd7 100644 --- a/flux/util.py +++ b/models/flux/util.py @@ -11,9 +11,9 @@ from huggingface_hub import hf_hub_download, login from PIL import ExifTags, Image from safetensors.torch import load_file as load_sft -from flux.model import Flux, FluxLoraWrapper, FluxParams -from flux.modules.autoencoder import AutoEncoder, AutoEncoderParams -from flux.modules.conditioner import HFEmbedder +from .model import Flux, FluxLoraWrapper, FluxParams +from .modules.autoencoder import AutoEncoder, AutoEncoderParams +from .modules.conditioner import HFEmbedder CHECKPOINTS_DIR = Path("checkpoints") diff --git a/models/hyvideo/__init__.py b/models/hyvideo/__init__.py new file mode 100644 index 0000000..d3a3700 --- /dev/null +++ b/models/hyvideo/__init__.py @@ -0,0 +1,2 @@ +from .hunyuan import HunyuanVideoSampler +from . import hunyuan_handler \ No newline at end of file diff --git a/hyvideo/config.py b/models/hyvideo/config.py similarity index 100% rename from hyvideo/config.py rename to models/hyvideo/config.py diff --git a/hyvideo/constants.py b/models/hyvideo/constants.py similarity index 100% rename from hyvideo/constants.py rename to models/hyvideo/constants.py diff --git a/hyvideo/data_kits/audio_dataset.py b/models/hyvideo/data_kits/audio_dataset.py similarity index 100% rename from hyvideo/data_kits/audio_dataset.py rename to models/hyvideo/data_kits/audio_dataset.py diff --git a/hyvideo/data_kits/audio_preprocessor.py b/models/hyvideo/data_kits/audio_preprocessor.py similarity index 100% rename from hyvideo/data_kits/audio_preprocessor.py rename to models/hyvideo/data_kits/audio_preprocessor.py diff --git a/hyvideo/data_kits/data_tools.py b/models/hyvideo/data_kits/data_tools.py similarity index 100% rename from hyvideo/data_kits/data_tools.py rename to models/hyvideo/data_kits/data_tools.py diff --git a/hyvideo/data_kits/face_align/__init__.py b/models/hyvideo/data_kits/face_align/__init__.py similarity index 100% rename from hyvideo/data_kits/face_align/__init__.py rename to models/hyvideo/data_kits/face_align/__init__.py diff --git a/hyvideo/data_kits/face_align/align.py b/models/hyvideo/data_kits/face_align/align.py similarity index 100% rename from hyvideo/data_kits/face_align/align.py rename to models/hyvideo/data_kits/face_align/align.py diff --git a/hyvideo/data_kits/face_align/detface.py b/models/hyvideo/data_kits/face_align/detface.py similarity index 100% rename from hyvideo/data_kits/face_align/detface.py rename to models/hyvideo/data_kits/face_align/detface.py diff --git a/hyvideo/diffusion/__init__.py b/models/hyvideo/diffusion/__init__.py similarity index 100% rename from hyvideo/diffusion/__init__.py rename to models/hyvideo/diffusion/__init__.py diff --git a/hyvideo/diffusion/pipelines/__init__.py b/models/hyvideo/diffusion/pipelines/__init__.py similarity index 100% rename from hyvideo/diffusion/pipelines/__init__.py rename to models/hyvideo/diffusion/pipelines/__init__.py diff --git a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py b/models/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py similarity index 100% rename from hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py rename to models/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py diff --git a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video_audio.py b/models/hyvideo/diffusion/pipelines/pipeline_hunyuan_video_audio.py similarity index 99% rename from hyvideo/diffusion/pipelines/pipeline_hunyuan_video_audio.py rename to models/hyvideo/diffusion/pipelines/pipeline_hunyuan_video_audio.py index 191f9ab..fd36f03 100644 --- a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video_audio.py +++ b/models/hyvideo/diffusion/pipelines/pipeline_hunyuan_video_audio.py @@ -41,9 +41,9 @@ from diffusers.utils import ( from diffusers.utils.torch_utils import randn_tensor from diffusers.pipelines.pipeline_utils import DiffusionPipeline -from hyvideo.constants import PRECISION_TO_TYPE -from hyvideo.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D -from hyvideo.text_encoder import TextEncoder +from ...constants import PRECISION_TO_TYPE +from ...vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D +from ...text_encoder import TextEncoder from einops import rearrange from ...modules import HYVideoDiffusionTransformer diff --git a/hyvideo/diffusion/schedulers/__init__.py b/models/hyvideo/diffusion/schedulers/__init__.py similarity index 100% rename from hyvideo/diffusion/schedulers/__init__.py rename to models/hyvideo/diffusion/schedulers/__init__.py diff --git a/hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py b/models/hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py similarity index 100% rename from hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py rename to models/hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py diff --git a/hyvideo/hunyuan.py b/models/hyvideo/hunyuan.py similarity index 79% rename from hyvideo/hunyuan.py rename to models/hyvideo/hunyuan.py index 380ec77..a38a7bd 100644 --- a/hyvideo/hunyuan.py +++ b/models/hyvideo/hunyuan.py @@ -8,24 +8,24 @@ from pathlib import Path from einops import rearrange import torch import torch.distributed as dist -from hyvideo.constants import PROMPT_TEMPLATE, NEGATIVE_PROMPT, PRECISION_TO_TYPE, NEGATIVE_PROMPT_I2V -from hyvideo.vae import load_vae -from hyvideo.modules import load_model -from hyvideo.text_encoder import TextEncoder -from hyvideo.utils.data_utils import align_to, get_closest_ratio, generate_crop_size_list -from hyvideo.modules.posemb_layers import get_nd_rotary_pos_embed, get_nd_rotary_pos_embed_new -from hyvideo.diffusion.schedulers import FlowMatchDiscreteScheduler -from hyvideo.diffusion.pipelines import HunyuanVideoPipeline -from hyvideo.diffusion.pipelines import HunyuanVideoAudioPipeline +from .constants import PROMPT_TEMPLATE, NEGATIVE_PROMPT, PRECISION_TO_TYPE, NEGATIVE_PROMPT_I2V +from .vae import load_vae +from .modules import load_model +from .text_encoder import TextEncoder +from .utils.data_utils import align_to, get_closest_ratio, generate_crop_size_list +from .modules.posemb_layers import get_nd_rotary_pos_embed, get_nd_rotary_pos_embed_new +from .diffusion.schedulers import FlowMatchDiscreteScheduler +from .diffusion.pipelines import HunyuanVideoPipeline +from .diffusion.pipelines import HunyuanVideoAudioPipeline from PIL import Image import numpy as np import torchvision.transforms as transforms import cv2 -from wan.utils.utils import calculate_new_dimensions, convert_tensor_to_image -from hyvideo.data_kits.audio_preprocessor import encode_audio, get_facemask +from shared.utils.utils import calculate_new_dimensions, convert_tensor_to_image +from .data_kits.audio_preprocessor import encode_audio, get_facemask from transformers import WhisperModel from transformers import AutoFeatureExtractor -from hyvideo.data_kits.face_align import AlignImage +from .data_kits.face_align import AlignImage import librosa def get_audio_feature(feature_extractor, audio_path, duration): @@ -66,174 +66,174 @@ def pad_image(crop_img, size, color=(255, 255, 255), resize_ratio=1): -def _merge_input_ids_with_image_features(self, image_features, inputs_embeds, input_ids, attention_mask, labels): - num_images, num_image_patches, embed_dim = image_features.shape - batch_size, sequence_length = input_ids.shape - left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(self.pad_token_id)) - # 1. Create a mask to know where special image tokens are - special_image_token_mask = input_ids == self.config.image_token_index - num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1) - # Compute the maximum embed dimension - max_embed_dim = (num_special_image_tokens.max() * (num_image_patches - 1)) + sequence_length - batch_indices, non_image_indices = torch.where(input_ids != self.config.image_token_index) +# def _merge_input_ids_with_image_features(self, image_features, inputs_embeds, input_ids, attention_mask, labels): +# num_images, num_image_patches, embed_dim = image_features.shape +# batch_size, sequence_length = input_ids.shape +# left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(self.pad_token_id)) +# # 1. Create a mask to know where special image tokens are +# special_image_token_mask = input_ids == self.config.image_token_index +# num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1) +# # Compute the maximum embed dimension +# max_embed_dim = (num_special_image_tokens.max() * (num_image_patches - 1)) + sequence_length +# batch_indices, non_image_indices = torch.where(input_ids != self.config.image_token_index) - # 2. Compute the positions where text should be written - # Calculate new positions for text tokens in merged image-text sequence. - # `special_image_token_mask` identifies image tokens. Each image token will be replaced by `nb_text_tokens_per_images - 1` text tokens. - # `torch.cumsum` computes how each image token shifts subsequent text token positions. - # - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one. - new_token_positions = torch.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1) - 1 - nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] - if left_padding: - new_token_positions += nb_image_pad[:, None] # offset for left padding - text_to_overwrite = new_token_positions[batch_indices, non_image_indices] +# # 2. Compute the positions where text should be written +# # Calculate new positions for text tokens in merged image-text sequence. +# # `special_image_token_mask` identifies image tokens. Each image token will be replaced by `nb_text_tokens_per_images - 1` text tokens. +# # `torch.cumsum` computes how each image token shifts subsequent text token positions. +# # - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one. +# new_token_positions = torch.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1) - 1 +# nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] +# if left_padding: +# new_token_positions += nb_image_pad[:, None] # offset for left padding +# text_to_overwrite = new_token_positions[batch_indices, non_image_indices] - # 3. Create the full embedding, already padded to the maximum position - final_embedding = torch.zeros( - batch_size, max_embed_dim, embed_dim, dtype=inputs_embeds.dtype, device=inputs_embeds.device - ) - final_attention_mask = torch.zeros( - batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device - ) - if labels is not None: - final_labels = torch.full( - (batch_size, max_embed_dim), self.config.ignore_index, dtype=input_ids.dtype, device=input_ids.device - ) - # In case the Vision model or the Language model has been offloaded to CPU, we need to manually - # set the corresponding tensors into their correct target device. - target_device = inputs_embeds.device - batch_indices, non_image_indices, text_to_overwrite = ( - batch_indices.to(target_device), - non_image_indices.to(target_device), - text_to_overwrite.to(target_device), - ) - attention_mask = attention_mask.to(target_device) +# # 3. Create the full embedding, already padded to the maximum position +# final_embedding = torch.zeros( +# batch_size, max_embed_dim, embed_dim, dtype=inputs_embeds.dtype, device=inputs_embeds.device +# ) +# final_attention_mask = torch.zeros( +# batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device +# ) +# if labels is not None: +# final_labels = torch.full( +# (batch_size, max_embed_dim), self.config.ignore_index, dtype=input_ids.dtype, device=input_ids.device +# ) +# # In case the Vision model or the Language model has been offloaded to CPU, we need to manually +# # set the corresponding tensors into their correct target device. +# target_device = inputs_embeds.device +# batch_indices, non_image_indices, text_to_overwrite = ( +# batch_indices.to(target_device), +# non_image_indices.to(target_device), +# text_to_overwrite.to(target_device), +# ) +# attention_mask = attention_mask.to(target_device) - # 4. Fill the embeddings based on the mask. If we have ["hey" "", "how", "are"] - # we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the image features - final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] - final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices] - if labels is not None: - final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices] +# # 4. Fill the embeddings based on the mask. If we have ["hey" "", "how", "are"] +# # we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the image features +# final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] +# final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices] +# if labels is not None: +# final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices] - # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) - image_to_overwrite = torch.full( - (batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device - ) - image_to_overwrite[batch_indices, text_to_overwrite] = False - image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device) +# # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) +# image_to_overwrite = torch.full( +# (batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device +# ) +# image_to_overwrite[batch_indices, text_to_overwrite] = False +# image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device) - if image_to_overwrite.sum() != image_features.shape[:-1].numel(): - raise ValueError( - f"The input provided to the model are wrong. The number of image tokens is {torch.sum(special_image_token_mask)} while" - f" the number of image given to the model is {num_images}. This prevents correct indexing and breaks batch generation." - ) +# if image_to_overwrite.sum() != image_features.shape[:-1].numel(): +# raise ValueError( +# f"The input provided to the model are wrong. The number of image tokens is {torch.sum(special_image_token_mask)} while" +# f" the number of image given to the model is {num_images}. This prevents correct indexing and breaks batch generation." +# ) - final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device) - final_attention_mask |= image_to_overwrite - position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1) +# final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device) +# final_attention_mask |= image_to_overwrite +# position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1) - # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. - batch_indices, pad_indices = torch.where(input_ids == self.pad_token_id) - indices_to_mask = new_token_positions[batch_indices, pad_indices] +# # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. +# batch_indices, pad_indices = torch.where(input_ids == self.pad_token_id) +# indices_to_mask = new_token_positions[batch_indices, pad_indices] - final_embedding[batch_indices, indices_to_mask] = 0 +# final_embedding[batch_indices, indices_to_mask] = 0 - if labels is None: - final_labels = None +# if labels is None: +# final_labels = None - return final_embedding, final_attention_mask, final_labels, position_ids +# return final_embedding, final_attention_mask, final_labels, position_ids -def patched_llava_forward( - self, - input_ids: torch.LongTensor = None, - pixel_values: torch.FloatTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - vision_feature_layer: Optional[int] = None, - vision_feature_select_strategy: Optional[str] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, -): - from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast +# def patched_llava_forward( +# self, +# input_ids: torch.LongTensor = None, +# pixel_values: torch.FloatTensor = None, +# attention_mask: Optional[torch.Tensor] = None, +# position_ids: Optional[torch.LongTensor] = None, +# past_key_values: Optional[List[torch.FloatTensor]] = None, +# inputs_embeds: Optional[torch.FloatTensor] = None, +# vision_feature_layer: Optional[int] = None, +# vision_feature_select_strategy: Optional[str] = None, +# labels: Optional[torch.LongTensor] = None, +# use_cache: Optional[bool] = None, +# output_attentions: Optional[bool] = None, +# output_hidden_states: Optional[bool] = None, +# return_dict: Optional[bool] = None, +# cache_position: Optional[torch.LongTensor] = None, +# num_logits_to_keep: int = 0, +# ): +# from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) +# output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions +# output_hidden_states = ( +# output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states +# ) +# return_dict = return_dict if return_dict is not None else self.config.use_return_dict +# vision_feature_layer = ( +# vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer +# ) +# vision_feature_select_strategy = ( +# vision_feature_select_strategy +# if vision_feature_select_strategy is not None +# else self.config.vision_feature_select_strategy +# ) - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") +# if (input_ids is None) ^ (inputs_embeds is not None): +# raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if pixel_values is not None and inputs_embeds is not None: - raise ValueError( - "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" - ) +# if pixel_values is not None and inputs_embeds is not None: +# raise ValueError( +# "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" +# ) - if inputs_embeds is None: - inputs_embeds = self.get_input_embeddings()(input_ids) +# if inputs_embeds is None: +# inputs_embeds = self.get_input_embeddings()(input_ids) - image_features = None - if pixel_values is not None: - image_features = self.get_image_features( - pixel_values=pixel_values, - vision_feature_layer=vision_feature_layer, - vision_feature_select_strategy=vision_feature_select_strategy, - ) +# image_features = None +# if pixel_values is not None: +# image_features = self.get_image_features( +# pixel_values=pixel_values, +# vision_feature_layer=vision_feature_layer, +# vision_feature_select_strategy=vision_feature_select_strategy, +# ) - inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features( - image_features, inputs_embeds, input_ids, attention_mask, labels - ) - cache_position = torch.arange(attention_mask.shape[1], device=attention_mask.device) +# inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features( +# image_features, inputs_embeds, input_ids, attention_mask, labels +# ) +# cache_position = torch.arange(attention_mask.shape[1], device=attention_mask.device) - outputs = self.language_model( - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, - ) +# outputs = self.language_model( +# attention_mask=attention_mask, +# position_ids=position_ids, +# past_key_values=past_key_values, +# inputs_embeds=inputs_embeds, +# use_cache=use_cache, +# output_attentions=output_attentions, +# output_hidden_states=output_hidden_states, +# return_dict=return_dict, +# cache_position=cache_position, +# num_logits_to_keep=num_logits_to_keep, +# ) - logits = outputs[0] +# logits = outputs[0] - loss = None +# loss = None - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output +# if not return_dict: +# output = (logits,) + outputs[1:] +# return (loss,) + output if loss is not None else output - return LlavaCausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - image_hidden_states=image_features if pixel_values is not None else None, - ) +# return LlavaCausalLMOutputWithPast( +# loss=loss, +# logits=logits, +# past_key_values=outputs.past_key_values, +# hidden_states=outputs.hidden_states, +# attentions=outputs.attentions, +# image_hidden_states=image_features if pixel_values is not None else None, +# ) def adapt_model(model, audio_block_name): modules_dict= { k: m for k, m in model.named_modules()} @@ -320,8 +320,8 @@ class Inference(object): device = "cuda" import transformers - transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.forward = patched_llava_forward # force legacy behaviour to be able to use tansformers v>(4.47) - transformers.models.llava.modeling_llava.LlavaForConditionalGeneration._merge_input_ids_with_image_features = _merge_input_ids_with_image_features + # transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.forward = patched_llava_forward # force legacy behaviour to be able to use tansformers v>(4.47) + # transformers.models.llava.modeling_llava.LlavaForConditionalGeneration._merge_input_ids_with_image_features = _merge_input_ids_with_image_features torch.set_grad_enabled(False) text_len = 512 @@ -778,7 +778,7 @@ class HunyuanVideoSampler(Inference): raise ValueError( f"Seed must be an integer, a list of integers, or None, got {seed}." ) - from wan.utils.utils import seed_everything + from shared.utils.utils import seed_everything seed_everything(seed) generator = [torch.Generator("cuda").manual_seed(seed) for seed in seeds] # generator = [torch.Generator(self.device).manual_seed(seed) for seed in seeds] @@ -956,7 +956,7 @@ class HunyuanVideoSampler(Inference): # out_latents= ref_latents / self.vae.config.scaling_factor # image = self.vae.decode(out_latents, return_dict=False, generator=generator)[0] # image = image.clamp(-1, 1) - # from wan.utils.utils import cache_video + # from shared.utils.utils import cache_video # cache_video( tensor=image, save_file="decode.mp4", fps=25, nrow=1, normalize=True, value_range=(-1, 1)) motion_pose = np.array([25] * 4) @@ -1040,5 +1040,4 @@ class HunyuanVideoSampler(Inference): return samples -def query_model_def(model_type, model_def): - return None \ No newline at end of file + diff --git a/models/hyvideo/hunyuan_handler.py b/models/hyvideo/hunyuan_handler.py new file mode 100644 index 0000000..30814b3 --- /dev/null +++ b/models/hyvideo/hunyuan_handler.py @@ -0,0 +1,147 @@ +import torch + +def get_hunyuan_text_encoder_filename(text_encoder_quantization): + if text_encoder_quantization =="int8": + text_encoder_filename = "ckpts/llava-llama-3-8b/llava-llama-3-8b-v1_1_vlm_quanto_int8.safetensors" + else: + text_encoder_filename = "ckpts/llava-llama-3-8b/llava-llama-3-8b-v1_1_vlm_fp16.safetensors" + + return text_encoder_filename + +class family_handler(): + @staticmethod + def query_model_def(base_model_type, model_def): + extra_model_def = {} + + if base_model_type in ["hunyuan_avatar", "hunyuan_custom_audio"]: + fps = 25 + elif base_model_type in ["hunyuan", "hunyuan_i2v", "hunyuan_custom_edit", "hunyuan_custom"]: + fps = 24 + else: + fps = 16 + extra_model_def["fps"] = fps + extra_model_def["frames_minimum"] = 5 + extra_model_def["frames_steps"] = 4 + extra_model_def["sliding_window"] = False + extra_model_def["embedded_guidance"] = base_model_type in ["hunyuan", "hunyuan_i2v"] + extra_model_def["cfg_star"] = base_model_type in [ "hunyuan_avatar", "hunyuan_custom_audio", "hunyuan_custom_edit", "hunyuan_custom"] + extra_model_def["skip_steps_cache"] = True + return extra_model_def + + @staticmethod + def query_supported_types(): + return ["hunyuan", "hunyuan_i2v", "hunyuan_custom", "hunyuan_custom_audio", "hunyuan_custom_edit", "hunyuan_avatar"] + + @staticmethod + def query_family_maps(): + models_eqv_map = { + } + + models_comp_map = { + "hunyuan_custom": ["hunyuan_custom_edit", "hunyuan_custom_audio"], + } + + return models_eqv_map, models_comp_map + + @staticmethod + def query_model_family(): + return "hunyuan" + + @staticmethod + def query_family_infos(): + return {"hunyuan":(20, "Hunyuan Video")} + + @staticmethod + def get_rgb_factors(model_type): + from shared.RGB_factors import get_rgb_factors + latent_rgb_factors, latent_rgb_factors_bias = get_rgb_factors("hunyuan") + return latent_rgb_factors, latent_rgb_factors_bias + + @staticmethod + def query_model_files(computeList, base_model_type, model_filename, text_encoder_quantization): + text_encoder_filename = get_hunyuan_text_encoder_filename(text_encoder_quantization) + return { + "repoId" : "DeepBeepMeep/HunyuanVideo", + "sourceFolderList" : [ "llava-llama-3-8b", "clip_vit_large_patch14", "whisper-tiny" , "det_align", "" ], + "fileList" :[ ["config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json", "preprocessor_config.json"] + computeList(text_encoder_filename) , + ["config.json", "merges.txt", "model.safetensors", "preprocessor_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json", "vocab.json"], + ["config.json", "model.safetensors", "preprocessor_config.json", "special_tokens_map.json", "tokenizer_config.json"], + ["detface.pt"], + [ "hunyuan_video_720_quanto_int8_map.json", "hunyuan_video_custom_VAE_fp32.safetensors", "hunyuan_video_custom_VAE_config.json", "hunyuan_video_VAE_fp32.safetensors", "hunyuan_video_VAE_config.json" , "hunyuan_video_720_quanto_int8_map.json" ] + computeList(model_filename) + ] + } + + @staticmethod + def load_model(model_filename, model_type = None, base_model_type = None, model_def = None, quantizeTransformer = False, text_encoder_quantization = None, dtype = torch.bfloat16, VAE_dtype = torch.float32, mixed_precision_transformer = False, save_quantized = False): + from .hunyuan import HunyuanVideoSampler + from mmgp import offload + + hunyuan_model = HunyuanVideoSampler.from_pretrained( + model_filepath = model_filename, + model_type = model_type, + base_model_type = base_model_type, + text_encoder_filepath = get_hunyuan_text_encoder_filename(text_encoder_quantization), + dtype = dtype, + quantizeTransformer = quantizeTransformer, + VAE_dtype = VAE_dtype, + mixed_precision_transformer = mixed_precision_transformer, + save_quantized = save_quantized + ) + + pipe = { "transformer" : hunyuan_model.model, "text_encoder" : hunyuan_model.text_encoder, "text_encoder_2" : hunyuan_model.text_encoder_2, "vae" : hunyuan_model.vae } + + if hunyuan_model.wav2vec != None: + pipe["wav2vec"] = hunyuan_model.wav2vec + + + # if hunyuan_model.align_instance != None: + # pipe["align_instance"] = hunyuan_model.align_instance.facedet.model + + + from .modules.models import get_linear_split_map + + split_linear_modules_map = get_linear_split_map() + hunyuan_model.model.split_linear_modules_map = split_linear_modules_map + offload.split_linear_modules(hunyuan_model.model, split_linear_modules_map ) + + + return hunyuan_model, pipe + + @staticmethod + def update_default_settings(base_model_type, model_def, ui_defaults): + ui_defaults["embedded_guidance_scale"]= 6.0 + + if base_model_type in ["hunyuan","hunyuan_i2v"]: + ui_defaults.update({ + "guidance_scale": 7.0, + }) + + elif base_model_type in ["hunyuan_custom"]: + ui_defaults.update({ + "guidance_scale": 7.5, + "flow_shift": 13, + "resolution": "1280x720", + "video_prompt_type": "I", + }) + elif base_model_type in ["hunyuan_custom_audio"]: + ui_defaults.update({ + "guidance_scale": 7.5, + "flow_shift": 13, + "video_prompt_type": "I", + }) + elif base_model_type in ["hunyuan_custom_edit"]: + ui_defaults.update({ + "guidance_scale": 7.5, + "flow_shift": 13, + "video_prompt_type": "MVAI", + "sliding_window_size": 129, + }) + elif base_model_type in ["hunyuan_avatar"]: + ui_defaults.update({ + "guidance_scale": 7.5, + "flow_shift": 5, + "remove_background_images_ref": 0, + "skip_steps_start_step_perc": 25, + "video_length": 129, + "video_prompt_type": "I", + }) diff --git a/hyvideo/modules/__init__.py b/models/hyvideo/modules/__init__.py similarity index 100% rename from hyvideo/modules/__init__.py rename to models/hyvideo/modules/__init__.py diff --git a/hyvideo/modules/activation_layers.py b/models/hyvideo/modules/activation_layers.py similarity index 100% rename from hyvideo/modules/activation_layers.py rename to models/hyvideo/modules/activation_layers.py diff --git a/hyvideo/modules/attenion.py b/models/hyvideo/modules/attenion.py similarity index 100% rename from hyvideo/modules/attenion.py rename to models/hyvideo/modules/attenion.py diff --git a/hyvideo/modules/audio_adapters.py b/models/hyvideo/modules/audio_adapters.py similarity index 100% rename from hyvideo/modules/audio_adapters.py rename to models/hyvideo/modules/audio_adapters.py diff --git a/hyvideo/modules/embed_layers.py b/models/hyvideo/modules/embed_layers.py similarity index 100% rename from hyvideo/modules/embed_layers.py rename to models/hyvideo/modules/embed_layers.py diff --git a/hyvideo/modules/mlp_layers.py b/models/hyvideo/modules/mlp_layers.py similarity index 100% rename from hyvideo/modules/mlp_layers.py rename to models/hyvideo/modules/mlp_layers.py diff --git a/hyvideo/modules/models.py b/models/hyvideo/modules/models.py similarity index 99% rename from hyvideo/modules/models.py rename to models/hyvideo/modules/models.py index 48978a9..92c1c08 100644 --- a/hyvideo/modules/models.py +++ b/models/hyvideo/modules/models.py @@ -18,7 +18,7 @@ from .modulate_layers import ModulateDiT, modulate, modulate_ , apply_gate, appl from .token_refiner import SingleTokenRefiner import numpy as np from mmgp import offload -from wan.modules.attention import pay_attention +from shared.attention import pay_attention from .audio_adapters import AudioProjNet2, PerceiverAttentionCA def get_linear_split_map(): diff --git a/hyvideo/modules/modulate_layers.py b/models/hyvideo/modules/modulate_layers.py similarity index 100% rename from hyvideo/modules/modulate_layers.py rename to models/hyvideo/modules/modulate_layers.py diff --git a/hyvideo/modules/norm_layers.py b/models/hyvideo/modules/norm_layers.py similarity index 100% rename from hyvideo/modules/norm_layers.py rename to models/hyvideo/modules/norm_layers.py diff --git a/hyvideo/modules/original models.py b/models/hyvideo/modules/original models.py similarity index 100% rename from hyvideo/modules/original models.py rename to models/hyvideo/modules/original models.py diff --git a/hyvideo/modules/placement.py b/models/hyvideo/modules/placement.py similarity index 100% rename from hyvideo/modules/placement.py rename to models/hyvideo/modules/placement.py diff --git a/hyvideo/modules/posemb_layers.py b/models/hyvideo/modules/posemb_layers.py similarity index 100% rename from hyvideo/modules/posemb_layers.py rename to models/hyvideo/modules/posemb_layers.py diff --git a/hyvideo/modules/token_refiner.py b/models/hyvideo/modules/token_refiner.py similarity index 100% rename from hyvideo/modules/token_refiner.py rename to models/hyvideo/modules/token_refiner.py diff --git a/hyvideo/modules/utils.py b/models/hyvideo/modules/utils.py similarity index 100% rename from hyvideo/modules/utils.py rename to models/hyvideo/modules/utils.py diff --git a/hyvideo/prompt_rewrite.py b/models/hyvideo/prompt_rewrite.py similarity index 100% rename from hyvideo/prompt_rewrite.py rename to models/hyvideo/prompt_rewrite.py diff --git a/hyvideo/text_encoder/__init__.py b/models/hyvideo/text_encoder/__init__.py similarity index 97% rename from hyvideo/text_encoder/__init__.py rename to models/hyvideo/text_encoder/__init__.py index 1376718..9bd47d4 100644 --- a/hyvideo/text_encoder/__init__.py +++ b/models/hyvideo/text_encoder/__init__.py @@ -15,6 +15,7 @@ from transformers.utils import ModelOutput from ..constants import TEXT_ENCODER_PATH, TOKENIZER_PATH from ..constants import PRECISION_TO_TYPE +from .llava.modeling_llava import LlavaForConditionalGeneration def use_default(value, default): @@ -188,10 +189,16 @@ class TextEncoder(nn.Module): if "llm" in text_encoder_type: from mmgp import offload - forcedConfigPath= None if "i2v" in text_encoder_type else "ckpts/llava-llama-3-8b/config.json" - self.model= offload.fast_load_transformers_model(self.model_path, modelPrefix="language_model" if forcedConfigPath != None else None, forcedConfigPath=forcedConfigPath) - if forcedConfigPath != None: + # forcedConfigPath= None if "i2v" in text_encoder_type else "ckpts/llava-llama-3-8b/config.json" + # self.model= offload.fast_load_transformers_model(self.model_path, modelPrefix="language_model" if forcedConfigPath != None else None, forcedConfigPath=forcedConfigPath) + + if "i2v" in text_encoder_type: + self.model= offload.fast_load_transformers_model(self.model_path, modelClass= LlavaForConditionalGeneration) + else: + self.model= offload.fast_load_transformers_model(self.model_path, modelPrefix="language_model", forcedConfigPath = "ckpts/llava-llama-3-8b/config.json") self.model.final_layer_norm = self.model.model.norm + + else: self.model, self.model_path = load_text_encoder( diff --git a/models/hyvideo/text_encoder/llava/__init__.py b/models/hyvideo/text_encoder/llava/__init__.py new file mode 100644 index 0000000..e6d2f52 --- /dev/null +++ b/models/hyvideo/text_encoder/llava/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# from typing import TYPE_CHECKING + +# from ...utils import _LazyModule +# from ...utils.import_utils import define_import_structure + + +# if TYPE_CHECKING: +# from .configuration_llava import * +# from .image_processing_llava_fast import * +# from .modeling_llava import * +# from .processing_llava import * +# else: +# import sys + +# _file = globals()["__file__"] + # sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/models/hyvideo/text_encoder/llava/configuration_llava.py b/models/hyvideo/text_encoder/llava/configuration_llava.py new file mode 100644 index 0000000..9c30798 --- /dev/null +++ b/models/hyvideo/text_encoder/llava/configuration_llava.py @@ -0,0 +1,137 @@ +# coding=utf-8 +# Copyright 2023 Microsoft Research & University of Wisconsin-Madison and the HuggingFace Inc. team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Llava model configuration""" + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging +from transformers.models.auto import CONFIG_MAPPING, AutoConfig + + +logger = logging.get_logger(__name__) + + +class LlavaConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`LlavaForConditionalGeneration`]. It is used to instantiate an + Llava model according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of the Llava-9B. + + e.g. [llava-hf/llava-9b](https://huggingface.co/llava-hf/llava-9b) + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + vision_config (`Union[AutoConfig, dict]`, *optional*, defaults to `CLIPVisionConfig`): + The config object or dictionary of the vision backbone. + text_config (`Union[AutoConfig, dict]`, *optional*, defaults to `LlamaConfig`): + The config object or dictionary of the text backbone. + image_token_index (`int`, *optional*, defaults to 32000): + The image token index to encode the image prompt. + projector_hidden_act (`str`, *optional*, defaults to `"gelu"`): + The activation function used by the multimodal projector. + vision_feature_select_strategy (`str`, *optional*, defaults to `"default"`): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"`. + vision_feature_layer (`Union[int, List[int]]`, *optional*, defaults to -2): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + image_seq_length (`int`, *optional*, defaults to 576): + Sequence length of one image embedding. + multimodal_projector_bias (`bool`, *optional*, defaults to `True`): + Whether to use bias in the multimodal projector. + + Example: + + ```python + >>> from transformers import LlavaForConditionalGeneration, LlavaConfig, CLIPVisionConfig, LlamaConfig + + >>> # Initializing a CLIP-vision config + >>> vision_config = CLIPVisionConfig() + + >>> # Initializing a Llama config + >>> text_config = LlamaConfig() + + >>> # Initializing a Llava llava-1.5-7b style configuration + >>> configuration = LlavaConfig(vision_config, text_config) + + >>> # Initializing a model from the llava-1.5-7b style configuration + >>> model = LlavaForConditionalGeneration(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "llava" + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig} + is_composition = True + + def __init__( + self, + vision_config=None, + text_config=None, + image_token_index=32000, + projector_hidden_act="gelu", + vision_feature_select_strategy="default", + vision_feature_layer=-2, + image_seq_length=576, + multimodal_projector_bias=True, + **kwargs, + ): + self.image_token_index = image_token_index + self.projector_hidden_act = projector_hidden_act + self.image_seq_length = image_seq_length + + if vision_feature_select_strategy not in ["default", "full"]: + raise ValueError( + "vision_feature_select_strategy should be one of 'default', 'full'." + f"Got: {vision_feature_select_strategy}" + ) + + self.vision_feature_select_strategy = vision_feature_select_strategy + self.vision_feature_layer = vision_feature_layer + + if isinstance(vision_config, dict): + vision_config["model_type"] = ( + vision_config["model_type"] if "model_type" in vision_config else "clip_vision_model" + ) + vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config) + elif vision_config is None: + vision_config = CONFIG_MAPPING["clip_vision_model"]( + intermediate_size=4096, + hidden_size=1024, + patch_size=14, + image_size=336, + num_hidden_layers=24, + num_attention_heads=16, + vocab_size=32000, + projection_dim=768, + ) + + self.vision_config = vision_config + + if isinstance(text_config, dict): + text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "llama" + text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config) + elif text_config is None: + text_config = CONFIG_MAPPING["llama"]() + + self.text_config = text_config + self.multimodal_projector_bias = multimodal_projector_bias + + super().__init__(**kwargs) + + +__all__ = ["LlavaConfig"] diff --git a/models/hyvideo/text_encoder/llava/image_processing_llava.py b/models/hyvideo/text_encoder/llava/image_processing_llava.py new file mode 100644 index 0000000..37ef079 --- /dev/null +++ b/models/hyvideo/text_encoder/llava/image_processing_llava.py @@ -0,0 +1,436 @@ +# coding=utf-8 +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Image processor class for LLaVa.""" + +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np + +from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict +from ...image_transforms import ( + convert_to_rgb, + get_resize_output_image_size, + resize, + to_channel_dimension_format, +) +from ...image_utils import ( + OPENAI_CLIP_MEAN, + OPENAI_CLIP_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + get_image_size, + infer_channel_dimension_format, + is_scaled_image, + make_list_of_images, + to_numpy_array, + valid_images, + validate_kwargs, + validate_preprocess_arguments, +) +from ...utils import TensorType, is_vision_available, logging + + +logger = logging.get_logger(__name__) + + +if is_vision_available(): + import PIL + + +class LlavaImageProcessor(BaseImageProcessor): + r""" + Constructs a LLaVa image processor. + + Args: + do_pad (`bool`, *optional*, defaults to `False`): + Whether to pad the image to a square based on the longest edge. + The padding value is determined by the `image_mean` parameter. + Can be overridden by `do_pad` in the `preprocess` method. + do_resize (`bool`, *optional*, defaults to `True`): + Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by + `do_resize` in the `preprocess` method. + size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 224}`): + Size of the image after resizing. The shortest edge of the image is resized to size["shortest_edge"], with + the longest edge resized to keep the input aspect ratio. Can be overridden by `size` in the `preprocess` + method. + resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`): + Resampling filter to use if resizing the image. Can be overridden by `resample` in the `preprocess` method. + do_center_crop (`bool`, *optional*, defaults to `True`): + Whether to center crop the image to the specified `crop_size`. Can be overridden by `do_center_crop` in the + `preprocess` method. + crop_size (`Dict[str, int]` *optional*, defaults to 224): + Size of the output image after applying `center_crop`. Can be overridden by `crop_size` in the `preprocess` + method. + do_rescale (`bool`, *optional*, defaults to `True`): + Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by `do_rescale` in + the `preprocess` method. + rescale_factor (`int` or `float`, *optional*, defaults to `1/255`): + Scale factor to use if rescaling the image. Can be overridden by `rescale_factor` in the `preprocess` + method. + do_normalize (`bool`, *optional*, defaults to `True`): + Whether to normalize the image. Can be overridden by `do_normalize` in the `preprocess` method. + image_mean (`float` or `List[float]`, *optional*, defaults to `[0.48145466, 0.4578275, 0.40821073]`): + Mean to use if normalizing the image. This is a float or list of floats the length of the number of + channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method. + image_std (`float` or `List[float]`, *optional*, defaults to `[0.26862954, 0.26130258, 0.27577711]`): + Standard deviation to use if normalizing the image. This is a float or list of floats the length of the + number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method. + Can be overridden by the `image_std` parameter in the `preprocess` method. + do_convert_rgb (`bool`, *optional*, defaults to `True`): + Whether to convert the image to RGB. + """ + + model_input_names = ["pixel_values"] + + def __init__( + self, + do_pad: bool = False, + do_resize: bool = True, + size: Dict[str, int] = None, + resample: PILImageResampling = PILImageResampling.BICUBIC, + do_center_crop: bool = True, + crop_size: Dict[str, int] = None, + do_rescale: bool = True, + rescale_factor: Union[int, float] = 1 / 255, + do_normalize: bool = True, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + do_convert_rgb: bool = True, + **kwargs, + ) -> None: + super().__init__(**kwargs) + size = size if size is not None else {"shortest_edge": 224} + size = get_size_dict(size, default_to_square=False) + crop_size = crop_size if crop_size is not None else {"height": 224, "width": 224} + crop_size = get_size_dict(crop_size, default_to_square=True, param_name="crop_size") + + self.do_pad = do_pad + self.do_resize = do_resize + self.size = size + self.resample = resample + self.do_center_crop = do_center_crop + self.crop_size = crop_size + self.do_rescale = do_rescale + self.rescale_factor = rescale_factor + self.do_normalize = do_normalize + self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN + self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD + self.do_convert_rgb = do_convert_rgb + self._valid_processor_keys = [ + "images", + "do_pad", + "do_resize", + "size", + "resample", + "do_center_crop", + "crop_size", + "do_rescale", + "rescale_factor", + "do_normalize", + "image_mean", + "image_std", + "do_convert_rgb", + "return_tensors", + "data_format", + "input_data_format", + ] + + def pad_to_square( + self, + image: np.ndarray, + background_color: Union[int, Tuple[int, int, int]] = 0, + data_format: Optional[Union[str, ChannelDimension]] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + ) -> np.array: + """ + Pads an image to a square based on the longest edge. + + Args: + image (`np.ndarray`): + The image to pad. + background_color (`int` or `Tuple[int, int, int]`, *optional*, defaults to 0): + The color to use for the padding. Can be an integer for single channel or a + tuple of integers representing for multi-channel images. If passed as integer + in mutli-channel mode, it will default to `0` in subsequent channels. + data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the output image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + If unset, will use same as the input image. + input_data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + If unset, will use the inferred format of the input image. + + Returns: + `np.ndarray`: The padded image. + """ + height, width = get_image_size(image, input_data_format) + num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1] + + if height == width: + image = ( + to_channel_dimension_format(image, data_format, input_data_format) + if data_format is not None + else image + ) + return image + + max_dim = max(height, width) + + # Ensure background_color is the correct shape + if isinstance(background_color, int): + background_color = [background_color] + elif len(background_color) != num_channels: + raise ValueError( + f"background_color must have no more than {num_channels} elements to match the number of channels" + ) + + if input_data_format == ChannelDimension.FIRST: + result = np.zeros((num_channels, max_dim, max_dim), dtype=image.dtype) + for i, color in enumerate(background_color): + result[i, :, :] = color + if width > height: + start = (max_dim - height) // 2 + result[:, start : start + height, :] = image + else: + start = (max_dim - width) // 2 + result[:, :, start : start + width] = image + else: + result = np.zeros((max_dim, max_dim, num_channels), dtype=image.dtype) + for i, color in enumerate(background_color): + result[:, :, i] = color + if width > height: + start = (max_dim - height) // 2 + result[start : start + height, :, :] = image + else: + start = (max_dim - width) // 2 + result[:, start : start + width, :] = image + + image = ( + to_channel_dimension_format(result, data_format, input_data_format) if data_format is not None else result + ) + return image + + # Copied from transformers.models.clip.image_processing_clip.CLIPImageProcessor.resize + def resize( + self, + image: np.ndarray, + size: Dict[str, int], + resample: PILImageResampling = PILImageResampling.BICUBIC, + data_format: Optional[Union[str, ChannelDimension]] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + **kwargs, + ) -> np.ndarray: + """ + Resize an image. The shortest edge of the image is resized to size["shortest_edge"], with the longest edge + resized to keep the input aspect ratio. + + Args: + image (`np.ndarray`): + Image to resize. + size (`Dict[str, int]`): + Size of the output image. + resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`): + Resampling filter to use when resiizing the image. + data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format of the image. If not provided, it will be the same as the input image. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format of the input image. If not provided, it will be inferred. + """ + default_to_square = True + if "shortest_edge" in size: + size = size["shortest_edge"] + default_to_square = False + elif "height" in size and "width" in size: + size = (size["height"], size["width"]) + else: + raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.") + + output_size = get_resize_output_image_size( + image, + size=size, + default_to_square=default_to_square, + input_data_format=input_data_format, + ) + return resize( + image, + size=output_size, + resample=resample, + data_format=data_format, + input_data_format=input_data_format, + **kwargs, + ) + + def preprocess( + self, + images: ImageInput, + do_pad: Optional[bool] = None, + do_resize: Optional[bool] = None, + size: Optional[Dict[str, int]] = None, + resample: Optional[PILImageResampling] = None, + do_center_crop: Optional[bool] = None, + crop_size: Optional[int] = None, + do_rescale: Optional[bool] = None, + rescale_factor: Optional[float] = None, + do_normalize: Optional[bool] = None, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + do_convert_rgb: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + **kwargs, + ) -> PIL.Image.Image: + """ + Preprocess an image or batch of images. + + Args: + images (`ImageInput`): + Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If + passing in images with pixel values between 0 and 1, set `do_rescale=False`. + do_pad (`bool`, *optional*, defaults to `self.do_pad`): + Whether to pad the image to a square based on the longest edge. + The padding value is determined by the `image_mean` parameter. + do_resize (`bool`, *optional*, defaults to `self.do_resize`): + Whether to resize the image. + size (`Dict[str, int]`, *optional*, defaults to `self.size`): + Size of the image after resizing. Shortest edge of the image is resized to size["shortest_edge"], with + the longest edge resized to keep the input aspect ratio. + resample (`int`, *optional*, defaults to `self.resample`): + Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only + has an effect if `do_resize` is set to `True`. + do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`): + Whether to center crop the image. + crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`): + Size of the center crop. Only has an effect if `do_center_crop` is set to `True`. + do_rescale (`bool`, *optional*, defaults to `self.do_rescale`): + Whether to rescale the image. + rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`): + Rescale factor to rescale the image by if `do_rescale` is set to `True`. + do_normalize (`bool`, *optional*, defaults to `self.do_normalize`): + Whether to normalize the image. + image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`): + Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`. + image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`): + Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to + `True`. + do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`): + Whether to convert the image to RGB. + return_tensors (`str` or `TensorType`, *optional*): + The type of tensors to return. Can be one of: + - Unset: Return a list of `np.ndarray`. + - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`. + - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`. + - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`. + - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`. + data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`): + The channel dimension format for the output image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + - Unset: Use the channel dimension format of the input image. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the input image. If unset, the channel dimension format is inferred + from the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. + """ + do_pad = do_pad if do_pad is not None else self.do_pad + do_resize = do_resize if do_resize is not None else self.do_resize + size = size if size is not None else self.size + size = get_size_dict(size, param_name="size", default_to_square=False) + resample = resample if resample is not None else self.resample + do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop + crop_size = crop_size if crop_size is not None else self.crop_size + crop_size = get_size_dict(crop_size, param_name="crop_size", default_to_square=True) + do_rescale = do_rescale if do_rescale is not None else self.do_rescale + rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor + do_normalize = do_normalize if do_normalize is not None else self.do_normalize + image_mean = image_mean if image_mean is not None else self.image_mean + image_std = image_std if image_std is not None else self.image_std + do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb + + validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys) + + images = make_list_of_images(images) + + if not valid_images(images): + raise ValueError( + "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, " + "torch.Tensor, tf.Tensor or jax.ndarray." + ) + # we don't pass `do_pad` here since LLaVa uses a custom padding to a square + validate_preprocess_arguments( + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + do_center_crop=do_center_crop, + crop_size=crop_size, + do_resize=do_resize, + size=size, + resample=resample, + ) + + if do_convert_rgb: + images = [convert_to_rgb(image) for image in images] + + # All transformations expect numpy arrays. + images = [to_numpy_array(image) for image in images] + + if is_scaled_image(images[0]) and do_rescale: + logger.warning_once( + "It looks like you are trying to rescale already rescaled images. If the input" + " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again." + ) + + if input_data_format is None: + # We assume that all images have the same channel dimension format. + input_data_format = infer_channel_dimension_format(images[0]) + + processed_images = [] + for image in images: + if do_pad: + image = self.pad_to_square( + image=image, + background_color=tuple(int(x * 255) for x in self.image_mean), + input_data_format=input_data_format, + ) + + if do_resize: + image = self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format) + + if do_center_crop: + image = self.center_crop(image=image, size=crop_size, input_data_format=input_data_format) + + if do_rescale: + image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format) + + if do_normalize: + image = self.normalize( + image=image, mean=image_mean, std=image_std, input_data_format=input_data_format + ) + + image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) + processed_images.append(image) + + return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + +__all__ = ["LlavaImageProcessor"] diff --git a/models/hyvideo/text_encoder/llava/image_processing_llava_fast.py b/models/hyvideo/text_encoder/llava/image_processing_llava_fast.py new file mode 100644 index 0000000..d85eb89 --- /dev/null +++ b/models/hyvideo/text_encoder/llava/image_processing_llava_fast.py @@ -0,0 +1,201 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Fast Image processor class for LLaVa.""" + +from typing import List, Optional, Tuple, Union + +from ...image_processing_utils import ( + BatchFeature, +) +from ...image_processing_utils_fast import ( + BASE_IMAGE_PROCESSOR_FAST_DOCSTRING, + BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS, + BaseImageProcessorFast, + DefaultFastImageProcessorKwargs, + group_images_by_shape, + reorder_images, +) +from ...image_utils import ( + OPENAI_CLIP_MEAN, + OPENAI_CLIP_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, +) +from ...processing_utils import Unpack +from ...utils import ( + TensorType, + add_start_docstrings, + is_torch_available, + is_torchvision_available, + is_torchvision_v2_available, + is_vision_available, +) + + +if is_vision_available(): + from ...image_utils import PILImageResampling + +if is_torch_available(): + import torch + +if is_torchvision_available(): + if is_torchvision_v2_available(): + from torchvision.transforms.v2 import functional as F + else: + from torchvision.transforms import functional as F + + +class LlavaFastImageProcessorKwargs(DefaultFastImageProcessorKwargs): + do_pad: Optional[bool] + + +@add_start_docstrings( + "Constructs a fast Llava image processor.", + BASE_IMAGE_PROCESSOR_FAST_DOCSTRING, + """ + do_pad (`bool`, *optional*, defaults to `self.do_pad`): + Whether to pad the image to a square based on the longest edge. Can be overridden by the `do_pad` parameter + """, +) +class LlavaImageProcessorFast(BaseImageProcessorFast): + resample = PILImageResampling.BICUBIC + image_mean = OPENAI_CLIP_MEAN + image_std = OPENAI_CLIP_STD + size = {"shortest_edge": 224} + default_to_square = False + crop_size = {"height": 224, "width": 224} + do_pad = False + do_resize = True + do_center_crop = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + valid_kwargs = LlavaFastImageProcessorKwargs + + def __init__(self, **kwargs: Unpack[LlavaFastImageProcessorKwargs]) -> None: + super().__init__(**kwargs) + + @add_start_docstrings( + BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS, + """ + do_pad (`bool`, *optional*, defaults to `self.do_pad`): + Whether to pad the image to a square based on the longest edge. Can be overridden by the `do_pad` parameter + """, + ) + def preprocess(self, images: ImageInput, **kwargs: Unpack[LlavaFastImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def pad_to_square( + self, + images: "torch.Tensor", + background_color: Union[int, Tuple[int, int, int]] = 0, + ) -> "torch.Tensor": + """ + Pads an image to a square based on the longest edge. + + Args: + images (`np.ndarray`): + The images to pad. + background_color (`int` or `Tuple[int, int, int]`, *optional*, defaults to 0): + The color to use for the padding. Can be an integer for single channel or a + tuple of integers representing for multi-channel images. If passed as integer + in mutli-channel mode, it will default to `0` in subsequent channels. + Returns: + `torch.Tensor`: The padded images. + """ + height, width = get_image_size(images, ChannelDimension.FIRST) + + if height == width: + return images + + num_channels = images.shape[1] if len(images.shape) == 4 else images.shape[0] + if isinstance(background_color, int): + background_color = [background_color] + [0] * (num_channels - 1) + elif len(background_color) != num_channels: + raise ValueError( + f"background_color must have no more than {num_channels} elements to match the number of channels" + ) + + max_dim = max(height, width) + paste_x_left = (max_dim - width) // 2 + paste_y_left = (max_dim - height) // 2 + paste_x_right = max_dim - width - paste_x_left + paste_y_right = max_dim - height - paste_y_left + padded_images = F.pad( + images, padding=[paste_x_left, paste_y_left, paste_x_right, paste_y_right], fill=background_color + ) + + return padded_images + + def _preprocess( + self, + images: List["torch.Tensor"], + do_resize: bool, + size: SizeDict, + interpolation: Optional["F.InterpolationMode"], + do_pad: bool, + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: Optional[Union[float, List[float]]], + image_std: Optional[Union[float, List[float]]], + return_tensors: Optional[Union[str, TensorType]], + ) -> BatchFeature: + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_pad: + stacked_images = self.pad_to_square( + images=stacked_images, background_color=tuple(int(x * 255) for x in self.image_mean) + ) + resized_images_grouped[shape] = stacked_images + padded_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for batched resizing + # Needed in case do_pad is False, or padding returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(padded_images) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(image=stacked_images, size=size, interpolation=interpolation) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + # Needed in case do_resize is False, or resize returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(resized_images) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images + + return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + +__all__ = ["LlavaImageProcessorFast"] diff --git a/models/hyvideo/text_encoder/llava/modeling_llava.py b/models/hyvideo/text_encoder/llava/modeling_llava.py new file mode 100644 index 0000000..f4ae058 --- /dev/null +++ b/models/hyvideo/text_encoder/llava/modeling_llava.py @@ -0,0 +1,531 @@ +# coding=utf-8 +# Copyright 2023 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Llava model.""" + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn + +from transformers.activations import ACT2FN +from transformers.generation import GenerationMixin +from transformers.modeling_outputs import ModelOutput +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_torchdynamo_compiling, + logging, + replace_return_docstrings, +) +from transformers.utils.deprecation import deprecate_kwarg +from transformers.models.auto import AutoModel, AutoModelForCausalLM +from .configuration_llava import LlavaConfig + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "LlavaConfig" + +# Base docstring +_CHECKPOINT_FOR_DOC = "llava-hf/llava-1.5-7b-hf" + + +@dataclass +class LlavaCausalLMOutputWithPast(ModelOutput): + """ + Base class for Llava causal language model (or autoregressive) outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size (batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[List[torch.FloatTensor]] = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + image_hidden_states: Optional[torch.FloatTensor] = None + + +class LlavaMultiModalProjector(nn.Module): + def __init__(self, config: LlavaConfig): + super().__init__() + # We have hidden_size * the number of vision feature layers + num_feature_layers = 1 if isinstance(config.vision_feature_layer, int) else len(config.vision_feature_layer) + self.linear_1 = nn.Linear( + config.vision_config.hidden_size * num_feature_layers, + config.text_config.hidden_size, + bias=config.multimodal_projector_bias, + ) + self.act = ACT2FN[config.projector_hidden_act] + self.linear_2 = nn.Linear( + config.text_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias + ) + + def forward(self, image_features): + hidden_states = self.linear_1(image_features) + hidden_states = self.act(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +LLAVA_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`LlavaConfig`] or [`LlavaVisionConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare LLaMA Model outputting raw hidden-states without any specific head on top.", + LLAVA_START_DOCSTRING, +) +class LlavaPreTrainedModel(PreTrainedModel): + config_class = LlavaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["LlavaVisionAttention"] + _skip_keys_device_placement = "past_key_values" + _supports_cache_class = True + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_quantized_cache = True + _supports_static_cache = True + + def _init_weights(self, module): + # important: this ported version of Llava isn't meant for training from scratch - only + # inference and fine-tuning - so the proper init weights code has been removed - the original codebase + # https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose + std = ( + self.config.initializer_range + if hasattr(self.config, "initializer_range") + else self.config.text_config.initializer_range + ) + + if hasattr(module, "class_embedding"): + module.class_embedding.data.normal_(mean=0.0, std=std) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +LLAVA_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)): + The tensors corresponding to the input images. Pixel values can be obtained using + [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details ([]`LlavaProcessor`] uses + [`CLIPImageProcessor`] for processing images). + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + vision_feature_layer (`Union[int, List[int]], *optional*, defaults to -2`): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*, defaults to `"default"`): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"`. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + """The LLAVA model which consists of a vision backbone and a language model.""", + LLAVA_START_DOCSTRING, +) +class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin): + def __init__(self, config: LlavaConfig): + super().__init__(config) + self.vision_tower = AutoModel.from_config(config.vision_config) + + self.multi_modal_projector = LlavaMultiModalProjector(config) + self.vocab_size = config.text_config.vocab_size + self.language_model = AutoModelForCausalLM.from_config(config.text_config) + + if self.language_model._tied_weights_keys is not None: + self._tied_weights_keys = [f"language_model.{k}" for k in self.language_model._tied_weights_keys] + + self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1 + + self.post_init() + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.language_model.get_output_embeddings() + + def set_output_embeddings(self, new_embeddings): + self.language_model.set_output_embeddings(new_embeddings) + + def set_decoder(self, decoder): + self.language_model.set_decoder(decoder) + + def get_decoder(self): + return self.language_model.get_decoder() + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + vision_feature_layer: Union[int, List[int]], + vision_feature_select_strategy: str, + **kwargs, + ): + """ + Obtains image last hidden states from the vision tower and apply multimodal projection. + + Args: + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + The tensors corresponding to the input images. + vision_feature_layer (`Union[int, List[int]]`): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` + Returns: + image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + """ + if vision_feature_select_strategy not in ["default", "full"]: + raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}") + + kwargs = {k: v for k, v in kwargs.items() if v is not None} + # this is not memory efficient at all (output_hidden_states=True) will save all the hidden states. + image_outputs = self.vision_tower(pixel_values, output_hidden_states=True, **kwargs) + + # If we have one vision feature layer, return the corresponding hidden states, + # otherwise, select the hidden states of each feature layer and concatenate them + if isinstance(vision_feature_layer, int): + selected_image_feature = image_outputs.hidden_states[vision_feature_layer] + if vision_feature_select_strategy == "default": + selected_image_feature = selected_image_feature[:, 1:] + else: + hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + # For default; crop CLS from each hidden state in the hidden state pool + if vision_feature_select_strategy == "default": + hs_pool = [hs[:, 1:] for hs in hs_pool] + selected_image_feature = torch.cat(hs_pool, dim=-1) + + image_features = self.multi_modal_projector(selected_image_feature) + return image_features + + + def _merge_input_ids_with_image_features(self, image_features, inputs_embeds, input_ids, attention_mask, labels): + num_images, num_image_patches, embed_dim = image_features.shape + batch_size, sequence_length = input_ids.shape + left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(self.pad_token_id)) + # 1. Create a mask to know where special image tokens are + special_image_token_mask = input_ids == self.config.image_token_index + num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1) + # Compute the maximum embed dimension + max_embed_dim = (num_special_image_tokens.max() * (num_image_patches - 1)) + sequence_length + batch_indices, non_image_indices = torch.where(input_ids != self.config.image_token_index) + + # 2. Compute the positions where text should be written + # Calculate new positions for text tokens in merged image-text sequence. + # `special_image_token_mask` identifies image tokens. Each image token will be replaced by `nb_text_tokens_per_images - 1` text tokens. + # `torch.cumsum` computes how each image token shifts subsequent text token positions. + # - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one. + new_token_positions = torch.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1) - 1 + nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] + if left_padding: + new_token_positions += nb_image_pad[:, None] # offset for left padding + text_to_overwrite = new_token_positions[batch_indices, non_image_indices] + + # 3. Create the full embedding, already padded to the maximum position + final_embedding = torch.zeros( + batch_size, max_embed_dim, embed_dim, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + final_attention_mask = torch.zeros( + batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device + ) + if labels is not None: + final_labels = torch.full( + (batch_size, max_embed_dim), self.config.ignore_index, dtype=input_ids.dtype, device=input_ids.device + ) + # In case the Vision model or the Language model has been offloaded to CPU, we need to manually + # set the corresponding tensors into their correct target device. + target_device = inputs_embeds.device + batch_indices, non_image_indices, text_to_overwrite = ( + batch_indices.to(target_device), + non_image_indices.to(target_device), + text_to_overwrite.to(target_device), + ) + attention_mask = attention_mask.to(target_device) + + # 4. Fill the embeddings based on the mask. If we have ["hey" "", "how", "are"] + # we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the image features + final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] + final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices] + if labels is not None: + final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices] + + # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) + image_to_overwrite = torch.full( + (batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device + ) + image_to_overwrite[batch_indices, text_to_overwrite] = False + image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device) + + if image_to_overwrite.sum() != image_features.shape[:-1].numel(): + raise ValueError( + f"The input provided to the model are wrong. The number of image tokens is {torch.sum(special_image_token_mask)} while" + f" the number of image given to the model is {num_images}. This prevents correct indexing and breaks batch generation." + ) + + final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device) + final_attention_mask |= image_to_overwrite + position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1) + + # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. + batch_indices, pad_indices = torch.where(input_ids == self.pad_token_id) + indices_to_mask = new_token_positions[batch_indices, pad_indices] + + final_embedding[batch_indices, indices_to_mask] = 0 + + if labels is None: + final_labels = None + + return final_embedding, final_attention_mask, final_labels, position_ids + + # @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + # @add_start_docstrings_to_model_forward(LLAVA_INPUTS_DOCSTRING) + # @replace_return_docstrings(output_type=LlavaCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + vision_feature_layer: Optional[int] = None, + vision_feature_select_strategy: Optional[str] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, + ): + from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast + + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + vision_feature_layer = ( + vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer + ) + vision_feature_select_strategy = ( + vision_feature_select_strategy + if vision_feature_select_strategy is not None + else self.config.vision_feature_select_strategy + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if pixel_values is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" + ) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + image_features = None + if pixel_values is not None: + image_features = self.get_image_features( + pixel_values=pixel_values, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + ) + + + inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features( + image_features, inputs_embeds, input_ids, attention_mask, labels + ) + cache_position = torch.arange(attention_mask.shape[1], device=attention_mask.device) + + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + num_logits_to_keep=num_logits_to_keep, + ) + + logits = outputs[0] + + loss = None + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return LlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features if pixel_values is not None else None, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + attention_mask=None, + cache_position=None, + logits_to_keep=None, + **kwargs, + ): + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model + + model_inputs = self.language_model.prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + logits_to_keep=logits_to_keep, + **kwargs, + ) + + if cache_position[0] == 0: + # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore + # Otherwise we need pixel values to be passed to model + model_inputs["pixel_values"] = pixel_values + + return model_inputs + + +__all__ = ["LlavaForConditionalGeneration", "LlavaPreTrainedModel"] diff --git a/models/hyvideo/text_encoder/llava/processing_llava.py b/models/hyvideo/text_encoder/llava/processing_llava.py new file mode 100644 index 0000000..6253e19 --- /dev/null +++ b/models/hyvideo/text_encoder/llava/processing_llava.py @@ -0,0 +1,203 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Processor class for Llava. +""" + +from typing import List, Union + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput, get_image_size, to_numpy_array +from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack, _validate_images_text_input_order +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import logging + + +logger = logging.get_logger(__name__) + + +class LlavaProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": False, + }, + "images_kwargs": {}, + } + + +class LlavaProcessor(ProcessorMixin): + r""" + Constructs a LLaVa processor which wraps a LLaVa image processor and a LLaMa tokenizer into a single processor. + + [`LlavaProcessor`] offers all the functionalities of [`LlavaImageProcessor`] and [`LlamaTokenizerFast`]. See the + [`~LlavaProcessor.__call__`] and [`~LlavaProcessor.decode`] for more information. + + Args: + image_processor ([`LlavaImageProcessor`], *optional*): + The image processor is a required input. + tokenizer ([`LlamaTokenizerFast`], *optional*): + The tokenizer is a required input. + patch_size (`int`, *optional*): + Patch size from the vision tower. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Shoudl be same as in model's config + chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages + in a chat into a tokenizable string. + image_token (`str`, *optional*, defaults to `""`): + Special token used to denote image location. + num_additional_image_tokens (`int`, *optional*, defaults to 0): + Number of additional tokens added to the image embeddings, such as CLS (+1). If the backbone has no CLS or other + extra tokens appended, no need to set this arg. + """ + + attributes = ["image_processor", "tokenizer"] + valid_kwargs = [ + "chat_template", + "patch_size", + "vision_feature_select_strategy", + "image_token", + "num_additional_image_tokens", + ] + image_processor_class = "AutoImageProcessor" + tokenizer_class = "AutoTokenizer" + + def __init__( + self, + image_processor=None, + tokenizer=None, + patch_size=None, + vision_feature_select_strategy=None, + chat_template=None, + image_token="", # set the default and let users change if they have peculiar special tokens in rare cases + num_additional_image_tokens=0, + **kwargs, + ): + self.patch_size = patch_size + self.num_additional_image_tokens = num_additional_image_tokens + self.vision_feature_select_strategy = vision_feature_select_strategy + self.image_token = tokenizer.image_token if hasattr(tokenizer, "image_token") else image_token + self.image_token_id = ( + tokenizer.image_token_id + if getattr(tokenizer, "image_token_id", None) + else tokenizer.convert_tokens_to_ids(self.image_token) + ) + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + def __call__( + self, + images: ImageInput = None, + text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None, + audio=None, + videos=None, + **kwargs: Unpack[LlavaProcessorKwargs], + ) -> BatchFeature: + """ + Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text` + and `kwargs` arguments to LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`] if `text` is not `None` to encode + the text. To prepare the image(s), this method forwards the `images` and `kwrags` arguments to + CLIPImageProcessor's [`~CLIPImageProcessor.__call__`] if `images` is not `None`. Please refer to the docstring + of the above two methods for more information. + + Args: + images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`): + The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch + tensor. Both channels-first and channels-last formats are supported. + text (`str`, `List[str]`, `List[List[str]]`): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Acceptable values are: + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + - `'jax'`: Return JAX `jnp.ndarray` objects. + + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. + """ + if images is None and text is None: + raise ValueError("You have to specify at least one of `images` or `text`.") + + # check if images and text inputs are reversed for BC + images, text = _validate_images_text_input_order(images, text) + + output_kwargs = self._merge_kwargs( + LlavaProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + if images is not None: + image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + else: + image_inputs = {} + + if isinstance(text, str): + text = [text] + elif not isinstance(text, list) and not isinstance(text[0], str): + raise ValueError("Invalid input text. Please provide a string, or a list of strings") + + # try to expand inputs in processing if we have the necessary parts + prompt_strings = text + if image_inputs.get("pixel_values") is not None: + # Replace the image token with the expanded image token sequence + pixel_values = image_inputs["pixel_values"] + height, width = get_image_size(to_numpy_array(pixel_values[0])) + num_image_tokens = (height // self.patch_size) * ( + width // self.patch_size + ) + self.num_additional_image_tokens + if self.vision_feature_select_strategy == "default": + num_image_tokens -= 1 + + prompt_strings = [] + for sample in text: + sample = sample.replace(self.image_token, self.image_token * num_image_tokens) + prompt_strings.append(sample) + + text_inputs = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"]) + return BatchFeature(data={**text_inputs, **image_inputs}) + + # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Llama + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please + refer to the docstring of this method for more information. + """ + return self.tokenizer.batch_decode(*args, **kwargs) + + # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode with CLIP->Llama + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + """ + return self.tokenizer.decode(*args, **kwargs) + + @property + # Copied from transformers.models.clip.processing_clip.CLIPProcessor.model_input_names + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) + + +__all__ = ["LlavaProcessor"] diff --git a/hyvideo/utils/__init__.py b/models/hyvideo/utils/__init__.py similarity index 100% rename from hyvideo/utils/__init__.py rename to models/hyvideo/utils/__init__.py diff --git a/hyvideo/utils/data_utils.py b/models/hyvideo/utils/data_utils.py similarity index 100% rename from hyvideo/utils/data_utils.py rename to models/hyvideo/utils/data_utils.py diff --git a/hyvideo/utils/file_utils.py b/models/hyvideo/utils/file_utils.py similarity index 100% rename from hyvideo/utils/file_utils.py rename to models/hyvideo/utils/file_utils.py diff --git a/hyvideo/utils/helpers.py b/models/hyvideo/utils/helpers.py similarity index 100% rename from hyvideo/utils/helpers.py rename to models/hyvideo/utils/helpers.py diff --git a/hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py b/models/hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py similarity index 100% rename from hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py rename to models/hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py diff --git a/hyvideo/vae/__init__.py b/models/hyvideo/vae/__init__.py similarity index 100% rename from hyvideo/vae/__init__.py rename to models/hyvideo/vae/__init__.py diff --git a/hyvideo/vae/autoencoder_kl_causal_3d.py b/models/hyvideo/vae/autoencoder_kl_causal_3d.py similarity index 100% rename from hyvideo/vae/autoencoder_kl_causal_3d.py rename to models/hyvideo/vae/autoencoder_kl_causal_3d.py diff --git a/hyvideo/vae/unet_causal_3d_blocks.py b/models/hyvideo/vae/unet_causal_3d_blocks.py similarity index 100% rename from hyvideo/vae/unet_causal_3d_blocks.py rename to models/hyvideo/vae/unet_causal_3d_blocks.py diff --git a/hyvideo/vae/vae.py b/models/hyvideo/vae/vae.py similarity index 100% rename from hyvideo/vae/vae.py rename to models/hyvideo/vae/vae.py diff --git a/models/ltx_video/__init__.py b/models/ltx_video/__init__.py new file mode 100644 index 0000000..3a3898e --- /dev/null +++ b/models/ltx_video/__init__.py @@ -0,0 +1,2 @@ +from .ltxv import LTXV +from . import ltxv_handler \ No newline at end of file diff --git a/ltx_video/configs/ltxv-13b-0.9.7-dev.original.yaml b/models/ltx_video/configs/ltxv-13b-0.9.7-dev.original.yaml similarity index 100% rename from ltx_video/configs/ltxv-13b-0.9.7-dev.original.yaml rename to models/ltx_video/configs/ltxv-13b-0.9.7-dev.original.yaml diff --git a/ltx_video/configs/ltxv-13b-0.9.7-dev.yaml b/models/ltx_video/configs/ltxv-13b-0.9.7-dev.yaml similarity index 100% rename from ltx_video/configs/ltxv-13b-0.9.7-dev.yaml rename to models/ltx_video/configs/ltxv-13b-0.9.7-dev.yaml diff --git a/ltx_video/configs/ltxv-13b-0.9.7-distilled.yaml b/models/ltx_video/configs/ltxv-13b-0.9.7-distilled.yaml similarity index 100% rename from ltx_video/configs/ltxv-13b-0.9.7-distilled.yaml rename to models/ltx_video/configs/ltxv-13b-0.9.7-distilled.yaml diff --git a/ltx_video/configs/ltxv-13b-0.9.8-dev.yaml b/models/ltx_video/configs/ltxv-13b-0.9.8-dev.yaml similarity index 100% rename from ltx_video/configs/ltxv-13b-0.9.8-dev.yaml rename to models/ltx_video/configs/ltxv-13b-0.9.8-dev.yaml diff --git a/ltx_video/configs/ltxv-13b-0.9.8-distilled.yaml b/models/ltx_video/configs/ltxv-13b-0.9.8-distilled.yaml similarity index 100% rename from ltx_video/configs/ltxv-13b-0.9.8-distilled.yaml rename to models/ltx_video/configs/ltxv-13b-0.9.8-distilled.yaml diff --git a/ltx_video/configs/ltxv-2b-0.9.6-dev.yaml b/models/ltx_video/configs/ltxv-2b-0.9.6-dev.yaml similarity index 100% rename from ltx_video/configs/ltxv-2b-0.9.6-dev.yaml rename to models/ltx_video/configs/ltxv-2b-0.9.6-dev.yaml diff --git a/ltx_video/configs/ltxv-2b-0.9.6-distilled.yaml b/models/ltx_video/configs/ltxv-2b-0.9.6-distilled.yaml similarity index 100% rename from ltx_video/configs/ltxv-2b-0.9.6-distilled.yaml rename to models/ltx_video/configs/ltxv-2b-0.9.6-distilled.yaml diff --git a/ltx_video/ltxv.py b/models/ltx_video/ltxv.py similarity index 98% rename from ltx_video/ltxv.py rename to models/ltx_video/ltxv.py index 34bae13..e71ac4f 100644 --- a/ltx_video/ltxv.py +++ b/models/ltx_video/ltxv.py @@ -7,7 +7,7 @@ from pathlib import Path from diffusers.utils import logging from typing import Optional, List, Union import yaml -from wan.utils.utils import calculate_new_dimensions +from shared.utils.utils import calculate_new_dimensions import imageio import json import numpy as np @@ -605,16 +605,4 @@ def load_media_file( raise Exception("video format not supported") return media_tensor -def query_model_def(model_type, model_def): - LTXV_config = model_def.get("LTXV_config", "") - distilled= "distilled" in LTXV_config - model_def_output = { - "no_guidance": True, - } - if distilled: - model_def_output.update({ - "lock_inference_steps": True, - "no_negative_prompt" : True, - }) - - return model_def_output \ No newline at end of file + diff --git a/models/ltx_video/ltxv_handler.py b/models/ltx_video/ltxv_handler.py new file mode 100644 index 0000000..b4190c8 --- /dev/null +++ b/models/ltx_video/ltxv_handler.py @@ -0,0 +1,92 @@ +import torch + + +def get_ltxv_text_encoder_filename(text_encoder_quantization): + text_encoder_filename = "ckpts/T5_xxl_1.1/T5_xxl_1.1_enc_bf16.safetensors" + if text_encoder_quantization =="int8": + text_encoder_filename = text_encoder_filename.replace("bf16", "quanto_bf16_int8") + return text_encoder_filename + +class family_handler(): + @staticmethod + def query_model_def(base_model_type, model_def): + LTXV_config = model_def.get("LTXV_config", "") + distilled= "distilled" in LTXV_config + extra_model_def = { + "no_guidance": True, + } + if distilled: + extra_model_def.update({ + "lock_inference_steps": True, + "no_negative_prompt" : True, + }) + + + extra_model_def["fps"] = 30 + extra_model_def["frames_minimum"] = 17 + extra_model_def["frames_steps"] = 8 + extra_model_def["sliding_window"] = True + + return extra_model_def + + @staticmethod + def query_supported_types(): + return ["ltxv_13B"] + + @staticmethod + def query_family_maps(): + return {}, {} + + @staticmethod + def get_rgb_factors(model_type): + from shared.RGB_factors import get_rgb_factors + latent_rgb_factors, latent_rgb_factors_bias = get_rgb_factors("ltxv") + return latent_rgb_factors, latent_rgb_factors_bias + + @staticmethod + def query_model_family(): + return "ltxv" + + @staticmethod + def query_family_infos(): + return {"ltxv":(10, "LTX Video")} + + @staticmethod + def get_vae_block_size(base_model_type): + return 32 + + @staticmethod + def query_model_files(computeList, base_model_type, model_filename, text_encoder_quantization): + text_encoder_filename = get_ltxv_text_encoder_filename(text_encoder_quantization) + return { + "repoId" : "DeepBeepMeep/LTX_Video", + "sourceFolderList" : ["T5_xxl_1.1", "" ], + "fileList" : [ ["added_tokens.json", "special_tokens_map.json", "spiece.model", "tokenizer_config.json"] + computeList(text_encoder_filename), ["ltxv_0.9.7_VAE.safetensors", "ltxv_0.9.7_spatial_upscaler.safetensors", "ltxv_scheduler.json"] + computeList(model_filename) ] + } + + + @staticmethod + def load_model(model_filename, model_type, base_model_type, model_def, quantizeTransformer = False, text_encoder_quantization = None, dtype = torch.bfloat16, VAE_dtype = torch.float32, mixed_precision_transformer = False, save_quantized = False): + from .ltxv import LTXV + + ltxv_model = LTXV( + model_filepath = model_filename, + text_encoder_filepath = get_ltxv_text_encoder_filename(text_encoder_quantization), + model_type = model_type, + base_model_type = base_model_type, + model_def = model_def, + dtype = dtype, + # quantizeTransformer = quantizeTransformer, + VAE_dtype = VAE_dtype, + mixed_precision_transformer = mixed_precision_transformer + ) + + pipeline = ltxv_model.pipeline + pipe = {"transformer" : pipeline.video_pipeline.transformer, "vae" : pipeline.vae, "text_encoder" : pipeline.video_pipeline.text_encoder, "latent_upsampler" : pipeline.latent_upsampler} + + return ltxv_model, pipe + + @staticmethod + def update_default_settings(base_model_type, model_def, ui_defaults): + pass + \ No newline at end of file diff --git a/ltx_video/__init__.py b/models/ltx_video/models/__init__.py similarity index 100% rename from ltx_video/__init__.py rename to models/ltx_video/models/__init__.py diff --git a/ltx_video/models/__init__.py b/models/ltx_video/models/autoencoders/__init__.py similarity index 100% rename from ltx_video/models/__init__.py rename to models/ltx_video/models/autoencoders/__init__.py diff --git a/ltx_video/models/autoencoders/causal_conv3d.py b/models/ltx_video/models/autoencoders/causal_conv3d.py similarity index 100% rename from ltx_video/models/autoencoders/causal_conv3d.py rename to models/ltx_video/models/autoencoders/causal_conv3d.py diff --git a/ltx_video/models/autoencoders/causal_video_autoencoder.py b/models/ltx_video/models/autoencoders/causal_video_autoencoder.py similarity index 99% rename from ltx_video/models/autoencoders/causal_video_autoencoder.py rename to models/ltx_video/models/autoencoders/causal_video_autoencoder.py index 5f05932..daed704 100644 --- a/ltx_video/models/autoencoders/causal_video_autoencoder.py +++ b/models/ltx_video/models/autoencoders/causal_video_autoencoder.py @@ -15,12 +15,12 @@ from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbedding from safetensors import safe_open -from ltx_video.models.autoencoders.conv_nd_factory import make_conv_nd, make_linear_nd -from ltx_video.models.autoencoders.pixel_norm import PixelNorm -from ltx_video.models.autoencoders.pixel_shuffle import PixelShuffleND -from ltx_video.models.autoencoders.vae import AutoencoderKLWrapper -from ltx_video.models.transformers.attention import Attention -from ltx_video.utils.diffusers_config_mapping import ( +from ..autoencoders.conv_nd_factory import make_conv_nd, make_linear_nd +from ...models.autoencoders.pixel_norm import PixelNorm +from ...models.autoencoders.pixel_shuffle import PixelShuffleND +from ...models.autoencoders.vae import AutoencoderKLWrapper +from ...models.transformers.attention import Attention +from ...utils.diffusers_config_mapping import ( diffusers_and_ours_config_mapping, make_hashable_key, VAE_KEYS_RENAME_DICT, diff --git a/ltx_video/models/autoencoders/conv_nd_factory.py b/models/ltx_video/models/autoencoders/conv_nd_factory.py similarity index 94% rename from ltx_video/models/autoencoders/conv_nd_factory.py rename to models/ltx_video/models/autoencoders/conv_nd_factory.py index 718c69b..59a3fc0 100644 --- a/ltx_video/models/autoencoders/conv_nd_factory.py +++ b/models/ltx_video/models/autoencoders/conv_nd_factory.py @@ -2,8 +2,8 @@ from typing import Tuple, Union import torch -from ltx_video.models.autoencoders.dual_conv3d import DualConv3d -from ltx_video.models.autoencoders.causal_conv3d import CausalConv3d +from ..autoencoders.dual_conv3d import DualConv3d +from ..autoencoders.causal_conv3d import CausalConv3d def make_conv_nd( diff --git a/ltx_video/models/autoencoders/dual_conv3d.py b/models/ltx_video/models/autoencoders/dual_conv3d.py similarity index 100% rename from ltx_video/models/autoencoders/dual_conv3d.py rename to models/ltx_video/models/autoencoders/dual_conv3d.py diff --git a/ltx_video/models/autoencoders/latent_upsampler.py b/models/ltx_video/models/autoencoders/latent_upsampler.py similarity index 98% rename from ltx_video/models/autoencoders/latent_upsampler.py rename to models/ltx_video/models/autoencoders/latent_upsampler.py index 4a76bc2..f666d2f 100644 --- a/ltx_video/models/autoencoders/latent_upsampler.py +++ b/models/ltx_video/models/autoencoders/latent_upsampler.py @@ -9,7 +9,7 @@ from einops import rearrange from diffusers import ConfigMixin, ModelMixin from safetensors.torch import safe_open -from ltx_video.models.autoencoders.pixel_shuffle import PixelShuffleND +from ...models.autoencoders.pixel_shuffle import PixelShuffleND class ResBlock(nn.Module): diff --git a/ltx_video/models/autoencoders/pixel_norm.py b/models/ltx_video/models/autoencoders/pixel_norm.py similarity index 100% rename from ltx_video/models/autoencoders/pixel_norm.py rename to models/ltx_video/models/autoencoders/pixel_norm.py diff --git a/ltx_video/models/autoencoders/pixel_shuffle.py b/models/ltx_video/models/autoencoders/pixel_shuffle.py similarity index 100% rename from ltx_video/models/autoencoders/pixel_shuffle.py rename to models/ltx_video/models/autoencoders/pixel_shuffle.py diff --git a/ltx_video/models/autoencoders/vae.py b/models/ltx_video/models/autoencoders/vae.py similarity index 99% rename from ltx_video/models/autoencoders/vae.py rename to models/ltx_video/models/autoencoders/vae.py index c1135ba..a0ce1c4 100644 --- a/ltx_video/models/autoencoders/vae.py +++ b/models/ltx_video/models/autoencoders/vae.py @@ -10,7 +10,7 @@ from diffusers.models.autoencoders.vae import ( DiagonalGaussianDistribution, ) from diffusers.models.modeling_outputs import AutoencoderKLOutput -from ltx_video.models.autoencoders.conv_nd_factory import make_conv_nd +from ...models.autoencoders.conv_nd_factory import make_conv_nd class AutoencoderKLWrapper(ModelMixin, ConfigMixin): diff --git a/ltx_video/models/autoencoders/vae_encode.py b/models/ltx_video/models/autoencoders/vae_encode.py similarity index 98% rename from ltx_video/models/autoencoders/vae_encode.py rename to models/ltx_video/models/autoencoders/vae_encode.py index b7d2476..4b6a5c4 100644 --- a/ltx_video/models/autoencoders/vae_encode.py +++ b/models/ltx_video/models/autoencoders/vae_encode.py @@ -5,10 +5,10 @@ from einops import rearrange from torch import Tensor -from ltx_video.models.autoencoders.causal_video_autoencoder import ( +from ...models.autoencoders.causal_video_autoencoder import ( CausalVideoAutoencoder, ) -from ltx_video.models.autoencoders.video_autoencoder import ( +from ...models.autoencoders.video_autoencoder import ( Downsample3D, VideoAutoencoder, ) diff --git a/ltx_video/models/autoencoders/video_autoencoder.py b/models/ltx_video/models/autoencoders/video_autoencoder.py similarity index 99% rename from ltx_video/models/autoencoders/video_autoencoder.py rename to models/ltx_video/models/autoencoders/video_autoencoder.py index 3c7926c..dbb2bcd 100644 --- a/ltx_video/models/autoencoders/video_autoencoder.py +++ b/models/ltx_video/models/autoencoders/video_autoencoder.py @@ -11,10 +11,10 @@ from torch.nn import functional from diffusers.utils import logging -from ltx_video.utils.torch_utils import Identity -from ltx_video.models.autoencoders.conv_nd_factory import make_conv_nd, make_linear_nd -from ltx_video.models.autoencoders.pixel_norm import PixelNorm -from ltx_video.models.autoencoders.vae import AutoencoderKLWrapper +from ...utils.torch_utils import Identity +from ...models.autoencoders.conv_nd_factory import make_conv_nd, make_linear_nd +from ...models.autoencoders.pixel_norm import PixelNorm +from ...models.autoencoders.vae import AutoencoderKLWrapper logger = logging.get_logger(__name__) diff --git a/ltx_video/models/autoencoders/__init__.py b/models/ltx_video/models/transformers/__init__.py similarity index 100% rename from ltx_video/models/autoencoders/__init__.py rename to models/ltx_video/models/transformers/__init__.py diff --git a/ltx_video/models/transformers/attention.py b/models/ltx_video/models/transformers/attention.py similarity index 99% rename from ltx_video/models/transformers/attention.py rename to models/ltx_video/models/transformers/attention.py index a7b4555..a87a8a0 100644 --- a/ltx_video/models/transformers/attention.py +++ b/models/ltx_video/models/transformers/attention.py @@ -19,15 +19,9 @@ from diffusers.utils import deprecate, logging from diffusers.utils.torch_utils import maybe_allow_in_graph from einops import rearrange from torch import nn -from wan.modules.attention import pay_attention -from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy +from shared.attention import pay_attention +from ...utils.skip_layer_strategy import SkipLayerStrategy -try: - from torch_xla.experimental.custom_kernel import flash_attention -except ImportError: - # workaround for automatic tests. Currently this function is manually patched - # to the torch_xla lib on setup of container - pass # code adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py diff --git a/ltx_video/models/transformers/embeddings.py b/models/ltx_video/models/transformers/embeddings.py similarity index 100% rename from ltx_video/models/transformers/embeddings.py rename to models/ltx_video/models/transformers/embeddings.py diff --git a/ltx_video/models/transformers/symmetric_patchifier.py b/models/ltx_video/models/transformers/symmetric_patchifier.py similarity index 100% rename from ltx_video/models/transformers/symmetric_patchifier.py rename to models/ltx_video/models/transformers/symmetric_patchifier.py diff --git a/ltx_video/models/transformers/transformer3d.py b/models/ltx_video/models/transformers/transformer3d.py similarity index 98% rename from ltx_video/models/transformers/transformer3d.py rename to models/ltx_video/models/transformers/transformer3d.py index e182f21..c90baeb 100644 --- a/ltx_video/models/transformers/transformer3d.py +++ b/models/ltx_video/models/transformers/transformer3d.py @@ -16,10 +16,10 @@ from diffusers.utils import BaseOutput, is_torch_version from diffusers.utils import logging from torch import nn from safetensors import safe_open -from ltx_video.models.transformers.attention import BasicTransformerBlock, reshape_hidden_states, restore_hidden_states_shape -from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy +from .attention import BasicTransformerBlock, reshape_hidden_states, restore_hidden_states_shape +from ...utils.skip_layer_strategy import SkipLayerStrategy -from ltx_video.utils.diffusers_config_mapping import ( +from ...utils.diffusers_config_mapping import ( diffusers_and_ours_config_mapping, make_hashable_key, TRANSFORMER_KEYS_RENAME_DICT, diff --git a/ltx_video/models/transformers/__init__.py b/models/ltx_video/pipelines/__init__.py similarity index 100% rename from ltx_video/models/transformers/__init__.py rename to models/ltx_video/pipelines/__init__.py diff --git a/ltx_video/pipelines/crf_compressor.py b/models/ltx_video/pipelines/crf_compressor.py similarity index 100% rename from ltx_video/pipelines/crf_compressor.py rename to models/ltx_video/pipelines/crf_compressor.py diff --git a/ltx_video/pipelines/pipeline_ltx_video.py b/models/ltx_video/pipelines/pipeline_ltx_video.py similarity index 99% rename from ltx_video/pipelines/pipeline_ltx_video.py rename to models/ltx_video/pipelines/pipeline_ltx_video.py index 8bb6d27..f98eb13 100644 --- a/ltx_video/pipelines/pipeline_ltx_video.py +++ b/models/ltx_video/pipelines/pipeline_ltx_video.py @@ -24,22 +24,22 @@ from transformers import ( AutoTokenizer, ) -from ltx_video.models.autoencoders.causal_video_autoencoder import ( +from ..models.autoencoders.causal_video_autoencoder import ( CausalVideoAutoencoder, ) -from ltx_video.models.autoencoders.vae_encode import ( +from ..models.autoencoders.vae_encode import ( get_vae_size_scale_factor, latent_to_pixel_coords, vae_decode, vae_encode, ) -from ltx_video.models.transformers.symmetric_patchifier import Patchifier -from ltx_video.models.transformers.transformer3d import Transformer3DModel -from ltx_video.schedulers.rf import TimestepShifter -from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy -from ltx_video.utils.prompt_enhance_utils import generate_cinematic_prompt -from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler -from ltx_video.models.autoencoders.vae_encode import ( +from ..models.transformers.symmetric_patchifier import Patchifier +from ..models.transformers.transformer3d import Transformer3DModel +from ..schedulers.rf import TimestepShifter +from ..utils.skip_layer_strategy import SkipLayerStrategy +from ..utils.prompt_enhance_utils import generate_cinematic_prompt +from ..models.autoencoders.latent_upsampler import LatentUpsampler +from ..models.autoencoders.vae_encode import ( un_normalize_latents, normalize_latents, ) diff --git a/ltx_video/pipelines/__init__.py b/models/ltx_video/schedulers/__init__.py similarity index 100% rename from ltx_video/pipelines/__init__.py rename to models/ltx_video/schedulers/__init__.py diff --git a/ltx_video/schedulers/rf.py b/models/ltx_video/schedulers/rf.py similarity index 99% rename from ltx_video/schedulers/rf.py rename to models/ltx_video/schedulers/rf.py index 2cf99da..bced26a 100644 --- a/ltx_video/schedulers/rf.py +++ b/models/ltx_video/schedulers/rf.py @@ -14,9 +14,9 @@ from torch import Tensor from safetensors import safe_open -from ltx_video.utils.torch_utils import append_dims +from ..utils.torch_utils import append_dims -from ltx_video.utils.diffusers_config_mapping import ( +from ..utils.diffusers_config_mapping import ( diffusers_and_ours_config_mapping, make_hashable_key, ) diff --git a/ltx_video/schedulers/__init__.py b/models/ltx_video/utils/__init__.py similarity index 100% rename from ltx_video/schedulers/__init__.py rename to models/ltx_video/utils/__init__.py diff --git a/ltx_video/utils/diffusers_config_mapping.py b/models/ltx_video/utils/diffusers_config_mapping.py similarity index 100% rename from ltx_video/utils/diffusers_config_mapping.py rename to models/ltx_video/utils/diffusers_config_mapping.py diff --git a/ltx_video/utils/prompt_enhance_utils.py b/models/ltx_video/utils/prompt_enhance_utils.py similarity index 100% rename from ltx_video/utils/prompt_enhance_utils.py rename to models/ltx_video/utils/prompt_enhance_utils.py diff --git a/ltx_video/utils/skip_layer_strategy.py b/models/ltx_video/utils/skip_layer_strategy.py similarity index 100% rename from ltx_video/utils/skip_layer_strategy.py rename to models/ltx_video/utils/skip_layer_strategy.py diff --git a/ltx_video/utils/torch_utils.py b/models/ltx_video/utils/torch_utils.py similarity index 100% rename from ltx_video/utils/torch_utils.py rename to models/ltx_video/utils/torch_utils.py diff --git a/wan/__init__.py b/models/wan/__init__.py similarity index 50% rename from wan/__init__.py rename to models/wan/__init__.py index 1688425..fe3be71 100644 --- a/wan/__init__.py +++ b/models/wan/__init__.py @@ -1,3 +1,4 @@ from . import configs, distributed, modules from .any2video import WanAny2V -from .diffusion_forcing import DTT2V \ No newline at end of file +from .diffusion_forcing import DTT2V +from . import wan_handler, df_handler diff --git a/wan/any2video.py b/models/wan/any2video.py similarity index 95% rename from wan/any2video.py rename to models/wan/any2video.py index bea70c5..4a20255 100644 --- a/wan/any2video.py +++ b/models/wan/any2video.py @@ -22,14 +22,16 @@ from .distributed.fsdp import shard_model from .modules.model import WanModel from .modules.t5 import T5EncoderModel from .modules.vae import WanVAE +from .modules.vae2_2 import Wan2_2_VAE + from .modules.clip import CLIPModel -from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler, +from shared.utils.fm_solvers import (FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps) -from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler -from wan.modules.posemb_layers import get_rotary_pos_embed -from .utils.vace_preprocessor import VaceVideoProcessor -from wan.utils.basic_flowmatch import FlowMatchScheduler -from wan.utils.utils import get_outpainting_frame_location, resize_lanczos, calculate_new_dimensions +from shared.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler +from .modules.posemb_layers import get_rotary_pos_embed +from shared.utils.vace_preprocessor import VaceVideoProcessor +from shared.utils.basic_flowmatch import FlowMatchScheduler +from shared.utils.utils import get_outpainting_frame_location, resize_lanczos, calculate_new_dimensions from .multitalk.multitalk_utils import MomentumBuffer, adaptive_projected_guidance, match_and_blend_colors, match_and_blend_colors_with_mask from mmgp import safetensors2 @@ -97,11 +99,19 @@ class WanAny2V: config.clip_checkpoint), tokenizer_path=os.path.join(checkpoint_dir , config.clip_tokenizer)) - self.vae_stride = config.vae_stride + + if base_model_type in ["ti2v_2_2"]: + self.vae_stride = (4, 16, 16) + vae_checkpoint = "Wan2.2_VAE.safetensors" + vae = Wan2_2_VAE + else: + self.vae_stride = config.vae_stride + vae_checkpoint = "Wan2.1_VAE.safetensors" + vae = WanVAE self.patch_size = config.patch_size - self.vae = WanVAE( - vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), dtype= VAE_dtype, + self.vae = vae( + vae_pth=os.path.join(checkpoint_dir, vae_checkpoint), dtype= VAE_dtype, device=self.device) # config_filename= "configs/t2v_1.3B.json" @@ -115,7 +125,11 @@ class WanAny2V: # forcedConfigPath = base_config_file = f"configs/flf2v_720p.json" # model_filename[1] = xmodel_filename - if self.transformer_switch: + source = model_def.get("source", None) + + if source is not None: + self.model = offload.fast_load_transformers_model(source, modelClass=WanModel, writable_tensors= False, forcedConfigPath= base_config_file) + elif self.transformer_switch: shared_modules= {} self.model = offload.fast_load_transformers_model(model_filename[:1], modules = model_filename[2:], modelClass=WanModel,do_quantize= quantizeTransformer and not save_quantized, writable_tensors= False, defaultConfigPath=base_config_file , forcedConfigPath= forcedConfigPath, return_shared_modules= shared_modules) self.model2 = offload.fast_load_transformers_model(model_filename[1:2], modules = shared_modules, modelClass=WanModel,do_quantize= quantizeTransformer and not save_quantized, writable_tensors= False, defaultConfigPath=base_config_file , forcedConfigPath= forcedConfigPath) @@ -138,6 +152,10 @@ class WanAny2V: self.model.eval().requires_grad_(False) if self.model2 is not None: self.model2.eval().requires_grad_(False) + if not source is None: + from wgp import save_model + save_model(self.model, model_type, dtype, None) + if save_quantized: from wgp import save_quantized_model save_quantized_model(self.model, model_type, model_filename[0], dtype, base_config_file) @@ -233,7 +251,7 @@ class WanAny2V: return [torch.cat([zz, mm], dim=0) for zz, mm in zip(z, m)] def fit_image_into_canvas(self, ref_img, image_size, canvas_tf_bg, device, fill_max = False, outpainting_dims = None, return_mask = False): - from wan.utils.utils import save_image + from shared.utils.utils import save_image ref_width, ref_height = ref_img.size if (ref_height, ref_width) == image_size and outpainting_dims == None: ref_img = TF.to_tensor(ref_img).sub_(0.5).div_(0.5).unsqueeze(1) @@ -478,11 +496,12 @@ class WanAny2V: fantasy = model_type in ["fantasy"] multitalk = model_type in ["multitalk", "vace_multitalk_14B"] recam = model_type in ["recam_1.3B"] + ti2v = model_type in ["ti2v_2_2"] ref_images_count = 0 trim_frames = 0 extended_overlapped_latents = None - + timestep_injection = False lat_frames = int((frame_num - 1) // self.vae_stride[0]) + 1 # image2video if model_type in ["i2v", "i2v_2_2", "fantasy", "multitalk", "flf2v_720p"]: @@ -589,7 +608,7 @@ class WanAny2V: source_latents = self.vae.encode([input_video])[0] #.to(dtype=self.dtype, device=self.device) del input_video # Process target camera (recammaster) - from wan.utils.cammmaster_tools import get_camera_embedding + from shared.utils.cammmaster_tools import get_camera_embedding cam_emb = get_camera_embedding(target_camera) cam_emb = cam_emb.to(dtype=self.dtype, device=self.device) kwargs['cam_emb'] = cam_emb @@ -632,6 +651,14 @@ class WanAny2V: ref_images_count = input_ref_images.shape[1] if input_ref_images != None else 0 trim_frames = input_ref_images.shape[1] + if ti2v: + if input_video is None: + height, width = (height // 32) * 32, (width // 32) * 32 + else: + height, width = input_video.shape[-2:] + source_latents = self.vae.encode([input_video])[0].unsqueeze(0) + timestep_injection = True + # Vace if vace : # vace context encode @@ -671,7 +698,7 @@ class WanAny2V: target_shape = (self.vae.model.z_dim, lat_frames + ref_images_count, height // self.vae_stride[1], width // self.vae_stride[2]) if multitalk and audio_proj != None: - from wan.multitalk.multitalk import get_target_masks + from .multitalk.multitalk import get_target_masks audio_proj = [audio.to(self.dtype) for audio in audio_proj] human_no = len(audio_proj[0]) token_ref_target_masks = get_target_masks(human_no, lat_h, lat_w, height, width, face_scale = 0.05, bbox = speakers_bboxes).to(self.dtype) if human_no > 1 else None @@ -717,7 +744,7 @@ class WanAny2V: # init denoising updated_num_steps= len(timesteps) if callback != None: - from wan.utils.loras_mutipliers import update_loras_slists + from shared.utils.loras_mutipliers import update_loras_slists model_switch_step = updated_num_steps for i, t in enumerate(timesteps): if t <= switch_threshold: @@ -748,6 +775,12 @@ class WanAny2V: offload.set_step_no_for_lora(trans, i) timestep = torch.stack([t]) + + if timestep_injection: + latents[:, :, :source_latents.shape[2]] = source_latents + timestep = torch.full((target_shape[-3],), t, dtype=torch.int64, device=latents.device) + timestep[:source_latents.shape[2]] = 0 + kwargs.update({"t": timestep, "current_step": i}) kwargs["slg_layers"] = slg_layers if int(slg_start * sampling_steps) <= i < int(slg_end * sampling_steps) else None @@ -887,6 +920,9 @@ class WanAny2V: callback(i, latents_preview[0], False) latents_preview = None + if timestep_injection: + latents[:, :, :source_latents.shape[2]] = source_latents + if vace and ref_images_count > 0: latents = latents[:, :, ref_images_count:] if trim_frames > 0: latents= latents[:, :,:-trim_frames] if return_latent_slice != None: @@ -923,8 +959,5 @@ class WanAny2V: setattr(target, "vace", module ) delattr(model, "vace_blocks") -def query_model_def(model_type, model_def): - if "URLs2" in model_def: - return { "no_steps_skipping":True} - else: - return None \ No newline at end of file + + diff --git a/wan/camera_extrinsics.json b/models/wan/camera_extrinsics.json similarity index 100% rename from wan/camera_extrinsics.json rename to models/wan/camera_extrinsics.json diff --git a/wan/configs/__init__.py b/models/wan/configs/__init__.py similarity index 100% rename from wan/configs/__init__.py rename to models/wan/configs/__init__.py diff --git a/wan/configs/shared_config.py b/models/wan/configs/shared_config.py similarity index 100% rename from wan/configs/shared_config.py rename to models/wan/configs/shared_config.py diff --git a/wan/configs/wan_i2v_14B.py b/models/wan/configs/wan_i2v_14B.py similarity index 100% rename from wan/configs/wan_i2v_14B.py rename to models/wan/configs/wan_i2v_14B.py diff --git a/wan/configs/wan_t2v_14B.py b/models/wan/configs/wan_t2v_14B.py similarity index 100% rename from wan/configs/wan_t2v_14B.py rename to models/wan/configs/wan_t2v_14B.py diff --git a/wan/configs/wan_t2v_1_3B.py b/models/wan/configs/wan_t2v_1_3B.py similarity index 100% rename from wan/configs/wan_t2v_1_3B.py rename to models/wan/configs/wan_t2v_1_3B.py diff --git a/models/wan/df_handler.py b/models/wan/df_handler.py new file mode 100644 index 0000000..9ddeab9 --- /dev/null +++ b/models/wan/df_handler.py @@ -0,0 +1,86 @@ +import torch + +class family_handler(): + @staticmethod + def query_model_def(base_model_type, model_def): + extra_model_def = {} + if base_model_type in ["sky_df_14B"]: + fps = 24 + else: + fps = 16 + extra_model_def["fps"] =fps + extra_model_def["frames_minimum"] = 17 + extra_model_def["frames_steps"] = 20 + extra_model_def["sliding_window"] = True + extra_model_def["skip_layer_guidance"] = True + return extra_model_def + + @staticmethod + def query_supported_types(): + return ["sky_df_1.3B", "sky_df_14B"] + + + @staticmethod + def query_family_maps(): + models_eqv_map = { + "sky_df_1.3B" : "sky_df_14B", + } + + models_comp_map = { + "sky_df_14B": ["sky_df_1.3B"], + } + return models_eqv_map, models_comp_map + + + + @staticmethod + def query_model_family(): + return "wan" + + @staticmethod + def query_family_infos(): + return {} + + + + @staticmethod + def query_model_files(computeList, base_model_type, model_filename, text_encoder_quantization): + from .wan_handler import family_handler + return family_handler.query_model_files(computeList, base_model_type, model_filename, text_encoder_quantization) + + @staticmethod + def load_model(model_filename, model_type, base_model_type, model_def, quantizeTransformer = False, text_encoder_quantization = None, dtype = torch.bfloat16, VAE_dtype = torch.float32, mixed_precision_transformer = False, save_quantized= False): + from .configs import WAN_CONFIGS + from .wan_handler import family_handler + cfg = WAN_CONFIGS['t2v-14B'] + from . import DTT2V + wan_model = DTT2V( + config=cfg, + checkpoint_dir="ckpts", + model_filename=model_filename, + model_type = model_type, + model_def = model_def, + base_model_type=base_model_type, + text_encoder_filename= family_handler.get_wan_text_encoder_filename(text_encoder_quantization), + quantizeTransformer = quantizeTransformer, + dtype = dtype, + VAE_dtype = VAE_dtype, + mixed_precision_transformer = mixed_precision_transformer, + save_quantized = save_quantized + ) + + pipe = {"transformer": wan_model.model, "text_encoder" : wan_model.text_encoder.model, "vae": wan_model.vae.model } + return wan_model, pipe + + @staticmethod + def update_default_settings(base_model_type, model_def, ui_defaults): + ui_defaults.update({ + "guidance_scale": 6.0, + "flow_shift": 8, + "sliding_window_discard_last_frames" : 0, + "resolution": "1280x720" if "720" in base_model_type else "960x544", + "sliding_window_size" : 121 if "720" in base_model_type else 97, + "RIFLEx_setting": 2, + "guidance_scale": 6, + "flow_shift": 8, + }) \ No newline at end of file diff --git a/models/wan/diffusion_forcing copy.py b/models/wan/diffusion_forcing copy.py new file mode 100644 index 0000000..753fd45 --- /dev/null +++ b/models/wan/diffusion_forcing copy.py @@ -0,0 +1,479 @@ +import math +import os +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union +import logging +import numpy as np +import torch +from diffusers.image_processor import PipelineImageInput +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from tqdm import tqdm +from .modules.model import WanModel +from .modules.t5 import T5EncoderModel +from .modules.vae import WanVAE +from wan.modules.posemb_layers import get_rotary_pos_embed +from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler, + get_sampling_sigmas, retrieve_timesteps) +from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler + +class DTT2V: + + + def __init__( + self, + config, + checkpoint_dir, + rank=0, + model_filename = None, + text_encoder_filename = None, + quantizeTransformer = False, + dtype = torch.bfloat16, + ): + self.device = torch.device(f"cuda") + self.config = config + self.rank = rank + self.dtype = dtype + self.num_train_timesteps = config.num_train_timesteps + self.param_dtype = config.param_dtype + + self.text_encoder = T5EncoderModel( + text_len=config.text_len, + dtype=config.t5_dtype, + device=torch.device('cpu'), + checkpoint_path=text_encoder_filename, + tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), + shard_fn= None) + + self.vae_stride = config.vae_stride + self.patch_size = config.patch_size + + + self.vae = WanVAE( + vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), + device=self.device) + + logging.info(f"Creating WanModel from {model_filename}") + from mmgp import offload + + self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel,do_quantize= quantizeTransformer, writable_tensors= False, forcedConfigPath="config.json") + # offload.load_model_data(self.model, "recam.ckpt") + # self.model.cpu() + # offload.save_model(self.model, "recam.safetensors") + if self.dtype == torch.float16 and not "fp16" in model_filename: + self.model.to(self.dtype) + # offload.save_model(self.model, "t2v_fp16.safetensors",do_quantize=True) + if self.dtype == torch.float16: + self.vae.model.to(self.dtype) + self.model.eval().requires_grad_(False) + + self.scheduler = FlowUniPCMultistepScheduler() + + @property + def do_classifier_free_guidance(self) -> bool: + return self._guidance_scale > 1 + + def encode_image( + self, image: PipelineImageInput, height: int, width: int, num_frames: int, tile_size = 0, causal_block_size = 0 + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + # prefix_video + prefix_video = np.array(image.resize((width, height))).transpose(2, 0, 1) + prefix_video = torch.tensor(prefix_video).unsqueeze(1) # .to(image_embeds.dtype).unsqueeze(1) + if prefix_video.dtype == torch.uint8: + prefix_video = (prefix_video.float() / (255.0 / 2.0)) - 1.0 + prefix_video = prefix_video.to(self.device) + prefix_video = [self.vae.encode(prefix_video.unsqueeze(0), tile_size = tile_size)[0]] # [(c, f, h, w)] + if prefix_video[0].shape[1] % causal_block_size != 0: + truncate_len = prefix_video[0].shape[1] % causal_block_size + print("the length of prefix video is truncated for the casual block size alignment.") + prefix_video[0] = prefix_video[0][:, : prefix_video[0].shape[1] - truncate_len] + predix_video_latent_length = prefix_video[0].shape[1] + return prefix_video, predix_video_latent_length + + def prepare_latents( + self, + shape: Tuple[int], + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + ) -> torch.Tensor: + return randn_tensor(shape, generator, device=device, dtype=dtype) + + def generate_timestep_matrix( + self, + num_frames, + step_template, + base_num_frames, + ar_step=5, + num_pre_ready=0, + casual_block_size=1, + shrink_interval_with_mask=False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[tuple]]: + step_matrix, step_index = [], [] + update_mask, valid_interval = [], [] + num_iterations = len(step_template) + 1 + num_frames_block = num_frames // casual_block_size + base_num_frames_block = base_num_frames // casual_block_size + if base_num_frames_block < num_frames_block: + infer_step_num = len(step_template) + gen_block = base_num_frames_block + min_ar_step = infer_step_num / gen_block + assert ar_step >= min_ar_step, f"ar_step should be at least {math.ceil(min_ar_step)} in your setting" + # print(num_frames, step_template, base_num_frames, ar_step, num_pre_ready, casual_block_size, num_frames_block, base_num_frames_block) + step_template = torch.cat( + [ + torch.tensor([999], dtype=torch.int64, device=step_template.device), + step_template.long(), + torch.tensor([0], dtype=torch.int64, device=step_template.device), + ] + ) # to handle the counter in row works starting from 1 + pre_row = torch.zeros(num_frames_block, dtype=torch.long) + if num_pre_ready > 0: + pre_row[: num_pre_ready // casual_block_size] = num_iterations + + while torch.all(pre_row >= (num_iterations - 1)) == False: + new_row = torch.zeros(num_frames_block, dtype=torch.long) + for i in range(num_frames_block): + if i == 0 or pre_row[i - 1] >= ( + num_iterations - 1 + ): # the first frame or the last frame is completely denoised + new_row[i] = pre_row[i] + 1 + else: + new_row[i] = new_row[i - 1] - ar_step + new_row = new_row.clamp(0, num_iterations) + + update_mask.append( + (new_row != pre_row) & (new_row != num_iterations) + ) # False: no need to update， True: need to update + step_index.append(new_row) + step_matrix.append(step_template[new_row]) + pre_row = new_row + + # for long video we split into several sequences, base_num_frames is set to the model max length (for training) + terminal_flag = base_num_frames_block + if shrink_interval_with_mask: + idx_sequence = torch.arange(num_frames_block, dtype=torch.int64) + update_mask = update_mask[0] + update_mask_idx = idx_sequence[update_mask] + last_update_idx = update_mask_idx[-1].item() + terminal_flag = last_update_idx + 1 + # for i in range(0, len(update_mask)): + for curr_mask in update_mask: + if terminal_flag < num_frames_block and curr_mask[terminal_flag]: + terminal_flag += 1 + valid_interval.append((max(terminal_flag - base_num_frames_block, 0), terminal_flag)) + + step_update_mask = torch.stack(update_mask, dim=0) + step_index = torch.stack(step_index, dim=0) + step_matrix = torch.stack(step_matrix, dim=0) + + if casual_block_size > 1: + step_update_mask = step_update_mask.unsqueeze(-1).repeat(1, 1, casual_block_size).flatten(1).contiguous() + step_index = step_index.unsqueeze(-1).repeat(1, 1, casual_block_size).flatten(1).contiguous() + step_matrix = step_matrix.unsqueeze(-1).repeat(1, 1, casual_block_size).flatten(1).contiguous() + valid_interval = [(s * casual_block_size, e * casual_block_size) for s, e in valid_interval] + + return step_matrix, step_index, step_update_mask, valid_interval + + @torch.no_grad() + def generate( + self, + prompt: Union[str, List[str]], + negative_prompt: Union[str, List[str]] = "", + image: PipelineImageInput = None, + height: int = 480, + width: int = 832, + num_frames: int = 97, + num_inference_steps: int = 50, + shift: float = 1.0, + guidance_scale: float = 5.0, + seed: float = 0.0, + overlap_history: int = 17, + addnoise_condition: int = 0, + base_num_frames: int = 97, + ar_step: int = 5, + causal_block_size: int = 1, + causal_attention: bool = False, + fps: int = 24, + VAE_tile_size = 0, + joint_pass = False, + callback = None, + ): + generator = torch.Generator(device=self.device) + generator.manual_seed(seed) + # if base_num_frames > base_num_frames: + # causal_block_size = 0 + self._guidance_scale = guidance_scale + + i2v_extra_kwrags = {} + prefix_video = None + predix_video_latent_length = 0 + if image: + frame_width, frame_height = image.size + scale = min(height / frame_height, width / frame_width) + height = (int(frame_height * scale) // 16) * 16 + width = (int(frame_width * scale) // 16) * 16 + + prefix_video, predix_video_latent_length = self.encode_image(image, height, width, num_frames, tile_size=VAE_tile_size, causal_block_size=causal_block_size) + + latent_length = (num_frames - 1) // 4 + 1 + latent_height = height // 8 + latent_width = width // 8 + + prompt_embeds = self.text_encoder([prompt], self.device) + prompt_embeds = [u.to(self.dtype).to(self.device) for u in prompt_embeds] + if self.do_classifier_free_guidance: + negative_prompt_embeds = self.text_encoder([negative_prompt], self.device) + negative_prompt_embeds = [u.to(self.dtype).to(self.device) for u in negative_prompt_embeds] + + + + self.scheduler.set_timesteps(num_inference_steps, device=self.device, shift=shift) + init_timesteps = self.scheduler.timesteps + fps_embeds = [fps] * prompt_embeds[0].shape[0] + fps_embeds = [0 if i == 16 else 1 for i in fps_embeds] + transformer_dtype = self.dtype + # with torch.cuda.amp.autocast(dtype=self.dtype), torch.no_grad(): + if overlap_history is None or base_num_frames is None or num_frames <= base_num_frames: + # short video generation + latent_shape = [16, latent_length, latent_height, latent_width] + latents = self.prepare_latents( + latent_shape, dtype=torch.float32, device=self.device, generator=generator + ) + latents = [latents] + if prefix_video is not None: + latents[0][:, :predix_video_latent_length] = prefix_video[0].to(torch.float32) + base_num_frames = (base_num_frames - 1) // 4 + 1 if base_num_frames is not None else latent_length + step_matrix, _, step_update_mask, valid_interval = self.generate_timestep_matrix( + latent_length, init_timesteps, base_num_frames, ar_step, predix_video_latent_length, causal_block_size + ) + sample_schedulers = [] + for _ in range(latent_length): + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=1000, shift=1, use_dynamic_shifting=False + ) + sample_scheduler.set_timesteps(num_inference_steps, device=self.device, shift=shift) + sample_schedulers.append(sample_scheduler) + sample_schedulers_counter = [0] * latent_length + + if callback != None: + callback(-1, None, True) + + freqs = get_rotary_pos_embed(latents[0].shape[1:], enable_RIFLEx= False) + for i, timestep_i in enumerate(tqdm(step_matrix)): + update_mask_i = step_update_mask[i] + valid_interval_i = valid_interval[i] + valid_interval_start, valid_interval_end = valid_interval_i + timestep = timestep_i[None, valid_interval_start:valid_interval_end].clone() + latent_model_input = [latents[0][:, valid_interval_start:valid_interval_end, :, :].clone()] + if addnoise_condition > 0 and valid_interval_start < predix_video_latent_length: + noise_factor = 0.001 * addnoise_condition + timestep_for_noised_condition = addnoise_condition + latent_model_input[0][:, valid_interval_start:predix_video_latent_length] = ( + latent_model_input[0][:, valid_interval_start:predix_video_latent_length] * (1.0 - noise_factor) + + torch.randn_like(latent_model_input[0][:, valid_interval_start:predix_video_latent_length]) + * noise_factor + ) + timestep[:, valid_interval_start:predix_video_latent_length] = timestep_for_noised_condition + kwrags = { + "x" : torch.stack([latent_model_input[0]]), + "t" : timestep, + "freqs" :freqs, + "fps" : fps_embeds, + # "causal_block_size" : causal_block_size, + "callback" : callback, + "pipeline" : self + } + kwrags.update(i2v_extra_kwrags) + + + if not self.do_classifier_free_guidance: + noise_pred = self.model( + context=prompt_embeds, + **kwrags, + )[0] + if self._interrupt: + return None + noise_pred= noise_pred.to(torch.float32) + else: + if joint_pass: + noise_pred_cond, noise_pred_uncond = self.model( + context=prompt_embeds, + context2=negative_prompt_embeds, + **kwrags, + ) + if self._interrupt: + return None + else: + noise_pred_cond = self.model( + context=prompt_embeds, + **kwrags, + )[0] + if self._interrupt: + return None + noise_pred_uncond = self.model( + context=negative_prompt_embeds, + **kwrags, + )[0] + if self._interrupt: + return None + noise_pred_cond= noise_pred_cond.to(torch.float32) + noise_pred_uncond= noise_pred_uncond.to(torch.float32) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + del noise_pred_cond, noise_pred_uncond + for idx in range(valid_interval_start, valid_interval_end): + if update_mask_i[idx].item(): + latents[0][:, idx] = sample_schedulers[idx].step( + noise_pred[:, idx - valid_interval_start], + timestep_i[idx], + latents[0][:, idx], + return_dict=False, + generator=generator, + )[0] + sample_schedulers_counter[idx] += 1 + if callback is not None: + callback(i, latents[0], False) + + x0 = latents[0].unsqueeze(0) + videos = self.vae.decode(x0, tile_size= VAE_tile_size) + videos = (videos / 2 + 0.5).clamp(0, 1) + videos = [video for video in videos] + videos = [video.permute(1, 2, 3, 0) * 255 for video in videos] + videos = [video.cpu().numpy().astype(np.uint8) for video in videos] + return videos + else: + # long video generation + base_num_frames = (base_num_frames - 1) // 4 + 1 if base_num_frames is not None else latent_length + overlap_history_frames = (overlap_history - 1) // 4 + 1 + n_iter = 1 + (latent_length - base_num_frames - 1) // (base_num_frames - overlap_history_frames) + 1 + print(f"n_iter:{n_iter}") + output_video = None + for i in range(n_iter): + if output_video is not None: # i !=0 + prefix_video = output_video[:, -overlap_history:].to(self.device) + prefix_video = [self.vae.encode(prefix_video.unsqueeze(0))[0]] # [(c, f, h, w)] + if prefix_video[0].shape[1] % causal_block_size != 0: + truncate_len = prefix_video[0].shape[1] % causal_block_size + print("the length of prefix video is truncated for the casual block size alignment.") + prefix_video[0] = prefix_video[0][:, : prefix_video[0].shape[1] - truncate_len] + predix_video_latent_length = prefix_video[0].shape[1] + finished_frame_num = i * (base_num_frames - overlap_history_frames) + overlap_history_frames + left_frame_num = latent_length - finished_frame_num + base_num_frames_iter = min(left_frame_num + overlap_history_frames, base_num_frames) + else: # i == 0 + base_num_frames_iter = base_num_frames + latent_shape = [16, base_num_frames_iter, latent_height, latent_width] + latents = self.prepare_latents( + latent_shape, dtype=torch.float32, device=self.device, generator=generator + ) + latents = [latents] + if prefix_video is not None: + latents[0][:, :predix_video_latent_length] = prefix_video[0].to(torch.float32) + step_matrix, _, step_update_mask, valid_interval = self.generate_timestep_matrix( + base_num_frames_iter, + init_timesteps, + base_num_frames_iter, + ar_step, + predix_video_latent_length, + causal_block_size, + ) + sample_schedulers = [] + for _ in range(base_num_frames_iter): + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=1000, shift=1, use_dynamic_shifting=False + ) + sample_scheduler.set_timesteps(num_inference_steps, device=self.device, shift=shift) + sample_schedulers.append(sample_scheduler) + sample_schedulers_counter = [0] * base_num_frames_iter + if callback != None: + callback(-1, None, True) + + freqs = get_rotary_pos_embed(latents[0].shape[1:], enable_RIFLEx= False) + for i, timestep_i in enumerate(tqdm(step_matrix)): + update_mask_i = step_update_mask[i] + valid_interval_i = valid_interval[i] + valid_interval_start, valid_interval_end = valid_interval_i + timestep = timestep_i[None, valid_interval_start:valid_interval_end].clone() + latent_model_input = [latents[0][:, valid_interval_start:valid_interval_end, :, :].clone()] + if addnoise_condition > 0 and valid_interval_start < predix_video_latent_length: + noise_factor = 0.001 * addnoise_condition + timestep_for_noised_condition = addnoise_condition + latent_model_input[0][:, valid_interval_start:predix_video_latent_length] = ( + latent_model_input[0][:, valid_interval_start:predix_video_latent_length] + * (1.0 - noise_factor) + + torch.randn_like( + latent_model_input[0][:, valid_interval_start:predix_video_latent_length] + ) + * noise_factor + ) + timestep[:, valid_interval_start:predix_video_latent_length] = timestep_for_noised_condition + kwrags = { + "x" : torch.stack([latent_model_input[0]]), + "t" : timestep, + "freqs" :freqs, + "fps" : fps_embeds, + "causal_block_size" : causal_block_size, + "causal_attention" : causal_attention, + "callback" : callback, + "pipeline" : self + } + kwrags.update(i2v_extra_kwrags) + + if not self.do_classifier_free_guidance: + noise_pred = self.model( + context=prompt_embeds, + **kwrags, + )[0] + if self._interrupt: + return None + noise_pred= noise_pred.to(torch.float32) + else: + if joint_pass: + noise_pred_cond, noise_pred_uncond = self.model( + context=prompt_embeds, + context2=negative_prompt_embeds, + **kwrags, + ) + if self._interrupt: + return None + else: + noise_pred_cond = self.model( + context=prompt_embeds, + **kwrags, + )[0] + if self._interrupt: + return None + noise_pred_uncond = self.model( + context=negative_prompt_embeds, + )[0] + if self._interrupt: + return None + noise_pred_cond= noise_pred_cond.to(torch.float32) + noise_pred_uncond= noise_pred_uncond.to(torch.float32) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + del noise_pred_cond, noise_pred_uncond + for idx in range(valid_interval_start, valid_interval_end): + if update_mask_i[idx].item(): + latents[0][:, idx] = sample_schedulers[idx].step( + noise_pred[:, idx - valid_interval_start], + timestep_i[idx], + latents[0][:, idx], + return_dict=False, + generator=generator, + )[0] + sample_schedulers_counter[idx] += 1 + if callback is not None: + callback(i, latents[0].squeeze(0), False) + + x0 = latents[0].unsqueeze(0) + videos = [self.vae.decode(x0, tile_size= VAE_tile_size)[0]] + if output_video is None: + output_video = videos[0].clamp(-1, 1).cpu() # c, f, h, w + else: + output_video = torch.cat( + [output_video, videos[0][:, overlap_history:].clamp(-1, 1).cpu()], 1 + ) # c, f, h, w + return output_video diff --git a/wan/diffusion_forcing.py b/models/wan/diffusion_forcing.py similarity index 98% rename from wan/diffusion_forcing.py rename to models/wan/diffusion_forcing.py index 6960bda..840330f 100644 --- a/wan/diffusion_forcing.py +++ b/models/wan/diffusion_forcing.py @@ -14,12 +14,12 @@ from tqdm import tqdm from .modules.model import WanModel from .modules.t5 import T5EncoderModel from .modules.vae import WanVAE -from wan.modules.posemb_layers import get_rotary_pos_embed -from wan.utils.utils import calculate_new_dimensions -from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler, +from .modules.posemb_layers import get_rotary_pos_embed +from shared.utils.utils import calculate_new_dimensions +from shared.utils.fm_solvers import (FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps) -from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler -from wan.utils.loras_mutipliers import update_loras_slists +from shared.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler +from shared.utils.loras_mutipliers import update_loras_slists class DTT2V: @@ -431,5 +431,3 @@ class DTT2V: return videos -def query_model_def(model_type, model_def): - return None \ No newline at end of file diff --git a/ltx_video/utils/__init__.py b/models/wan/distributed/__init__.py similarity index 100% rename from ltx_video/utils/__init__.py rename to models/wan/distributed/__init__.py diff --git a/wan/distributed/fsdp.py b/models/wan/distributed/fsdp.py similarity index 100% rename from wan/distributed/fsdp.py rename to models/wan/distributed/fsdp.py diff --git a/wan/distributed/xdit_context_parallel.py b/models/wan/distributed/xdit_context_parallel.py similarity index 100% rename from wan/distributed/xdit_context_parallel.py rename to models/wan/distributed/xdit_context_parallel.py diff --git a/wan/fantasytalking/infer.py b/models/wan/fantasytalking/infer.py similarity index 100% rename from wan/fantasytalking/infer.py rename to models/wan/fantasytalking/infer.py diff --git a/wan/fantasytalking/model.py b/models/wan/fantasytalking/model.py similarity index 99% rename from wan/fantasytalking/model.py rename to models/wan/fantasytalking/model.py index 5ec3655..d0eb74d 100644 --- a/wan/fantasytalking/model.py +++ b/models/wan/fantasytalking/model.py @@ -1,7 +1,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from wan.modules.attention import pay_attention +from shared.attention import pay_attention class AudioProjModel(nn.Module): diff --git a/wan/fantasytalking/utils.py b/models/wan/fantasytalking/utils.py similarity index 100% rename from wan/fantasytalking/utils.py rename to models/wan/fantasytalking/utils.py diff --git a/wan/modules/__init__.py b/models/wan/modules/__init__.py similarity index 81% rename from wan/modules/__init__.py rename to models/wan/modules/__init__.py index 38c29ce..56aea65 100644 --- a/wan/modules/__init__.py +++ b/models/wan/modules/__init__.py @@ -1,8 +1,9 @@ -from .attention import pay_attention +from shared.attention import pay_attention from .model import WanModel from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model from .tokenizers import HuggingfaceTokenizer from .vae import WanVAE +from .vae2_2 import Wan2_2_VAE __all__ = [ 'WanVAE', diff --git a/wan/modules/clip.py b/models/wan/modules/clip.py similarity index 99% rename from wan/modules/clip.py rename to models/wan/modules/clip.py index da91a00..fc29893 100644 --- a/wan/modules/clip.py +++ b/models/wan/modules/clip.py @@ -8,7 +8,7 @@ import torch.nn as nn import torch.nn.functional as F import torchvision.transforms as T -from .attention import pay_attention +from shared.attention import pay_attention from .tokenizers import HuggingfaceTokenizer from .xlm_roberta import XLMRoberta diff --git a/wan/modules/model.py b/models/wan/modules/model.py similarity index 99% rename from wan/modules/model.py rename to models/wan/modules/model.py index 19eb5c3..6dd0199 100644 --- a/wan/modules/model.py +++ b/models/wan/modules/model.py @@ -12,9 +12,9 @@ from diffusers.models.modeling_utils import ModelMixin import numpy as np from typing import Union,Optional from mmgp import offload -from .attention import pay_attention +from shared.attention import pay_attention from torch.backends.cuda import sdp_kernel -from wan.multitalk.multitalk_utils import get_attn_map_with_target +from ..multitalk.multitalk_utils import get_attn_map_with_target __all__ = ['WanModel'] @@ -150,7 +150,7 @@ class WanLayerNorm(nn.LayerNorm): return x # return super().forward(x).type_as(x) -from wan.modules.posemb_layers import apply_rotary_emb +from .posemb_layers import apply_rotary_emb class WanSelfAttention(nn.Module): @@ -429,7 +429,7 @@ class WanAttentionBlock(nn.Module): self.block_id = block_id if output_dim > 0: - from wan.multitalk.attention import SingleStreamMutiAttention + from ..multitalk.attention import SingleStreamMutiAttention # init audio module self.audio_cross_attn = SingleStreamMutiAttention( dim=dim, @@ -762,15 +762,17 @@ class WanModel(ModelMixin, ConfigMixin): offload.shared_state["_chipmunk_layers"] = None def preprocess_loras(self, model_type, sd): - # new_sd = {} - # for k,v in sd.items(): - # if not k.endswith(".modulation.diff"): - # new_sd[ k] = v - # sd = new_sd + first = next(iter(sd), None) if first == None: return sd - + + # if first.startswith("blocks."): + # new_sd = {} + # for k,v in sd.items(): + # new_sd["diffusion_model." + k] = v + # sd = new_sd + if first.startswith("lora_unet_"): new_sd = {} print("Converting Lora Safetensors format to Lora Diffusers format") @@ -846,7 +848,6 @@ class WanModel(ModelMixin, ConfigMixin): super().__init__() - assert model_type in ['t2v', 'i2v', 'i2v2_2'] self.model_type = model_type self.patch_size = patch_size @@ -893,7 +894,7 @@ class WanModel(ModelMixin, ConfigMixin): # blocks if vace_layers == None: - cross_attn_type = 't2v_cross_attn' if model_type in ['t2v','i2v2_2'] else 'i2v_cross_attn' + cross_attn_type = 't2v_cross_attn' if model_type in ['t2v','i2v2_2', 'ti2v2_2'] else 'i2v_cross_attn' self.blocks = nn.ModuleList([ WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, block_no =i, output_dim=multitalk_output_dim, norm_input_visual=norm_input_visual) @@ -962,7 +963,7 @@ class WanModel(ModelMixin, ConfigMixin): block.projector.bias = nn.Parameter(torch.zeros(dim)) if fantasytalking_dim > 0: - from wan.fantasytalking.model import WanCrossAttentionProcessor + from ..fantasytalking.model import WanCrossAttentionProcessor for block in self.blocks: block.cross_attn.processor = WanCrossAttentionProcessor(fantasytalking_dim, dim) diff --git a/wan/modules/motion_patch.py b/models/wan/modules/motion_patch.py similarity index 100% rename from wan/modules/motion_patch.py rename to models/wan/modules/motion_patch.py diff --git a/wan/modules/posemb_layers.py b/models/wan/modules/posemb_layers.py similarity index 100% rename from wan/modules/posemb_layers.py rename to models/wan/modules/posemb_layers.py diff --git a/wan/modules/t5.py b/models/wan/modules/t5.py similarity index 100% rename from wan/modules/t5.py rename to models/wan/modules/t5.py diff --git a/wan/modules/tokenizers.py b/models/wan/modules/tokenizers.py similarity index 100% rename from wan/modules/tokenizers.py rename to models/wan/modules/tokenizers.py diff --git a/wan/modules/vae.py b/models/wan/modules/vae.py similarity index 100% rename from wan/modules/vae.py rename to models/wan/modules/vae.py diff --git a/models/wan/modules/vae2_2.py b/models/wan/modules/vae2_2.py new file mode 100644 index 0000000..1592e13 --- /dev/null +++ b/models/wan/modules/vae2_2.py @@ -0,0 +1,1186 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import logging + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +__all__ = [ + "Wan2_2_VAE", +] + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = ( + self.padding[2], + self.padding[2], + self.padding[1], + self.padding[1], + 2 * self.padding[0], + 0, + ) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return (F.normalize(x, dim=(1 if self.channel_first else -1)) * + self.scale * self.gamma + self.bias) + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ( + "none", + "upsample2d", + "upsample3d", + "downsample2d", + "downsample3d", + ) + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + # nn.Conv2d(dim, dim//2, 3, padding=1) + ) + self.time_conv = CausalConv3d( + dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and + feat_cache[idx] != "Rep"): + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and + feat_cache[idx] == "Rep"): + cache_x = torch.cat( + [ + torch.zeros_like(cache_x).to(cache_x.device), + cache_x + ], + dim=2, + ) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight.detach().clone() + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix # * 0.5 + conv.weight = nn.Parameter(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data.detach().clone() + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight = nn.Parameter(conv_weight) + nn.init.zeros_(conv.bias.data) + + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), + nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1), + ) + self.shortcut = ( + CausalConv3d(in_dim, out_dim, 1) + if in_dim != out_dim else nn.Identity()) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm(x) + # compute query, key, value + q, k, v = ( + self.to_qkv(x).reshape(b * t, 1, c * 3, + -1).permute(0, 1, 3, + 2).contiguous().chunk(3, dim=-1)) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, "(b t) c h w-> b c t h w", t=t) + return x + identity + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange( + x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange( + x, + "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, + r=patch_size, + ) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() == 4: + x = rearrange( + x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange( + x, + "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, + r=patch_size, + ) + return x + + +class AvgDown3D(nn.Module): + + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1:, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + + def __init__(self, + in_dim, + out_dim, + dropout, + mult, + temperal_downsample=False, + down_flag=False): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample(out_dim, mode=mode)) + + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + + def __init__(self, + in_dim, + out_dim, + dropout, + mult, + temperal_upsample=False, + up_flag=False): + super().__init__() + # Shortcut path with upsample + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + + # Main path with residual blocks and upsample + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final upsample block + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample(out_dim, mode=mode)) + + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut + else: + return x_main + + +class Encoder3d(nn.Module): + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = ( + temperal_downsample[i] + if i < len(temperal_downsample) else False) + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + )) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + # # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + return x + + +class Decoder3d(nn.Module): + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout), + ) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len( + temperal_upsample) else False + upsamples.append( + Up_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1, + )) + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class WanVAE_(nn.Module): + + def __init__( + self, + dim=160, + dec_dim=256, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d( + dim, + z_dim * 2, + dim_mult, + num_res_blocks, + attn_scales, + self.temperal_downsample, + dropout, + ) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d( + dec_dim, + z_dim, + dim_mult, + num_res_blocks, + attn_scales, + self.temperal_upsample, + dropout, + ) + + def forward(self, x, scale=[0, 1]): + mu = self.encode(x, scale) + x_recon = self.decode(mu, scale) + return x_recon, mu + + def encode(self, x, scale, any_end_frame = False): + self.clear_cache() + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + def decode(self, z, scale,any_end_frame = False): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + else: + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + ) + out = torch.cat([out, out_], 2) + out = unpatchify(out, patch_size=2) + self.clear_cache() + return out + + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent) + return b + + def spatial_tiled_decode(self, z, scale, tile_size, any_end_frame= False): + tile_sample_min_size = tile_size + tile_latent_min_size = int(tile_sample_min_size / 16) + tile_overlap_factor = 0.25 + + # z: [b,c,t,h,w] + + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + + + overlap_size = int(tile_latent_min_size * (1 - tile_overlap_factor)) #8 0.75 + blend_extent = int(tile_sample_min_size * tile_overlap_factor) #256 0.25 + row_limit = tile_sample_min_size - blend_extent + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, z.shape[-2], overlap_size): + row = [] + for j in range(0, z.shape[-1], overlap_size): + tile = z[:, :, :, i: i + tile_latent_min_size, j: j + tile_latent_min_size] + decoded = self.decode(tile, scale, any_end_frame= any_end_frame) + row.append(decoded) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + return torch.cat(result_rows, dim=-2) + + + def spatial_tiled_encode(self, x, scale, tile_size, any_end_frame = False) : + tile_sample_min_size = tile_size + tile_latent_min_size = int(tile_sample_min_size / 16) + tile_overlap_factor = 0.25 + + overlap_size = int(tile_sample_min_size * (1 - tile_overlap_factor)) + blend_extent = int(tile_latent_min_size * tile_overlap_factor) + row_limit = tile_latent_min_size - blend_extent + + # Split video into tiles and encode them separately. + rows = [] + for i in range(0, x.shape[-2], overlap_size): + row = [] + for j in range(0, x.shape[-1], overlap_size): + tile = x[:, :, :, i: i + tile_sample_min_size, j: j + tile_sample_min_size] + tile = self.encode(tile, scale, any_end_frame= any_end_frame) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + mu = torch.cat(result_rows, dim=-2) + + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + + return mu + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs): + # params + cfg = dict( + dim=dim, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, True], + dropout=0.0, + ) + cfg.update(**kwargs) + + # init model + with torch.device("meta"): + model = WanVAE_(**cfg) + + from mmgp import offload + # load checkpoint + logging.info(f"loading {pretrained_path}") + # model.load_state_dict( + # torch.load(pretrained_path, map_location=device), assign=True) + # offload.save_model(model, "Wan_vae_2_2.safetensors") + # model.to(torch.bfloat16) + # offload.save_model(model, "Wan_vae_2_2_bf16.safetensors") + offload.load_model_data(model, pretrained_path.replace(".pth", ".safetensors"), writable_tensors= False) + + return model + + +class Wan2_2_VAE: + + def __init__( + self, + z_dim=48, + c_dim=160, + vae_pth=None, + dim_mult=[1, 2, 4, 4], + temperal_downsample=[False, True, True], + dtype=torch.float, + device="cuda", + ): + + self.dtype = dtype + self.device = device + + mean = torch.tensor( + [ + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ], + dtype=dtype, + device=device, + ) + std = torch.tensor( + [ + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ], + dtype=dtype, + device=device, + ) + self.scale = [mean, 1.0 / std] + + # init model + self.model = ( + _video_vae( + pretrained_path=vae_pth, + z_dim=z_dim, + dim=c_dim, + dim_mult=dim_mult, + temperal_downsample=temperal_downsample, + ).eval().requires_grad_(False).to(device)) + + + @staticmethod + def get_VAE_tile_size(vae_config, device_mem_capacity, mixed_precision): + # VAE Tiling + if vae_config == 0: + if mixed_precision: + device_mem_capacity = device_mem_capacity / 2 + if device_mem_capacity >= 24000: + use_vae_config = 1 + elif device_mem_capacity >= 8000: + use_vae_config = 2 + else: + use_vae_config = 3 + else: + use_vae_config = vae_config + + if use_vae_config == 1: + VAE_tile_size = 0 + elif use_vae_config == 2: + VAE_tile_size = 256 + else: + VAE_tile_size = 128 + + return VAE_tile_size + + def encode(self, videos, tile_size = 256, any_end_frame = False): + """ + videos: A list of videos each with shape [C, T, H, W]. + """ + original_dtype = videos[0].dtype + + if tile_size > 0 and False: + return [ self.model.spatial_tiled_encode(u.to(self.dtype).unsqueeze(0), self.scale, tile_size, any_end_frame=any_end_frame).float().squeeze(0) for u in videos ] + else: + return [ self.model.encode(u.to(self.dtype).unsqueeze(0), self.scale, any_end_frame=any_end_frame).float().squeeze(0) for u in videos ] + + + def decode(self, zs, tile_size, any_end_frame = False): + if tile_size > 0 and False: + return [ self.model.spatial_tiled_decode(u.to(self.dtype).unsqueeze(0), self.scale, tile_size, any_end_frame=any_end_frame).clamp_(-1, 1).float().squeeze(0) for u in zs ] + else: + return [ self.model.decode(u.to(self.dtype).unsqueeze(0), self.scale, any_end_frame=any_end_frame).clamp_(-1, 1).float().squeeze(0) for u in zs ] + + + # def encode(self, videos, VAE_tile_size = 0, any_end_frame = False ): + # with amp.autocast(dtype=self.dtype): + # return [ + # self.model.encode(u.unsqueeze(0), + # self.scale).float().squeeze(0) + # for u in videos + # ] + + # def decode(self, zs, VAE_tile_size = 0, any_end_frame = False): + # with amp.autocast(dtype=self.dtype): + # return [ + # self.model.decode(u.unsqueeze(0), + # self.scale).float().clamp_(-1, + # 1).squeeze(0) + # for u in zs + # ] diff --git a/wan/modules/xlm_roberta.py b/models/wan/modules/xlm_roberta.py similarity index 100% rename from wan/modules/xlm_roberta.py rename to models/wan/modules/xlm_roberta.py diff --git a/wan/multitalk/attention.py b/models/wan/multitalk/attention.py similarity index 99% rename from wan/multitalk/attention.py rename to models/wan/multitalk/attention.py index 12fb317..27d488f 100644 --- a/wan/multitalk/attention.py +++ b/models/wan/multitalk/attention.py @@ -3,7 +3,7 @@ import torch import torch.nn as nn from einops import rearrange, repeat from .multitalk_utils import RotaryPositionalEmbedding1D, normalize_and_scale, split_token_counts_and_frame_ids -from wan.modules.attention import pay_attention +from shared.attention import pay_attention # import xformers.ops diff --git a/wan/multitalk/kokoro/__init__.py b/models/wan/multitalk/kokoro/__init__.py similarity index 100% rename from wan/multitalk/kokoro/__init__.py rename to models/wan/multitalk/kokoro/__init__.py diff --git a/wan/multitalk/kokoro/__main__.py b/models/wan/multitalk/kokoro/__main__.py similarity index 100% rename from wan/multitalk/kokoro/__main__.py rename to models/wan/multitalk/kokoro/__main__.py diff --git a/wan/multitalk/kokoro/custom_stft.py b/models/wan/multitalk/kokoro/custom_stft.py similarity index 100% rename from wan/multitalk/kokoro/custom_stft.py rename to models/wan/multitalk/kokoro/custom_stft.py diff --git a/wan/multitalk/kokoro/istftnet.py b/models/wan/multitalk/kokoro/istftnet.py similarity index 100% rename from wan/multitalk/kokoro/istftnet.py rename to models/wan/multitalk/kokoro/istftnet.py diff --git a/wan/multitalk/kokoro/model.py b/models/wan/multitalk/kokoro/model.py similarity index 100% rename from wan/multitalk/kokoro/model.py rename to models/wan/multitalk/kokoro/model.py diff --git a/wan/multitalk/kokoro/modules.py b/models/wan/multitalk/kokoro/modules.py similarity index 100% rename from wan/multitalk/kokoro/modules.py rename to models/wan/multitalk/kokoro/modules.py diff --git a/wan/multitalk/kokoro/pipeline.py b/models/wan/multitalk/kokoro/pipeline.py similarity index 100% rename from wan/multitalk/kokoro/pipeline.py rename to models/wan/multitalk/kokoro/pipeline.py diff --git a/wan/multitalk/multitalk.py b/models/wan/multitalk/multitalk.py similarity index 98% rename from wan/multitalk/multitalk.py rename to models/wan/multitalk/multitalk.py index 56ba16b..cadb78b 100644 --- a/wan/multitalk/multitalk.py +++ b/models/wan/multitalk/multitalk.py @@ -7,10 +7,8 @@ import subprocess import torchvision.transforms as transforms import torch.nn.functional as F import torch.nn as nn -import wan -from wan.configs import SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS -from wan.utils.utils import cache_image, cache_video, str2bool -# from wan.utils.multitalk_utils import save_video_ffmpeg +from shared.utils.utils import cache_image, cache_video, str2bool +# from shared.utils.multitalk_utils import save_video_ffmpeg # from .kokoro import KPipeline from transformers import Wav2Vec2FeatureExtractor from .wav2vec2 import Wav2Vec2Model diff --git a/wan/multitalk/multitalk_model.py b/models/wan/multitalk/multitalk_model.py similarity index 100% rename from wan/multitalk/multitalk_model.py rename to models/wan/multitalk/multitalk_model.py diff --git a/wan/multitalk/multitalk_utils.py b/models/wan/multitalk/multitalk_utils.py similarity index 100% rename from wan/multitalk/multitalk_utils.py rename to models/wan/multitalk/multitalk_utils.py diff --git a/wan/multitalk/torch_utils.py b/models/wan/multitalk/torch_utils.py similarity index 100% rename from wan/multitalk/torch_utils.py rename to models/wan/multitalk/torch_utils.py diff --git a/wan/multitalk/wav2vec2.py b/models/wan/multitalk/wav2vec2.py similarity index 100% rename from wan/multitalk/wav2vec2.py rename to models/wan/multitalk/wav2vec2.py diff --git a/models/wan/text2video fuse attempt.py b/models/wan/text2video fuse attempt.py new file mode 100644 index 0000000..8af9458 --- /dev/null +++ b/models/wan/text2video fuse attempt.py @@ -0,0 +1,698 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import gc +import logging +import math +import os +import random +import sys +import types +from contextlib import contextmanager +from functools import partial +from mmgp import offload +import torch +import torch.nn as nn +import torch.cuda.amp as amp +import torch.distributed as dist +from tqdm import tqdm +from PIL import Image +import torchvision.transforms.functional as TF +import torch.nn.functional as F +from .distributed.fsdp import shard_model +from .modules.model import WanModel +from .modules.t5 import T5EncoderModel +from .modules.vae import WanVAE +from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler, + get_sampling_sigmas, retrieve_timesteps) +from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler +from wan.modules.posemb_layers import get_rotary_pos_embed +from .utils.vace_preprocessor import VaceVideoProcessor + + +def optimized_scale(positive_flat, negative_flat): + + # Calculate dot production + dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) + + # Squared norm of uncondition + squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8 + + # st_star = v_cond^T * v_uncond / ||v_uncond||^2 + st_star = dot_product / squared_norm + + return st_star + + +class WanT2V: + + def __init__( + self, + config, + checkpoint_dir, + rank=0, + model_filename = None, + text_encoder_filename = None, + quantizeTransformer = False, + dtype = torch.bfloat16 + ): + self.device = torch.device(f"cuda") + self.config = config + self.rank = rank + self.dtype = dtype + self.num_train_timesteps = config.num_train_timesteps + self.param_dtype = config.param_dtype + + self.text_encoder = T5EncoderModel( + text_len=config.text_len, + dtype=config.t5_dtype, + device=torch.device('cpu'), + checkpoint_path=text_encoder_filename, + tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), + shard_fn= None) + + self.vae_stride = config.vae_stride + self.patch_size = config.patch_size + + + self.vae = WanVAE( + vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), + device=self.device) + + logging.info(f"Creating WanModel from {model_filename}") + from mmgp import offload + + self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel,do_quantize= quantizeTransformer, writable_tensors= False) + # offload.load_model_data(self.model, "recam.ckpt") + # self.model.cpu() + # offload.save_model(self.model, "recam.safetensors") + if self.dtype == torch.float16 and not "fp16" in model_filename: + self.model.to(self.dtype) + # offload.save_model(self.model, "t2v_fp16.safetensors",do_quantize=True) + if self.dtype == torch.float16: + self.vae.model.to(self.dtype) + self.model.eval().requires_grad_(False) + + + self.sample_neg_prompt = config.sample_neg_prompt + + if "Vace" in model_filename: + self.vid_proc = VaceVideoProcessor(downsample=tuple([x * y for x, y in zip(config.vae_stride, self.patch_size)]), + min_area=480*832, + max_area=480*832, + min_fps=config.sample_fps, + max_fps=config.sample_fps, + zero_start=True, + seq_len=32760, + keep_last=True) + + self.adapt_vace_model() + + self.scheduler = FlowUniPCMultistepScheduler() + + def vace_encode_frames(self, frames, ref_images, masks=None, tile_size = 0): + if ref_images is None: + ref_images = [None] * len(frames) + else: + assert len(frames) == len(ref_images) + + if masks is None: + latents = self.vae.encode(frames, tile_size = tile_size) + else: + inactive = [i * (1 - m) + 0 * m for i, m in zip(frames, masks)] + reactive = [i * m + 0 * (1 - m) for i, m in zip(frames, masks)] + inactive = self.vae.encode(inactive, tile_size = tile_size) + reactive = self.vae.encode(reactive, tile_size = tile_size) + latents = [torch.cat((u, c), dim=0) for u, c in zip(inactive, reactive)] + + cat_latents = [] + for latent, refs in zip(latents, ref_images): + if refs is not None: + if masks is None: + ref_latent = self.vae.encode(refs, tile_size = tile_size) + else: + ref_latent = self.vae.encode(refs, tile_size = tile_size) + ref_latent = [torch.cat((u, torch.zeros_like(u)), dim=0) for u in ref_latent] + assert all([x.shape[1] == 1 for x in ref_latent]) + latent = torch.cat([*ref_latent, latent], dim=1) + cat_latents.append(latent) + return cat_latents + + def vace_encode_masks(self, masks, ref_images=None): + if ref_images is None: + ref_images = [None] * len(masks) + else: + assert len(masks) == len(ref_images) + + result_masks = [] + for mask, refs in zip(masks, ref_images): + c, depth, height, width = mask.shape + new_depth = int((depth + 3) // self.vae_stride[0]) + height = 2 * (int(height) // (self.vae_stride[1] * 2)) + width = 2 * (int(width) // (self.vae_stride[2] * 2)) + + # reshape + mask = mask[0, :, :, :] + mask = mask.view( + depth, height, self.vae_stride[1], width, self.vae_stride[1] + ) # depth, height, 8, width, 8 + mask = mask.permute(2, 4, 0, 1, 3) # 8, 8, depth, height, width + mask = mask.reshape( + self.vae_stride[1] * self.vae_stride[2], depth, height, width + ) # 8*8, depth, height, width + + # interpolation + mask = F.interpolate(mask.unsqueeze(0), size=(new_depth, height, width), mode='nearest-exact').squeeze(0) + + if refs is not None: + length = len(refs) + mask_pad = torch.zeros_like(mask[:, :length, :, :]) + mask = torch.cat((mask_pad, mask), dim=1) + result_masks.append(mask) + return result_masks + + def vace_latent(self, z, m): + return [torch.cat([zz, mm], dim=0) for zz, mm in zip(z, m)] + + def prepare_source(self, src_video, src_mask, src_ref_images, total_frames, image_size, device, original_video = False, keep_frames= [], start_frame = 0, pre_src_video = None): + image_sizes = [] + trim_video = len(keep_frames) + + for i, (sub_src_video, sub_src_mask, sub_pre_src_video) in enumerate(zip(src_video, src_mask,pre_src_video)): + prepend_count = 0 if sub_pre_src_video == None else sub_pre_src_video.shape[1] + num_frames = total_frames - prepend_count + if sub_src_mask is not None and sub_src_video is not None: + src_video[i], src_mask[i], _, _, _ = self.vid_proc.load_video_pair(sub_src_video, sub_src_mask, max_frames= num_frames, trim_video = trim_video - prepend_count, start_frame = start_frame) + # src_video is [-1, 1], 0 = inpainting area (in fact 127 in [0, 255]) + # src_mask is [-1, 1], 0 = preserve original video (in fact 127 in [0, 255]) and 1 = Inpainting (in fact 255 in [0, 255]) + src_video[i] = src_video[i].to(device) + src_mask[i] = src_mask[i].to(device) + if prepend_count > 0: + src_video[i] = torch.cat( [sub_pre_src_video, src_video[i]], dim=1) + src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), src_mask[i]] ,1) + src_video_shape = src_video[i].shape + if src_video_shape[1] != total_frames: + src_video[i] = torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) + src_mask[i] = torch.cat( [src_mask[i], src_mask[i].new_ones(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) + src_mask[i] = torch.clamp((src_mask[i][:1, :, :, :] + 1) / 2, min=0, max=1) + image_sizes.append(src_video[i].shape[2:]) + elif sub_src_video is None: + if prepend_count > 0: + src_video[i] = torch.cat( [sub_pre_src_video, torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device)], dim=1) + src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), torch.ones((3, num_frames, image_size[0], image_size[1]), device=device)] ,1) + else: + src_video[i] = torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device) + src_mask[i] = torch.ones_like(src_video[i], device=device) + image_sizes.append(image_size) + else: + src_video[i], _, _, _ = self.vid_proc.load_video(sub_src_video, max_frames= num_frames, trim_video = trim_video - prepend_count, start_frame = start_frame) + src_video[i] = src_video[i].to(device) + src_mask[i] = torch.zeros_like(src_video[i], device=device) if original_video else torch.ones_like(src_video[i], device=device) + if prepend_count > 0: + src_video[i] = torch.cat( [sub_pre_src_video, src_video[i]], dim=1) + src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), src_mask[i]] ,1) + src_video_shape = src_video[i].shape + if src_video_shape[1] != total_frames: + src_video[i] = torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) + src_mask[i] = torch.cat( [src_mask[i], src_mask[i].new_ones(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) + image_sizes.append(src_video[i].shape[2:]) + for k, keep in enumerate(keep_frames): + if not keep: + src_video[i][:, k:k+1] = 0 + src_mask[i][:, k:k+1] = 1 + + for i, ref_images in enumerate(src_ref_images): + if ref_images is not None: + image_size = image_sizes[i] + for j, ref_img in enumerate(ref_images): + if ref_img is not None: + ref_img = TF.to_tensor(ref_img).sub_(0.5).div_(0.5).unsqueeze(1) + if ref_img.shape[-2:] != image_size: + canvas_height, canvas_width = image_size + ref_height, ref_width = ref_img.shape[-2:] + white_canvas = torch.ones((3, 1, canvas_height, canvas_width), device=device) # [-1, 1] + scale = min(canvas_height / ref_height, canvas_width / ref_width) + new_height = int(ref_height * scale) + new_width = int(ref_width * scale) + resized_image = F.interpolate(ref_img.squeeze(1).unsqueeze(0), size=(new_height, new_width), mode='bilinear', align_corners=False).squeeze(0).unsqueeze(1) + top = (canvas_height - new_height) // 2 + left = (canvas_width - new_width) // 2 + white_canvas[:, :, top:top + new_height, left:left + new_width] = resized_image + ref_img = white_canvas + src_ref_images[i][j] = ref_img.to(device) + return src_video, src_mask, src_ref_images + + def decode_latent(self, zs, ref_images=None, tile_size= 0 ): + if ref_images is None: + ref_images = [None] * len(zs) + else: + assert len(zs) == len(ref_images) + + trimed_zs = [] + for z, refs in zip(zs, ref_images): + if refs is not None: + z = z[:, len(refs):, :, :] + trimed_zs.append(z) + + return self.vae.decode(trimed_zs, tile_size= tile_size) + + def generate_timestep_matrix( + self, + num_frames, + step_template, + base_num_frames, + ar_step=5, + num_pre_ready=0, + casual_block_size=1, + shrink_interval_with_mask=False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[tuple]]: + step_matrix, step_index = [], [] + update_mask, valid_interval = [], [] + num_iterations = len(step_template) + 1 + num_frames_block = num_frames // casual_block_size + base_num_frames_block = base_num_frames // casual_block_size + if base_num_frames_block < num_frames_block: + infer_step_num = len(step_template) + gen_block = base_num_frames_block + min_ar_step = infer_step_num / gen_block + assert ar_step >= min_ar_step, f"ar_step should be at least {math.ceil(min_ar_step)} in your setting" + # print(num_frames, step_template, base_num_frames, ar_step, num_pre_ready, casual_block_size, num_frames_block, base_num_frames_block) + step_template = torch.cat( + [ + torch.tensor([999], dtype=torch.int64, device=step_template.device), + step_template.long(), + torch.tensor([0], dtype=torch.int64, device=step_template.device), + ] + ) # to handle the counter in row works starting from 1 + pre_row = torch.zeros(num_frames_block, dtype=torch.long) + if num_pre_ready > 0: + pre_row[: num_pre_ready // casual_block_size] = num_iterations + + while torch.all(pre_row >= (num_iterations - 1)) == False: + new_row = torch.zeros(num_frames_block, dtype=torch.long) + for i in range(num_frames_block): + if i == 0 or pre_row[i - 1] >= ( + num_iterations - 1 + ): # the first frame or the last frame is completely denoised + new_row[i] = pre_row[i] + 1 + else: + new_row[i] = new_row[i - 1] - ar_step + new_row = new_row.clamp(0, num_iterations) + + update_mask.append( + (new_row != pre_row) & (new_row != num_iterations) + ) # False: no need to update， True: need to update + step_index.append(new_row) + step_matrix.append(step_template[new_row]) + pre_row = new_row + + # for long video we split into several sequences, base_num_frames is set to the model max length (for training) + terminal_flag = base_num_frames_block + if shrink_interval_with_mask: + idx_sequence = torch.arange(num_frames_block, dtype=torch.int64) + update_mask = update_mask[0] + update_mask_idx = idx_sequence[update_mask] + last_update_idx = update_mask_idx[-1].item() + terminal_flag = last_update_idx + 1 + # for i in range(0, len(update_mask)): + for curr_mask in update_mask: + if terminal_flag < num_frames_block and curr_mask[terminal_flag]: + terminal_flag += 1 + valid_interval.append((max(terminal_flag - base_num_frames_block, 0), terminal_flag)) + + step_update_mask = torch.stack(update_mask, dim=0) + step_index = torch.stack(step_index, dim=0) + step_matrix = torch.stack(step_matrix, dim=0) + + if casual_block_size > 1: + step_update_mask = step_update_mask.unsqueeze(-1).repeat(1, 1, casual_block_size).flatten(1).contiguous() + step_index = step_index.unsqueeze(-1).repeat(1, 1, casual_block_size).flatten(1).contiguous() + step_matrix = step_matrix.unsqueeze(-1).repeat(1, 1, casual_block_size).flatten(1).contiguous() + valid_interval = [(s * casual_block_size, e * casual_block_size) for s, e in valid_interval] + + return step_matrix, step_index, step_update_mask, valid_interval + + def generate(self, + input_prompt, + input_frames= None, + input_masks = None, + input_ref_images = None, + source_video=None, + target_camera=None, + context_scale=1.0, + size=(1280, 720), + frame_num=81, + shift=5.0, + sample_solver='unipc', + sampling_steps=50, + guide_scale=5.0, + n_prompt="", + seed=-1, + offload_model=True, + callback = None, + enable_RIFLEx = None, + VAE_tile_size = 0, + joint_pass = False, + slg_layers = None, + slg_start = 0.0, + slg_end = 1.0, + cfg_star_switch = True, + cfg_zero_step = 5, + ): + r""" + Generates video frames from text prompt using diffusion process. + + Args: + input_prompt (`str`): + Text prompt for content generation + size (tupele[`int`], *optional*, defaults to (1280,720)): + Controls video resolution, (width,height). + frame_num (`int`, *optional*, defaults to 81): + How many frames to sample from a video. The number should be 4n+1 + shift (`float`, *optional*, defaults to 5.0): + Noise schedule shift parameter. Affects temporal dynamics + sample_solver (`str`, *optional*, defaults to 'unipc'): + Solver used to sample the video. + sampling_steps (`int`, *optional*, defaults to 40): + Number of diffusion sampling steps. Higher values improve quality but slow generation + guide_scale (`float`, *optional*, defaults 5.0): + Classifier-free guidance scale. Controls prompt adherence vs. creativity + n_prompt (`str`, *optional*, defaults to ""): + Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` + seed (`int`, *optional*, defaults to -1): + Random seed for noise generation. If -1, use random seed. + offload_model (`bool`, *optional*, defaults to True): + If True, offloads models to CPU during generation to save VRAM + + Returns: + torch.Tensor: + Generated video frames tensor. Dimensions: (C, N H, W) where: + - C: Color channels (3 for RGB) + - N: Number of frames (81) + - H: Frame height (from size) + - W: Frame width from size) + """ + # preprocess + + if n_prompt == "": + n_prompt = self.sample_neg_prompt + seed = seed if seed >= 0 else random.randint(0, sys.maxsize) + seed_g = torch.Generator(device=self.device) + seed_g.manual_seed(seed) + + frame_num = max(17, frame_num) # must match causal_block_size for value of 5 + frame_num = int( round( (frame_num - 17) / 20)* 20 + 17 ) + num_frames = frame_num + addnoise_condition = 20 + causal_attention = True + fps = 16 + ar_step = 5 + + + + context = self.text_encoder([input_prompt], self.device) + context_null = self.text_encoder([n_prompt], self.device) + if target_camera != None: + size = (source_video.shape[2], source_video.shape[1]) + source_video = source_video.to(dtype=self.dtype , device=self.device) + source_video = source_video.permute(3, 0, 1, 2).div_(127.5).sub_(1.) + source_latents = self.vae.encode([source_video]) #.to(dtype=self.dtype, device=self.device) + del source_video + # Process target camera (recammaster) + from wan.utils.cammmaster_tools import get_camera_embedding + cam_emb = get_camera_embedding(target_camera) + cam_emb = cam_emb.to(dtype=self.dtype, device=self.device) + + if input_frames != None: + # vace context encode + input_frames = [u.to(self.device) for u in input_frames] + input_ref_images = [ None if u == None else [v.to(self.device) for v in u] for u in input_ref_images] + input_masks = [u.to(self.device) for u in input_masks] + + z0 = self.vace_encode_frames(input_frames, input_ref_images, masks=input_masks, tile_size = VAE_tile_size) + m0 = self.vace_encode_masks(input_masks, input_ref_images) + z = self.vace_latent(z0, m0) + + target_shape = list(z0[0].shape) + target_shape[0] = int(target_shape[0] / 2) + else: + F = frame_num + target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1, + size[1] // self.vae_stride[1], + size[0] // self.vae_stride[2]) + + seq_len = math.ceil((target_shape[2] * target_shape[3]) / + (self.patch_size[1] * self.patch_size[2]) * + target_shape[1]) + + context = [u.to(self.dtype) for u in context] + context_null = [u.to(self.dtype) for u in context_null] + + noise = [ torch.randn( *target_shape, dtype=torch.float32, device=self.device, generator=seed_g) ] + + # evaluation mode + + # if sample_solver == 'unipc': + # sample_scheduler = FlowUniPCMultistepScheduler( + # num_train_timesteps=self.num_train_timesteps, + # shift=1, + # use_dynamic_shifting=False) + # sample_scheduler.set_timesteps( + # sampling_steps, device=self.device, shift=shift) + # timesteps = sample_scheduler.timesteps + # elif sample_solver == 'dpm++': + # sample_scheduler = FlowDPMSolverMultistepScheduler( + # num_train_timesteps=self.num_train_timesteps, + # shift=1, + # use_dynamic_shifting=False) + # sampling_sigmas = get_sampling_sigmas(sampling_steps, shift) + # timesteps, _ = retrieve_timesteps( + # sample_scheduler, + # device=self.device, + # sigmas=sampling_sigmas) + # else: + # raise NotImplementedError("Unsupported solver.") + + # sample videos + latents = noise + del noise + batch_size =len(latents) + if target_camera != None: + shape = list(latents[0].shape[1:]) + shape[0] *= 2 + freqs = get_rotary_pos_embed(shape, enable_RIFLEx= False) + else: + freqs = get_rotary_pos_embed(latents[0].shape[1:], enable_RIFLEx= enable_RIFLEx) + # arg_c = {'context': context, 'freqs': freqs, 'pipeline': self, 'callback': callback} + # arg_null = {'context': context_null, 'freqs': freqs, 'pipeline': self, 'callback': callback} + # arg_both = {'context': context, 'context2': context_null, 'freqs': freqs, 'pipeline': self, 'callback': callback} + + i2v_extra_kwrags = {} + + if target_camera != None: + recam_dict = {'cam_emb': cam_emb} + i2v_extra_kwrags.update(recam_dict) + + if input_frames != None: + vace_dict = {'vace_context' : z, 'vace_context_scale' : context_scale} + i2v_extra_kwrags.update(vace_dict) + + + latent_length = (num_frames - 1) // 4 + 1 + latent_height = height // 8 + latent_width = width // 8 + if ar_step == 0: + causal_block_size = 1 + fps_embeds = [fps] #* prompt_embeds[0].shape[0] + fps_embeds = [0 if i == 16 else 1 for i in fps_embeds] + + self.scheduler.set_timesteps(sampling_steps, device=self.device, shift=shift) + init_timesteps = self.scheduler.timesteps + base_num_frames_iter = latent_length + latent_shape = [16, base_num_frames_iter, latent_height, latent_width] + + prefix_video = None + predix_video_latent_length = 0 + + if prefix_video is not None: + latents[0][:, :predix_video_latent_length] = prefix_video[0].to(torch.float32) + step_matrix, _, step_update_mask, valid_interval = self.generate_timestep_matrix( + base_num_frames_iter, + init_timesteps, + base_num_frames_iter, + ar_step, + predix_video_latent_length, + causal_block_size, + ) + sample_schedulers = [] + for _ in range(base_num_frames_iter): + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=1000, shift=1, use_dynamic_shifting=False + ) + sample_scheduler.set_timesteps(sampling_steps, device=self.device, shift=shift) + sample_schedulers.append(sample_scheduler) + sample_schedulers_counter = [0] * base_num_frames_iter + + updated_num_steps= len(step_matrix) + + if callback != None: + callback(-1, None, True, override_num_inference_steps = updated_num_steps) + if self.model.enable_teacache: + self.model.compute_teacache_threshold(self.model.teacache_start_step, timesteps, self.model.teacache_multiplier) + # if callback != None: + # callback(-1, None, True) + + for i, timestep_i in enumerate(tqdm(step_matrix)): + update_mask_i = step_update_mask[i] + valid_interval_i = valid_interval[i] + valid_interval_start, valid_interval_end = valid_interval_i + timestep = timestep_i[None, valid_interval_start:valid_interval_end].clone() + latent_model_input = [latents[0][:, valid_interval_start:valid_interval_end, :, :].clone()] + if addnoise_condition > 0 and valid_interval_start < predix_video_latent_length: + noise_factor = 0.001 * addnoise_condition + timestep_for_noised_condition = addnoise_condition + latent_model_input[0][:, valid_interval_start:predix_video_latent_length] = ( + latent_model_input[0][:, valid_interval_start:predix_video_latent_length] + * (1.0 - noise_factor) + + torch.randn_like( + latent_model_input[0][:, valid_interval_start:predix_video_latent_length] + ) + * noise_factor + ) + timestep[:, valid_interval_start:predix_video_latent_length] = timestep_for_noised_condition + kwrags = { + "x" : torch.stack([latent_model_input[0]]), + "t" : timestep, + "freqs" :freqs, + "fps" : fps_embeds, + "causal_block_size" : causal_block_size, + "causal_attention" : causal_attention, + "callback" : callback, + "pipeline" : self, + "current_step" : i, + } + kwrags.update(i2v_extra_kwrags) + + if not self.do_classifier_free_guidance: + noise_pred = self.model( + context=context, + **kwrags, + )[0] + if self._interrupt: + return None + noise_pred= noise_pred.to(torch.float32) + else: + if joint_pass: + noise_pred_cond, noise_pred_uncond = self.model( + context=context, + context2=context_null, + **kwrags, + ) + if self._interrupt: + return None + else: + noise_pred_cond = self.model( + context=context, + **kwrags, + )[0] + if self._interrupt: + return None + noise_pred_uncond = self.model( + context=context_null, + )[0] + if self._interrupt: + return None + noise_pred_cond= noise_pred_cond.to(torch.float32) + noise_pred_uncond= noise_pred_uncond.to(torch.float32) + noise_pred = noise_pred_uncond + guide_scale * (noise_pred_cond - noise_pred_uncond) + del noise_pred_cond, noise_pred_uncond + for idx in range(valid_interval_start, valid_interval_end): + if update_mask_i[idx].item(): + latents[0][:, idx] = sample_schedulers[idx].step( + noise_pred[:, idx - valid_interval_start], + timestep_i[idx], + latents[0][:, idx], + return_dict=False, + generator=seed_g, + )[0] + sample_schedulers_counter[idx] += 1 + if callback is not None: + callback(i, latents[0].squeeze(0), False) + + # for i, t in enumerate(tqdm(timesteps)): + # if target_camera != None: + # latent_model_input = [torch.cat([u,v], dim=1) for u,v in zip(latents,source_latents )] + # else: + # latent_model_input = latents + # slg_layers_local = None + # if int(slg_start * sampling_steps) <= i < int(slg_end * sampling_steps): + # slg_layers_local = slg_layers + # timestep = [t] + # offload.set_step_no_for_lora(self.model, i) + # timestep = torch.stack(timestep) + + # if joint_pass: + # noise_pred_cond, noise_pred_uncond = self.model( + # latent_model_input, t=timestep, current_step=i, slg_layers=slg_layers_local, **arg_both) + # if self._interrupt: + # return None + # else: + # noise_pred_cond = self.model( + # latent_model_input, t=timestep,current_step=i, is_uncond = False, **arg_c)[0] + # if self._interrupt: + # return None + # noise_pred_uncond = self.model( + # latent_model_input, t=timestep,current_step=i, is_uncond = True, slg_layers=slg_layers_local, **arg_null)[0] + # if self._interrupt: + # return None + + # # del latent_model_input + + # # CFG Zero *. Thanks to https://github.com/WeichenFan/CFG-Zero-star/ + # noise_pred_text = noise_pred_cond + # if cfg_star_switch: + # positive_flat = noise_pred_text.view(batch_size, -1) + # negative_flat = noise_pred_uncond.view(batch_size, -1) + + # alpha = optimized_scale(positive_flat,negative_flat) + # alpha = alpha.view(batch_size, 1, 1, 1) + + # if (i <= cfg_zero_step): + # noise_pred = noise_pred_text*0. # it would be faster not to compute noise_pred... + # else: + # noise_pred_uncond *= alpha + # noise_pred = noise_pred_uncond + guide_scale * (noise_pred_text - noise_pred_uncond) + # del noise_pred_uncond + + # temp_x0 = sample_scheduler.step( + # noise_pred[:, :target_shape[1]].unsqueeze(0), + # t, + # latents[0].unsqueeze(0), + # return_dict=False, + # generator=seed_g)[0] + # latents = [temp_x0.squeeze(0)] + # del temp_x0 + + # if callback is not None: + # callback(i, latents[0], False) + + x0 = latents + + if input_frames == None: + videos = self.vae.decode(x0, VAE_tile_size) + else: + videos = self.decode_latent(x0, input_ref_images, VAE_tile_size) + + del latents + del sample_scheduler + + return videos[0] if self.rank == 0 else None + + def adapt_vace_model(self): + model = self.model + modules_dict= { k: m for k, m in model.named_modules()} + for model_layer, vace_layer in model.vace_layers_mapping.items(): + module = modules_dict[f"vace_blocks.{vace_layer}"] + target = modules_dict[f"blocks.{model_layer}"] + setattr(target, "vace", module ) + delattr(model, "vace_blocks") + + \ No newline at end of file diff --git a/wan/trajectory_editor/app.py b/models/wan/trajectory_editor/app.py similarity index 100% rename from wan/trajectory_editor/app.py rename to models/wan/trajectory_editor/app.py diff --git a/models/wan/wan_handler.py b/models/wan/wan_handler.py new file mode 100644 index 0000000..240db67 --- /dev/null +++ b/models/wan/wan_handler.py @@ -0,0 +1,186 @@ +import torch + +def test_class_i2v(base_model_type): + return base_model_type in ["i2v", "i2v_2_2", "fun_inp_1.3B", "fun_inp", "flf2v_720p", "fantasy", "multitalk", ] #"hunyuan_i2v", + +class family_handler(): + @staticmethod + def get_wan_text_encoder_filename(text_encoder_quantization): + text_encoder_filename = "ckpts/umt5-xxl/models_t5_umt5-xxl-enc-bf16.safetensors" + if text_encoder_quantization =="int8": + text_encoder_filename = text_encoder_filename.replace("bf16", "quanto_int8") + return text_encoder_filename + + + + @staticmethod + def query_modules_files(): + return { + "vace_14B" : ["ckpts/wan2.1_Vace_14B_module_mbf16.safetensors", "ckpts/wan2.1_Vace_14B_module_quanto_mbf16_int8.safetensors", "ckpts/wan2.1_Vace_14B_module_quanto_mfp16_int8.safetensors"], + "vace_1.3B" : ["ckpts/wan2.1_Vace_1_3B_module.safetensors"], + "fantasy": ["ckpts/wan2.1_fantasy_speaking_14B_bf16.safetensors"], + "multitalk": ["ckpts/wan2.1_multitalk_14B_mbf16.safetensors", "ckpts/wan2.1_multitalk_14B_quanto_mbf16_int8.safetensors", "ckpts/wan2.1_multitalk_14B_quanto_mfp16_int8.safetensors"] +} + + @staticmethod + def query_model_def(base_model_type, model_def): + extra_model_def = {} + if "URLs2" in model_def: + extra_model_def["no_steps_skipping"] = True + extra_model_def["i2v_class"] = test_class_i2v(base_model_type) + + vace_class = base_model_type in ["vace_14B", "vace_1.3B", "vace_multitalk_14B"] + extra_model_def["vace_class"] = vace_class + + if base_model_type in ["multitalk", "vace_multitalk_14B"]: + fps = 25 + elif base_model_type in ["fantasy"]: + fps = 23 + elif base_model_type in ["ti2v_2_2"]: + fps = 24 + else: + fps = 16 + extra_model_def["fps"] =fps + + if vace_class: + frames_minimum, frames_steps = 17, 4 + else: + frames_minimum, frames_steps = 5, 4 + extra_model_def.update({ + "frames_minimum" : frames_minimum, + "frames_steps" : frames_steps, + "sliding_window" : base_model_type in ["multitalk", "t2v", "fantasy"] or test_class_i2v(base_model_type) or vace_class, #"ti2v_2_2", + "guidance_max_phases" : 2, + "skip_layer_guidance" : True, + "cfg_zero" : True, + "cfg_star" : True, + "adaptive_projected_guidance" : True, + "skip_steps_cache" : True, + }) + + return extra_model_def + + @staticmethod + def query_supported_types(): + return ["multitalk", "fantasy", "vace_14B", "vace_multitalk_14B", + "t2v_1.3B", "t2v", "vace_1.3B", "phantom_1.3B", "phantom_14B", + "recam_1.3B", + "i2v", "i2v_2_2", "ti2v_2_2", "flf2v_720p", "fun_inp_1.3B", "fun_inp"] + + + @staticmethod + def query_family_maps(): + + models_eqv_map = { + "flf2v_720p" : "i2v", + "t2v_1.3B" : "t2v", + } + + models_comp_map = { + "vace_14B" : [ "vace_multitalk_14B"], + "t2v" : [ "vace_14B", "vace_1.3B" "vace_multitalk_14B", "t2v_1.3B", "phantom_1.3B","phantom_14B"], + "i2v" : [ "fantasy", "multitalk", "flf2v_720p" ], + "fantasy": ["multitalk"], + } + return models_eqv_map, models_comp_map + + @staticmethod + def query_model_family(): + return "wan" + + @staticmethod + def query_family_infos(): + return {"wan":(0, "Wan2.1"), "wan2_2":(1, "Wan2.2") } + + @staticmethod + def get_vae_block_size(base_model_type): + return 32 if base_model_type == "ti2v_2_2" else 16 + + @staticmethod + def get_rgb_factors(model_type): + from shared.RGB_factors import get_rgb_factors + if model_type == "ti2v_2_2": return None, None + latent_rgb_factors, latent_rgb_factors_bias = get_rgb_factors("wan") + return latent_rgb_factors, latent_rgb_factors_bias + + @staticmethod + def query_model_files(computeList, base_model_type, model_filename, text_encoder_quantization): + text_encoder_filename = family_handler.get_wan_text_encoder_filename(text_encoder_quantization) + vae_filepath = "Wan2.2_VAE.safetensors" if base_model_type == "ti2v_2_2" else "Wan2.1_VAE.safetensors" + return { + "repoId" : "DeepBeepMeep/Wan2.1", + "sourceFolderList" : ["xlm-roberta-large", "umt5-xxl", "" ], + "fileList" : [ [ "models_clip_open-clip-xlm-roberta-large-vit-huge-14-bf16.safetensors", "sentencepiece.bpe.model", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json"], ["special_tokens_map.json", "spiece.model", "tokenizer.json", "tokenizer_config.json"] + computeList(text_encoder_filename) , [vae_filepath, "fantasy_proj_model.safetensors" ] + computeList(model_filename) ] + } + + + @staticmethod + def load_model(model_filename, model_type, base_model_type, model_def, quantizeTransformer = False, text_encoder_quantization = None, dtype = torch.bfloat16, VAE_dtype = torch.float32, mixed_precision_transformer = False, save_quantized= False): + from .configs import WAN_CONFIGS + + if test_class_i2v(base_model_type): + cfg = WAN_CONFIGS['i2v-14B'] + else: + cfg = WAN_CONFIGS['t2v-14B'] + # cfg = WAN_CONFIGS['t2v-1.3B'] + from . import WanAny2V + wan_model = WanAny2V( + config=cfg, + checkpoint_dir="ckpts", + model_filename=model_filename, + model_type = model_type, + model_def = model_def, + base_model_type=base_model_type, + text_encoder_filename= family_handler.get_wan_text_encoder_filename(text_encoder_quantization), + quantizeTransformer = quantizeTransformer, + dtype = dtype, + VAE_dtype = VAE_dtype, + mixed_precision_transformer = mixed_precision_transformer, + save_quantized = save_quantized + ) + + pipe = {"transformer": wan_model.model, "text_encoder" : wan_model.text_encoder.model, "vae": wan_model.vae.model } + if hasattr(wan_model,"model2") and wan_model.model2 is not None: + pipe["transformer2"] = wan_model.model2 + if hasattr(wan_model, "clip"): + pipe["text_encoder_2"] = wan_model.clip.model + return wan_model, pipe + + @staticmethod + def update_default_settings(base_model_type, model_def, ui_defaults): + if base_model_type in ["fantasy"]: + ui_defaults.update({ + "audio_guidance_scale": 5.0, + "sliding_window_size": 1, + }) + + elif base_model_type in ["multitalk"]: + ui_defaults.update({ + "guidance_scale": 5.0, + "flow_shift": 7, # 11 for 720p + "audio_guidance_scale": 4, + "sliding_window_discard_last_frames" : 4, + "sample_solver" : "euler", + "adaptive_switch" : 1, + }) + + elif base_model_type in ["phantom_1.3B", "phantom_14B"]: + ui_defaults.update({ + "guidance_scale": 7.5, + "flow_shift": 5, + "remove_background_images_ref": 1, + "video_prompt_type": "I", + # "resolution": "1280x720" + }) + + elif base_model_type in ["vace_14B", "vace_multitalk_14B"]: + ui_defaults.update({ + "sliding_window_discard_last_frames": 0, + }) + + elif base_model_type in ["ti2v_2_2"]: + ui_defaults.update({ + "image_prompt_type": "T", + }) + + \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 44bda57..827b29b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,8 @@ -torch>=2.4.0 +torch>=2.6.0 torchvision>=0.19.0 opencv-python>=4.9.0.80 -diffusers>=0.31.0 -transformers==4.51.3 +diffusers==0.34.0 +transformers==4.52.1 #transformers==4.46.3 # was needed by llamallava used by i2v hunyuan before patch tokenizers>=0.20.3 accelerate>=1.1.1 @@ -46,6 +46,6 @@ soundfile ffmpeg-python pyannote.audio pynvml -huggingface_hub[hf_xet] +#huggingface_hub[hf_xet] #slow down everything !!!! # num2words # spacy \ No newline at end of file diff --git a/shared/RGB_factors.py b/shared/RGB_factors.py new file mode 100644 index 0000000..5ec1e59 --- /dev/null +++ b/shared/RGB_factors.py @@ -0,0 +1,213 @@ +# thanks Comfyui for the rgb factors +def get_rgb_factors(model_family): + if model_family == "wan": + latent_channels = 16 + latent_dimensions = 3 + latent_rgb_factors = [ + [-0.1299, -0.1692, 0.2932], + [ 0.0671, 0.0406, 0.0442], + [ 0.3568, 0.2548, 0.1747], + [ 0.0372, 0.2344, 0.1420], + [ 0.0313, 0.0189, -0.0328], + [ 0.0296, -0.0956, -0.0665], + [-0.3477, -0.4059, -0.2925], + [ 0.0166, 0.1902, 0.1975], + [-0.0412, 0.0267, -0.1364], + [-0.1293, 0.0740, 0.1636], + [ 0.0680, 0.3019, 0.1128], + [ 0.0032, 0.0581, 0.0639], + [-0.1251, 0.0927, 0.1699], + [ 0.0060, -0.0633, 0.0005], + [ 0.3477, 0.2275, 0.2950], + [ 0.1984, 0.0913, 0.1861] + ] + + latent_rgb_factors_bias = [-0.1835, -0.0868, -0.3360] + + # latent_rgb_factors_bias = [0.0259, -0.0192, -0.0761] + elif model_family =="flux": + scale_factor = 0.3611 + shift_factor = 0.1159 + latent_rgb_factors =[ + [-0.0346, 0.0244, 0.0681], + [ 0.0034, 0.0210, 0.0687], + [ 0.0275, -0.0668, -0.0433], + [-0.0174, 0.0160, 0.0617], + [ 0.0859, 0.0721, 0.0329], + [ 0.0004, 0.0383, 0.0115], + [ 0.0405, 0.0861, 0.0915], + [-0.0236, -0.0185, -0.0259], + [-0.0245, 0.0250, 0.1180], + [ 0.1008, 0.0755, -0.0421], + [-0.0515, 0.0201, 0.0011], + [ 0.0428, -0.0012, -0.0036], + [ 0.0817, 0.0765, 0.0749], + [-0.1264, -0.0522, -0.1103], + [-0.0280, -0.0881, -0.0499], + [-0.1262, -0.0982, -0.0778] + ] + latent_rgb_factors_bias = [-0.0329, -0.0718, -0.0851] + + elif model_family == "ltxv": + latent_channels = 128 + latent_dimensions = 3 + + latent_rgb_factors = [ + [ 1.1202e-02, -6.3815e-04, -1.0021e-02], + [ 8.6031e-02, 6.5813e-02, 9.5409e-04], + [-1.2576e-02, -7.5734e-03, -4.0528e-03], + [ 9.4063e-03, -2.1688e-03, 2.6093e-03], + [ 3.7636e-03, 1.2765e-02, 9.1548e-03], + [ 2.1024e-02, -5.2973e-03, 3.4373e-03], + [-8.8896e-03, -1.9703e-02, -1.8761e-02], + [-1.3160e-02, -1.0523e-02, 1.9709e-03], + [-1.5152e-03, -6.9891e-03, -7.5810e-03], + [-1.7247e-03, 4.6560e-04, -3.3839e-03], + [ 1.3617e-02, 4.7077e-03, -2.0045e-03], + [ 1.0256e-02, 7.7318e-03, 1.3948e-02], + [-1.6108e-02, -6.2151e-03, 1.1561e-03], + [ 7.3407e-03, 1.5628e-02, 4.4865e-04], + [ 9.5357e-04, -2.9518e-03, -1.4760e-02], + [ 1.9143e-02, 1.0868e-02, 1.2264e-02], + [ 4.4575e-03, 3.6682e-05, -6.8508e-03], + [-4.5681e-04, 3.2570e-03, 7.7929e-03], + [ 3.3902e-02, 3.3405e-02, 3.7454e-02], + [-2.3001e-02, -2.4877e-03, -3.1033e-03], + [ 5.0265e-02, 3.8841e-02, 3.3539e-02], + [-4.1018e-03, -1.1095e-03, 1.5859e-03], + [-1.2689e-01, -1.3107e-01, -2.1005e-01], + [ 2.6276e-02, 1.4189e-02, -3.5963e-03], + [-4.8679e-03, 8.8486e-03, 7.8029e-03], + [-1.6610e-03, -4.8597e-03, -5.2060e-03], + [-2.1010e-03, 2.3610e-03, 9.3796e-03], + [-2.2482e-02, -2.1305e-02, -1.5087e-02], + [-1.5753e-02, -1.0646e-02, -6.5083e-03], + [-4.6975e-03, 5.0288e-03, -6.7390e-03], + [ 1.1951e-02, 2.0712e-02, 1.6191e-02], + [-6.3704e-03, -8.4827e-03, -9.5483e-03], + [ 7.2610e-03, -9.9326e-03, -2.2978e-02], + [-9.1904e-04, 6.2882e-03, 9.5720e-03], + [-3.7178e-02, -3.7123e-02, -5.6713e-02], + [-1.3373e-01, -1.0720e-01, -5.3801e-02], + [-5.3702e-03, 8.1256e-03, 8.8397e-03], + [-1.5247e-01, -2.1437e-01, -2.1843e-01], + [ 3.1441e-02, 7.0335e-03, -9.7541e-03], + [ 2.1528e-03, -8.9817e-03, -2.1023e-02], + [ 3.8461e-03, -5.8957e-03, -1.5014e-02], + [-4.3470e-03, -1.2940e-02, -1.5972e-02], + [-5.4781e-03, -1.0842e-02, -3.0204e-03], + [-6.5347e-03, 3.0806e-03, -1.0163e-02], + [-5.0414e-03, -7.1503e-03, -8.9686e-04], + [-8.5851e-03, -2.4351e-03, 1.0674e-03], + [-9.0016e-03, -9.6493e-03, 1.5692e-03], + [ 5.0914e-03, 1.2099e-02, 1.9968e-02], + [ 1.3758e-02, 1.1669e-02, 8.1958e-03], + [-1.0518e-02, -1.1575e-02, -4.1307e-03], + [-2.8410e-02, -3.1266e-02, -2.2149e-02], + [ 2.9336e-03, 3.6511e-02, 1.8717e-02], + [-1.6703e-02, -1.6696e-02, -4.4529e-03], + [ 4.8818e-02, 4.0063e-02, 8.7410e-03], + [-1.5066e-02, -5.7328e-04, 2.9785e-03], + [-1.7613e-02, -8.1034e-03, 1.3086e-02], + [-9.2633e-03, 1.0803e-02, -6.3489e-03], + [ 3.0851e-03, 4.7750e-04, 1.2347e-02], + [-2.2785e-02, -2.3043e-02, -2.6005e-02], + [-2.4787e-02, -1.5389e-02, -2.2104e-02], + [-2.3572e-02, 1.0544e-03, 1.2361e-02], + [-7.8915e-03, -1.2271e-03, -6.0968e-03], + [-1.1478e-02, -1.2543e-03, 6.2679e-03], + [-5.4229e-02, 2.6644e-02, 6.3394e-03], + [ 4.4216e-03, -7.3338e-03, -1.0464e-02], + [-4.5013e-03, 1.6082e-03, 1.4420e-02], + [ 1.3673e-02, 8.8877e-03, 4.1253e-03], + [-1.0145e-02, 9.0072e-03, 1.5695e-02], + [-5.6234e-03, 1.1847e-03, 8.1261e-03], + [-3.7171e-03, -5.3538e-03, 1.2590e-03], + [ 2.9476e-02, 2.1424e-02, 3.0424e-02], + [-3.4925e-02, -2.4340e-02, -2.5316e-02], + [-3.4127e-02, -2.2406e-02, -1.0589e-02], + [-1.7342e-02, -1.3249e-02, -1.0719e-02], + [-2.1478e-03, -8.6051e-03, -2.9878e-03], + [ 1.2089e-03, -4.2391e-03, -6.8569e-03], + [ 9.0411e-04, -6.6886e-03, -6.7547e-05], + [ 1.6048e-02, -1.0057e-02, -2.8929e-02], + [ 1.2290e-03, 1.0163e-02, 1.8861e-02], + [ 1.7264e-02, 2.7257e-04, 1.3785e-02], + [-1.3482e-02, -3.6427e-03, 6.7481e-04], + [ 4.6782e-03, -5.2423e-03, 2.4467e-03], + [-5.9113e-03, -6.2244e-03, -1.8162e-03], + [ 1.5496e-02, 1.4582e-02, 1.9514e-03], + [ 7.4958e-03, 1.5886e-03, -8.2305e-03], + [ 1.9086e-02, 1.6360e-03, -3.9674e-03], + [-5.7021e-03, -2.7307e-03, -4.1066e-03], + [ 1.7450e-03, 1.4602e-02, 2.5794e-02], + [-8.2788e-04, 2.2902e-03, 4.5161e-03], + [ 1.1632e-02, 8.9193e-03, -7.2813e-03], + [ 7.5721e-03, 2.6784e-03, 1.1393e-02], + [ 5.1939e-03, 3.6903e-03, 1.4049e-02], + [-1.8383e-02, -2.2529e-02, -2.4477e-02], + [ 5.8842e-04, -5.7874e-03, -1.4770e-02], + [-1.6125e-02, -8.6101e-03, -1.4533e-02], + [ 2.0540e-02, 2.0729e-02, 6.4338e-03], + [ 3.3587e-03, -1.1226e-02, -1.6444e-02], + [-1.4742e-03, -1.0489e-02, 1.7097e-03], + [ 2.8130e-02, 2.3546e-02, 3.2791e-02], + [-1.8532e-02, -1.2842e-02, -8.7756e-03], + [-8.0533e-03, -1.0771e-02, -1.7536e-02], + [-3.9009e-03, 1.6150e-02, 3.3359e-02], + [-7.4554e-03, -1.4154e-02, -6.1910e-03], + [ 3.4734e-03, -1.1370e-02, -1.0581e-02], + [ 1.1476e-02, 3.9281e-03, 2.8231e-03], + [ 7.1639e-03, -1.4741e-03, -3.8066e-03], + [ 2.2250e-03, -8.7552e-03, -9.5719e-03], + [ 2.4146e-02, 2.1696e-02, 2.8056e-02], + [-5.4365e-03, -2.4291e-02, -1.7802e-02], + [ 7.4263e-03, 1.0510e-02, 1.2705e-02], + [ 6.2669e-03, 6.2658e-03, 1.9211e-02], + [ 1.6378e-02, 9.4933e-03, 6.6971e-03], + [ 1.7173e-02, 2.3601e-02, 2.3296e-02], + [-1.4568e-02, -9.8279e-03, -1.1556e-02], + [ 1.4431e-02, 1.4430e-02, 6.6362e-03], + [-6.8230e-03, 1.8863e-02, 1.4555e-02], + [ 6.1156e-03, 3.4700e-03, -2.6662e-03], + [-2.6983e-03, -5.9402e-03, -9.2276e-03], + [ 1.0235e-02, 7.4173e-03, -7.6243e-03], + [-1.3255e-02, 1.9322e-02, -9.2153e-04], + [ 2.4222e-03, -4.8039e-03, -1.5759e-02], + [ 2.6244e-02, 2.5951e-02, 2.0249e-02], + [ 1.5711e-02, 1.8498e-02, 2.7407e-03], + [-2.1714e-03, 4.7214e-03, -2.2443e-02], + [-7.4747e-03, 7.4166e-03, 1.4430e-02], + [-8.3906e-03, -7.9776e-03, 9.7927e-03], + [ 3.8321e-02, 9.6622e-03, -1.9268e-02], + [-1.4605e-02, -6.7032e-03, 3.9675e-03] + ] + latent_rgb_factors_bias = [-0.0571, -0.1657, -0.2512] + + elif model_family == "hunyuan": + latent_channels = 16 + latent_dimensions = 3 + scale_factor = 0.476986 + latent_rgb_factors = [ + [-0.0395, -0.0331, 0.0445], + [ 0.0696, 0.0795, 0.0518], + [ 0.0135, -0.0945, -0.0282], + [ 0.0108, -0.0250, -0.0765], + [-0.0209, 0.0032, 0.0224], + [-0.0804, -0.0254, -0.0639], + [-0.0991, 0.0271, -0.0669], + [-0.0646, -0.0422, -0.0400], + [-0.0696, -0.0595, -0.0894], + [-0.0799, -0.0208, -0.0375], + [ 0.1166, 0.1627, 0.0962], + [ 0.1165, 0.0432, 0.0407], + [-0.2315, -0.1920, -0.1355], + [-0.0270, 0.0401, -0.0821], + [-0.0616, -0.0997, -0.0727], + [ 0.0249, -0.0469, -0.1703] + ] + + latent_rgb_factors_bias = [ 0.0259, -0.0192, -0.0761] + else: + latent_rgb_factors_bias = latent_rgb_factors = None + return latent_rgb_factors, latent_rgb_factors_bias \ No newline at end of file diff --git a/wan/modules/attention.py b/shared/attention.py similarity index 100% rename from wan/modules/attention.py rename to shared/attention.py diff --git a/wan/modules/sage2_core.py b/shared/sage2_core.py similarity index 100% rename from wan/modules/sage2_core.py rename to shared/sage2_core.py diff --git a/wan/utils/__init__.py b/shared/utils/__init__.py similarity index 100% rename from wan/utils/__init__.py rename to shared/utils/__init__.py diff --git a/wan/utils/basic_flowmatch.py b/shared/utils/basic_flowmatch.py similarity index 100% rename from wan/utils/basic_flowmatch.py rename to shared/utils/basic_flowmatch.py diff --git a/wan/utils/cammmaster_tools.py b/shared/utils/cammmaster_tools.py similarity index 100% rename from wan/utils/cammmaster_tools.py rename to shared/utils/cammmaster_tools.py diff --git a/wan/utils/fm_solvers.py b/shared/utils/fm_solvers.py similarity index 100% rename from wan/utils/fm_solvers.py rename to shared/utils/fm_solvers.py diff --git a/wan/utils/fm_solvers_unipc.py b/shared/utils/fm_solvers_unipc.py similarity index 100% rename from wan/utils/fm_solvers_unipc.py rename to shared/utils/fm_solvers_unipc.py diff --git a/wan/utils/loras_mutipliers.py b/shared/utils/loras_mutipliers.py similarity index 100% rename from wan/utils/loras_mutipliers.py rename to shared/utils/loras_mutipliers.py diff --git a/wan/utils/motion.py b/shared/utils/motion.py similarity index 100% rename from wan/utils/motion.py rename to shared/utils/motion.py diff --git a/wan/utils/notification_sound.py b/shared/utils/notification_sound.py similarity index 100% rename from wan/utils/notification_sound.py rename to shared/utils/notification_sound.py diff --git a/wan/utils/prompt_extend.py b/shared/utils/prompt_extend.py similarity index 100% rename from wan/utils/prompt_extend.py rename to shared/utils/prompt_extend.py diff --git a/wan/utils/prompt_parser.py b/shared/utils/prompt_parser.py similarity index 100% rename from wan/utils/prompt_parser.py rename to shared/utils/prompt_parser.py diff --git a/wan/utils/qwen_vl_utils.py b/shared/utils/qwen_vl_utils.py similarity index 100% rename from wan/utils/qwen_vl_utils.py rename to shared/utils/qwen_vl_utils.py diff --git a/wan/utils/stats.py b/shared/utils/stats.py similarity index 100% rename from wan/utils/stats.py rename to shared/utils/stats.py diff --git a/wan/utils/thread_utils.py b/shared/utils/thread_utils.py similarity index 100% rename from wan/utils/thread_utils.py rename to shared/utils/thread_utils.py diff --git a/wan/utils/utils.py b/shared/utils/utils.py similarity index 100% rename from wan/utils/utils.py rename to shared/utils/utils.py diff --git a/wan/utils/vace_preprocessor.py b/shared/utils/vace_preprocessor.py similarity index 100% rename from wan/utils/vace_preprocessor.py rename to shared/utils/vace_preprocessor.py diff --git a/wan/distributed/__init__.py b/wan/distributed/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/wan/trajectory_editor/templates/index.html b/wan/trajectory_editor/templates/index.html deleted file mode 100644 index 2ac8c78..0000000 --- a/wan/trajectory_editor/templates/index.html +++ /dev/null @@ -1,571 +0,0 @@ - - - - - - - Track Point Editor - - - -

Track Point Editor

- - -
- - -
- - - - - -
-
- - - -
-
- - - -
-
- - -
-
- -
-
- - - -
- -
- - - - -
-
- - - - - diff --git a/wgp.py b/wgp.py index 5a54b75..48ccc1b 100644 --- a/wgp.py +++ b/wgp.py @@ -13,14 +13,13 @@ from datetime import datetime import gradio as gr import random import json -import wan -from wan.utils import notification_sound -from wan.configs import MAX_AREA_CONFIGS, WAN_CONFIGS, SUPPORTED_SIZES, VACE_SIZE_CONFIGS -from wan.utils.loras_mutipliers import preparse_loras_multipliers, parse_loras_multipliers -from wan.utils.utils import cache_video, convert_tensor_to_image, save_image, get_video_info, get_file_creation_date, convert_image_to_video -from wan.utils.utils import extract_audio_tracks, combine_video_with_audio_tracks, combine_and_concatenate_video_with_audio_tracks, cleanup_temp_audio_files, calculate_new_dimensions +import importlib +from shared.utils import notification_sound +from shared.utils.loras_mutipliers import preparse_loras_multipliers, parse_loras_multipliers +from shared.utils.utils import cache_video, convert_tensor_to_image, save_image, get_video_info, get_file_creation_date, convert_image_to_video +from shared.utils.utils import extract_audio_tracks, combine_video_with_audio_tracks, combine_and_concatenate_video_with_audio_tracks, cleanup_temp_audio_files, calculate_new_dimensions -from wan.modules.attention import get_attention_modes, get_supported_attention_modes +from shared.attention import get_attention_modes, get_supported_attention_modes from huggingface_hub import hf_hub_download, snapshot_download import torch import gc @@ -29,7 +28,7 @@ import math import typing import asyncio import inspect -from wan.utils import prompt_parser +from shared.utils import prompt_parser import base64 import io from PIL import Image @@ -51,7 +50,7 @@ AUTOSAVE_FILENAME = "queue.zip" PROMPT_VARS_MAX = 10 target_mmgp_version = "3.5.6" -WanGP_version = "7.61" +WanGP_version = "7.7S" settings_version = 2.23 max_source_video_frames = 3000 prompt_enhancer_image_caption_model, prompt_enhancer_image_caption_processor, prompt_enhancer_llm_model, prompt_enhancer_llm_tokenizer = None, None, None, None @@ -118,7 +117,7 @@ def pil_to_base64_uri(pil_image, format="png", quality=75): return None if isinstance(pil_image, str): - from wan.utils.utils import get_video_frame + from shared.utils.utils import get_video_frame pil_image = get_video_frame(pil_image, 0) buffer = io.BytesIO() @@ -178,7 +177,7 @@ def process_prompt_and_add_tasks(state, model_choice): return get_queue_table(queue) model_def = get_model_def(model_type) image_outputs = inputs["image_mode"] == 1 - no_steps_skipping = model_def.get("no_steps_skipping", False) + any_steps_skipping = model_def.get("skip_steps_cache", False) model_type = get_base_model_type(model_type) inputs["model_filename"] = model_filename @@ -302,7 +301,7 @@ def process_prompt_and_add_tasks(state, model_choice): gr.Info(f"Error parsing Loras Multipliers: {errors}") return - if no_steps_skipping: skip_steps_cache_type = "" + if not any_steps_skipping: skip_steps_cache_type = "" if switch_threshold is not None and switch_threshold != 0 and len(skip_steps_cache_type) > 0: gr.Info("Steps skipping is not yet supported if Switch Threshold is not null") return @@ -321,7 +320,7 @@ def process_prompt_and_add_tasks(state, model_choice): audio_prompt_type = "" if "B" in audio_prompt_type or "X" in audio_prompt_type: - from wan.multitalk.multitalk import parse_speakers_locations + from models.wan.multitalk.multitalk import parse_speakers_locations speakers_bboxes, error = parse_speakers_locations(speakers_locations) if len(error) > 0: gr.Info(error) @@ -1398,6 +1397,12 @@ def _parse_args(): help="Path to a directory that contains flux images Loras" ) + parser.add_argument( + "--lora-dir-qwen", + type=str, + default="loras_qwen", + help="Path to a directory that contains qwen images Loras" + ) parser.add_argument( "--check-loras", @@ -1630,6 +1635,10 @@ def get_lora_dir(model_type): lora_dir_1_3B = os.path.join(root_lora_dir, "1.3B") if os.path.isdir(lora_dir_1_3B ): return lora_dir_1_3B + elif model_type == "ti2v_2_2": + lora_dir_5B = os.path.join(root_lora_dir, "5B") + if os.path.isdir(lora_dir_5B ): + return lora_dir_5B else: lora_dir_14B = os.path.join(root_lora_dir, "14B") if os.path.isdir(lora_dir_14B ): @@ -1644,6 +1653,8 @@ def get_lora_dir(model_type): return args.lora_dir_hunyuan_i2v else: return args.lora_dir_hunyuan + elif model_family =="qwen": + return args.lora_dir_qwen else: raise Exception("loras unknown") @@ -1706,7 +1717,6 @@ if not Path(server_config_filename).is_file(): "save_path": "outputs", #os.path.join(os.getcwd(), "compile" : "", "metadata_type": "metadata", - "default_ui": "t2v", "boost" : 1, "clear_file_list" : 5, "vae_config": 0, @@ -1734,24 +1744,10 @@ for path in ["wan2.1_Vace_1.3B_preview_bf16.safetensors", "sky_reels2_diffusion print(f"Removing old version of model '{path}'. A new version of this model will be downloaded next time you use it.") os.remove( os.path.join("ckpts" , path)) -families_infos = {"wan":(0, "Wan2.1"), "wan2_2":(1, "Wan2.2"), "ltxv":(10, "LTX Video"), "hunyuan":(20, "Hunyuan Video"), "flux":(30, "Flux 1"), "unknown": (100, "Unknown") } models_def = {} +family_handlers = ["models.wan.wan_handler", "models.wan.df_handler", "models.hyvideo.hunyuan_handler", "models.ltx_video.ltxv_handler", "models.flux.flux_handler", "models.qwen.qwen_handler"] -modules_files = { - "vace_14B" : ["ckpts/wan2.1_Vace_14B_module_mbf16.safetensors", "ckpts/wan2.1_Vace_14B_module_quanto_mbf16_int8.safetensors", "ckpts/wan2.1_Vace_14B_module_quanto_mfp16_int8.safetensors"], - "vace_1.3B" : ["ckpts/wan2.1_Vace_1_3B_module.safetensors"], - "fantasy": ["ckpts/wan2.1_fantasy_speaking_14B_bf16.safetensors"], - "multitalk": ["ckpts/wan2.1_multitalk_14B_mbf16.safetensors", "ckpts/wan2.1_multitalk_14B_quanto_mbf16_int8.safetensors", "ckpts/wan2.1_multitalk_14B_quanto_mfp16_int8.safetensors"] -} - -# architectures supported -base_types = ["multitalk", "fantasy", "vace_14B", "vace_multitalk_14B", - "t2v_1.3B", "t2v", "vace_1.3B", "phantom_1.3B", "phantom_14B", - "recam_1.3B", "sky_df_1.3B", "sky_df_14B", - "i2v", "i2v_2_2", "flf2v_720p", "fun_inp_1.3B", "fun_inp", "ltxv_13B", - "hunyuan", "hunyuan_i2v", "hunyuan_custom", "hunyuan_custom_audio", "hunyuan_custom_edit", "hunyuan_avatar", "flux" - ] # only needed for imported old settings files model_signatures = {"t2v": "text2video_14B", "t2v_1.3B" : "text2video_1.3B", "fun_inp_1.3B" : "Fun_InP_1.3B", "fun_inp" : "Fun_InP_14B", @@ -1762,36 +1758,50 @@ model_signatures = {"t2v": "text2video_14B", "t2v_1.3B" : "text2video_1.3B", " "hunyuan" : "hunyuan_video_720", "hunyuan_i2v" : "hunyuan_video_i2v_720", "hunyuan_custom" : "hunyuan_video_custom_720", "hunyuan_custom_audio" : "hunyuan_video_custom_audio", "hunyuan_custom_edit" : "hunyuan_video_custom_edit", "hunyuan_avatar" : "hunyuan_video_avatar" } + +def map_family_handlers(family_handlers): + base_types_handlers, families_infos, models_eqv_map, models_comp_map = {}, {"unknown": (100, "Unknown")}, {}, {} + for path in family_handlers: + handler = importlib.import_module(path).family_handler + for model_type in handler.query_supported_types(): + if model_type in base_types_handlers: + prev = base_types_handlers[model_type].__name__ + raise Exception(f"Model type {model_type} supported by {prev} and {handler.__name__}") + base_types_handlers[model_type] = handler + families_infos.update(handler.query_family_infos()) + eq_map, comp_map = handler.query_family_maps() + models_eqv_map.update(eq_map); models_comp_map.update(comp_map) + return base_types_handlers, families_infos, models_eqv_map, models_comp_map + +model_types_handlers, families_infos, models_eqv_map, models_comp_map = map_family_handlers(family_handlers) + def get_base_model_type(model_type): model_def = get_model_def(model_type) if model_def == None: - return model_type if model_type in base_types else None + return model_type if model_type in model_types_handlers else None # return model_type else: return model_def["architecture"] +def get_model_handler(model_type): + base_model_type = get_base_model_type(model_type) + if base_model_type is None: + raise Exception(f"Unknown model type {model_type}") + model_handler = model_types_handlers.get(base_model_type, None) + if model_handler is None: + raise Exception(f"No model handler found for base model type {base_model_type}") + return model_handler + def are_model_types_compatible(imported_model_type, current_model_type): imported_base_model_type = get_base_model_type(imported_model_type) curent_base_model_type = get_base_model_type(current_model_type) if imported_base_model_type == curent_base_model_type: return True - eqv_map = { - "flf2v_720p" : "i2v", - "t2v_1.3B" : "t2v", - "sky_df_1.3B" : "sky_df_14B", - } - if imported_base_model_type in eqv_map: - imported_base_model_type = eqv_map[imported_base_model_type] - comp_map = { - "vace_14B" : [ "vace_multitalk_14B"], - "t2v" : [ "vace_14B", "vace_1.3B" "vace_multitalk_14B", "t2v_1.3B", "phantom_1.3B","phantom_14B"], - "i2v" : [ "fantasy", "multitalk", "flf2v_720p" ], - "fantasy": ["multitalk"], - "sky_df_14B": ["sky_df_1.3B"], - "hunyuan_custom": ["hunyuan_custom_edit", "hunyuan_custom_audio"], - } - comp_list= comp_map.get(imported_base_model_type, None) + if imported_base_model_type in models_eqv_map: + imported_base_model_type = models_eqv_map[imported_base_model_type] + + comp_list= models_comp_map.get(imported_base_model_type, None) if comp_list == None: return False return curent_base_model_type in comp_list @@ -1817,51 +1827,32 @@ def get_model_family(model_type, for_ui = False): model_family = model_def.get("group", None) if model_family is not None and model_family in families_infos: return model_family - - if "hunyuan" in base_model_type : - return "hunyuan" - elif "ltxv" in base_model_type: - return "ltxv" - elif "flux" in base_model_type: - return "flux" - else: - return "wan" + handler = model_types_handlers.get(base_model_type, None) + if handler is None: + return "unknown" + return handler.query_model_family() -def test_class_i2v(model_type): - model_type = get_base_model_type(model_type) - return model_type in ["i2v", "i2v_2_2", "fun_inp_1.3B", "fun_inp", "flf2v_720p", "fantasy", "multitalk" ] #"hunyuan_i2v", +def test_class_i2v(model_type): + model_def = get_model_def(model_type) + return model_def.get("i2v_class", False) def test_vace_module(model_type): - model_type = get_base_model_type(model_type) - return model_type in ["vace_14B", "vace_1.3B", "vace_multitalk_14B"] + model_def = get_model_def(model_type) + return model_def.get("vace_class", False) def test_any_sliding_window(model_type): - model_type = get_base_model_type(model_type) - return test_vace_module(model_type) or model_type in ["sky_df_1.3B", "sky_df_14B", "ltxv_13B", "multitalk", "t2v", "fantasy"] or test_class_i2v(model_type) + model_def = get_model_def(model_type) + return model_def.get("sliding_window", False) def get_model_min_frames_and_step(model_type): - model_type = get_base_model_type(model_type) - if model_type in ["sky_df_14B"]: - return 17, 20 - elif model_type in ["ltxv_13B"]: - return 17, 8 - elif test_vace_module(model_type): - return 17, 4 - else: - return 5, 4 - + mode_def = get_model_def(model_type) + frames_minimum = mode_def.get("frames_minimum", 5) + frames_steps = mode_def.get("frames_steps", 4) + return frames_minimum, frames_steps + def get_model_fps(model_type): - model_type = get_base_model_type(model_type) - if model_type in ["hunyuan_avatar", "hunyuan_custom_audio", "multitalk", "vace_multitalk_14B"]: - fps = 25 - elif model_type in ["sky_df_14B", "hunyuan", "hunyuan_i2v", "hunyuan_custom_edit", "hunyuan_custom"]: - fps = 24 - elif model_type in ["fantasy"]: - fps = 23 - elif model_type in ["ltxv_13B"]: - fps = 30 - else: - fps = 16 + mode_def = get_model_def(model_type) + fps= mode_def.get("fps", 16) return fps def get_computed_fps(force_fps, base_model_type , video_guide, video_source ): @@ -1912,10 +1903,13 @@ def get_model_recursive_prop(model_type, prop = "URLs", return_list = True, sta raise Exception(f"Unknown model type '{model_type}'") -def get_model_filename(model_type, quantization ="int8", dtype_policy = "", is_module = False, submodel_no = 1, stack=[]): - if is_module: - choices = modules_files.get(model_type, None) - if choices == None: raise Exception(f"Invalid Module Id '{model_type}'") +def get_model_filename(model_type, quantization ="int8", dtype_policy = "", module_type = None, submodel_no = 1, stack=[]): + if module_type is not None: + base_model_type = get_base_model_type(model_type) + model_type_handler = model_types_handlers[base_model_type] + modules_files = model_type_handler.query_modules_files() if hasattr(model_type_handler, "query_modules_files") else {} + choices = modules_files.get(module_type, None) + if choices == None: raise Exception(f"Invalid Module Id '{module_type}'") else: key_name = "URLs" if submodel_no <= 1 else f"URLs{submodel_no}" @@ -2063,7 +2057,6 @@ def get_default_settings(model_type): "repeat_generation": 1, "multi_images_gen_type": 0, "guidance_scale": 5.0, - "embedded_guidance_scale" : 6.0, "flow_shift": 7.0 if not "720" in base_model_type and i2v else 5.0, "negative_prompt": "", "activated_loras": [], @@ -2076,87 +2069,8 @@ def get_default_settings(model_type): "slg_start_perc": 10, "slg_end_perc": 90 } - if base_model_type in ["fantasy"]: - ui_defaults["audio_guidance_scale"] = 5.0 - elif base_model_type in ["multitalk"]: - ui_defaults.update({ - "guidance_scale": 5.0, - "flow_shift": 7, # 11 for 720p - "audio_guidance_scale": 4, - "sliding_window_discard_last_frames" : 4, - "sample_solver" : "euler", - "adaptive_switch" : 1, - }) - - elif base_model_type in ["hunyuan","hunyuan_i2v"]: - ui_defaults.update({ - "guidance_scale": 7.0, - }) - - elif base_model_type in ["flux"]: - ui_defaults.update({ - "embedded_guidance": 2.5, - }) - if model_def.get("reference_image", False): - ui_defaults.update({ - "video_prompt_type": "KI", - }) - elif base_model_type in ["sky_df_1.3B", "sky_df_14B"]: - ui_defaults.update({ - "guidance_scale": 6.0, - "flow_shift": 8, - "sliding_window_discard_last_frames" : 0, - "resolution": "1280x720" if "720" in base_model_type else "960x544", - "sliding_window_size" : 121 if "720" in base_model_type else 97, - "RIFLEx_setting": 2, - "guidance_scale": 6, - "flow_shift": 8, - }) - - - elif base_model_type in ["phantom_1.3B", "phantom_14B"]: - ui_defaults.update({ - "guidance_scale": 7.5, - "flow_shift": 5, - "remove_background_images_ref": 1, - "video_prompt_type": "I", - # "resolution": "1280x720" - }) - - elif base_model_type in ["hunyuan_custom"]: - ui_defaults.update({ - "guidance_scale": 7.5, - "flow_shift": 13, - "resolution": "1280x720", - "video_prompt_type": "I", - }) - elif base_model_type in ["hunyuan_custom_audio"]: - ui_defaults.update({ - "guidance_scale": 7.5, - "flow_shift": 13, - "video_prompt_type": "I", - }) - elif base_model_type in ["hunyuan_custom_edit"]: - ui_defaults.update({ - "guidance_scale": 7.5, - "flow_shift": 13, - "video_prompt_type": "MVAI", - "sliding_window_size": 129, - }) - elif base_model_type in ["hunyuan_avatar"]: - ui_defaults.update({ - "guidance_scale": 7.5, - "flow_shift": 5, - "remove_background_images_ref": 0, - "skip_steps_start_step_perc": 25, - "video_length": 129, - "video_prompt_type": "I", - }) - elif base_model_type in ["vace_14B", "vace_multitalk_14B"]: - ui_defaults.update({ - "sliding_window_discard_last_frames": 0, - }) - + model_handler = get_model_handler(model_type) + model_handler.update_default_settings(base_model_type, model_def, ui_defaults) ui_defaults_update = model_def.get("settings", None) if ui_defaults_update is not None: ui_defaults.update(ui_defaults_update) @@ -2182,27 +2096,13 @@ def get_default_settings(model_type): ui_defaults["num_inference_steps"] = default_number_steps return ui_defaults -def get_model_query_handler(model_type): - base_model_type = get_base_model_type(model_type) - model_family= get_model_family(base_model_type) - if model_family == "wan": - if base_model_type in ("sky_df_1.3B", "sky_df_14B"): - from wan.diffusion_forcing import query_model_def - else: - from wan.any2video import query_model_def - elif model_family == "hunyuan": - from hyvideo.hunyuan import query_model_def - elif model_family == "ltxv": - from ltx_video.ltxv import query_model_def - elif model_family == "flux": - from flux.flux_main import query_model_def - else: - raise Exception(f"Unknown / unsupported model type {model_type}") - return query_model_def def init_model_def(model_type, model_def): - query_handler = get_model_query_handler(model_type) - default_model_def = query_handler(model_type, model_def) + base_model_type = get_base_model_type(model_type) + family_handler = model_types_handlers.get(base_model_type, None) + if family_handler is None: + raise Exception(f"Unknown model type {model_type}") + default_model_def = family_handler.query_model_def(base_model_type, model_def) if default_model_def is None: return model_def default_model_def.update(model_def) return default_model_def @@ -2282,7 +2182,6 @@ if len(args.vae_config) > 0: vae_config = int(args.vae_config) reload_needed = False -default_ui = server_config.get("default_ui", "t2v") save_path = server_config.get("save_path", os.path.join(os.getcwd(), "gradio_outputs")) preload_model_policy = server_config.get("preload_model_policy", []) @@ -2331,7 +2230,7 @@ def save_model(model, model_type, dtype, config_file, submodel_no = 1): model_filename = os.path.basename(url) break if model_filename is None: - print(f"No target filename mentioned in {url_key}") + print(f"No target filename with bf16 or fp16 in its name is mentioned in {url_key}") return if not os.path.isfile(model_filename): offload.save_model(model, os.path.join("ckpts",model_filename), config_file_path=config_file) @@ -2395,27 +2294,6 @@ def get_loras_preprocessor(transformer, model_type): return preprocessor_wrapper -def get_wan_text_encoder_filename(text_encoder_quantization): - text_encoder_filename = "ckpts/umt5-xxl/models_t5_umt5-xxl-enc-bf16.safetensors" - if text_encoder_quantization =="int8": - text_encoder_filename = text_encoder_filename.replace("bf16", "quanto_int8") - return text_encoder_filename - -def get_ltxv_text_encoder_filename(text_encoder_quantization): - text_encoder_filename = "ckpts/T5_xxl_1.1/T5_xxl_1.1_enc_bf16.safetensors" - if text_encoder_quantization =="int8": - text_encoder_filename = text_encoder_filename.replace("bf16", "quanto_bf16_int8") - return text_encoder_filename - -def get_hunyuan_text_encoder_filename(text_encoder_quantization): - if text_encoder_quantization =="int8": - text_encoder_filename = "ckpts/llava-llama-3-8b/llava-llama-3-8b-v1_1_vlm_quanto_int8.safetensors" - else: - text_encoder_filename = "ckpts/llava-llama-3-8b/llava-llama-3-8b-v1_1_vlm_fp16.safetensors" - - return text_encoder_filename - - def process_files_def(repoId, sourceFolderList, fileList): targetRoot = "ckpts/" for sourceFolder, files in zip(sourceFolderList,fileList ): @@ -2440,7 +2318,7 @@ def download_mmaudio(): } process_files_def(**enhancer_def) -def download_models(model_filename, model_type, submodel_no = 1): +def download_models(model_filename, model_type, module_type = None, submodel_no = 1): def computeList(filename): if filename == None: return [] @@ -2451,7 +2329,7 @@ def download_models(model_filename, model_type, submodel_no = 1): from urllib.request import urlretrieve - from wan.utils.utils import create_progress_hook + from shared.utils.utils import create_progress_hook shared_def = { "repoId" : "DeepBeepMeep/Wan2.1", @@ -2495,17 +2373,21 @@ def download_models(model_filename, model_type, submodel_no = 1): else: urlretrieve(url,filename, create_progress_hook(filename)) - model_family = get_model_family(model_type) + base_model_type = get_base_model_type(model_type) model_def = get_model_def(model_type) source = model_def.get("source", None) - - + model_type_handler = model_types_handlers[base_model_type] + key_name = "URLs" if submodel_no <= 1 else f"URLs{submodel_no}" if source is not None: model_filename = None - elif not model_type in modules_files: - if not os.path.isfile(model_filename ): + elif module_type is not None: + modules_files = model_type_handler.query_modules_files() if hasattr(model_type_handler, "query_modules_files") else {} + if module_type not in modules_files: + raise Exception(f"Unknown module {model_type} for model type {model_type}") + else: + if not os.path.isfile(model_filename): URLs = get_model_recursive_prop(model_type, key_name, return_list= False) if isinstance(URLs, str): raise Exception("Missing model " + URLs) @@ -2547,55 +2429,7 @@ def download_models(model_filename, model_type, submodel_no = 1): except Exception as e: if os.path.isfile(filename): os.remove(filename) raise Exception(f"Lora URL '{url}' is invalid: {str(e)}'") - - if model_family == "wan": - text_encoder_filename = get_wan_text_encoder_filename(text_encoder_quantization) - model_files = { - "repoId" : "DeepBeepMeep/Wan2.1", - "sourceFolderList" : ["xlm-roberta-large", "umt5-xxl", "" ], - "fileList" : [ [ "models_clip_open-clip-xlm-roberta-large-vit-huge-14-bf16.safetensors", "sentencepiece.bpe.model", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json"], ["special_tokens_map.json", "spiece.model", "tokenizer.json", "tokenizer_config.json"] + computeList(text_encoder_filename) , ["Wan2.1_VAE.safetensors", "fantasy_proj_model.safetensors" ] + computeList(model_filename) ] - } - elif model_family == "ltxv": - text_encoder_filename = get_ltxv_text_encoder_filename(text_encoder_quantization) - model_files = { - "repoId" : "DeepBeepMeep/LTX_Video", - "sourceFolderList" : ["T5_xxl_1.1", "" ], - "fileList" : [ ["added_tokens.json", "special_tokens_map.json", "spiece.model", "tokenizer_config.json"] + computeList(text_encoder_filename), ["ltxv_0.9.7_VAE.safetensors", "ltxv_0.9.7_spatial_upscaler.safetensors", "ltxv_scheduler.json"] + computeList(model_filename) ] - } - elif model_family == "hunyuan": - text_encoder_filename = get_hunyuan_text_encoder_filename(text_encoder_quantization) - model_files = { - "repoId" : "DeepBeepMeep/HunyuanVideo", - "sourceFolderList" : [ "llava-llama-3-8b", "clip_vit_large_patch14", "whisper-tiny" , "det_align", "" ], - "fileList" :[ ["config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json", "preprocessor_config.json"] + computeList(text_encoder_filename) , - ["config.json", "merges.txt", "model.safetensors", "preprocessor_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json", "vocab.json"], - ["config.json", "model.safetensors", "preprocessor_config.json", "special_tokens_map.json", "tokenizer_config.json"], - ["detface.pt"], - [ "hunyuan_video_720_quanto_int8_map.json", "hunyuan_video_custom_VAE_fp32.safetensors", "hunyuan_video_custom_VAE_config.json", "hunyuan_video_VAE_fp32.safetensors", "hunyuan_video_VAE_config.json" , "hunyuan_video_720_quanto_int8_map.json" ] + computeList(model_filename) - ] - } - elif model_family == "flux": - text_encoder_filename = get_ltxv_text_encoder_filename(text_encoder_quantization) - model_files = [ - { - "repoId" : "DeepBeepMeep/Flux", - "sourceFolderList" : [""], - "fileList" : [ ["flux_vae.safetensors"] ] - }, - { - "repoId" : "DeepBeepMeep/LTX_Video", - "sourceFolderList" : ["T5_xxl_1.1"], - "fileList" : [ ["added_tokens.json", "special_tokens_map.json", "spiece.model", "tokenizer_config.json"] + computeList(text_encoder_filename) ] - }, - { - "repoId" : "DeepBeepMeep/HunyuanVideo", - "sourceFolderList" : [ "clip_vit_large_patch14", ], - "fileList" :[ - ["config.json", "merges.txt", "model.safetensors", "preprocessor_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json", "vocab.json"], - ] - } - ] - + model_files = model_type_handler.query_model_files(computeList, base_model_type, model_filename, text_encoder_quantization) if not isinstance(model_files, list): model_files = [model_files] for one_repo in model_files: process_files_def(**one_repo) @@ -2690,116 +2524,6 @@ def setup_loras(model_type, transformer, lora_dir, lora_preselected_preset, spl print(error[:200]) return loras, loras_names, loras_presets, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset - -def load_wan_model(model_filename, model_type, base_model_type, model_def, quantizeTransformer = False, dtype = torch.bfloat16, VAE_dtype = torch.float32, mixed_precision_transformer = False, save_quantized= False): - if test_class_i2v(base_model_type): - cfg = WAN_CONFIGS['i2v-14B'] - else: - cfg = WAN_CONFIGS['t2v-14B'] - # cfg = WAN_CONFIGS['t2v-1.3B'] - if base_model_type in ("sky_df_1.3B", "sky_df_14B"): - model_factory = wan.DTT2V - else: - model_factory = wan.WanAny2V - - wan_model = model_factory( - config=cfg, - checkpoint_dir="ckpts", - model_filename=model_filename, - model_type = model_type, - model_def = model_def, - base_model_type=base_model_type, - text_encoder_filename= get_wan_text_encoder_filename(text_encoder_quantization), - quantizeTransformer = quantizeTransformer, - dtype = dtype, - VAE_dtype = VAE_dtype, - mixed_precision_transformer = mixed_precision_transformer, - save_quantized = save_quantized - ) - - pipe = {"transformer": wan_model.model, "text_encoder" : wan_model.text_encoder.model, "vae": wan_model.vae.model } - if hasattr(wan_model,"model2") and wan_model.model2 is not None: - pipe["transformer2"] = wan_model.model2 - if hasattr(wan_model, "clip"): - pipe["text_encoder_2"] = wan_model.clip.model - return wan_model, pipe - -def load_ltxv_model(model_filename, model_type, base_model_type, model_def, quantizeTransformer = False, dtype = torch.bfloat16, VAE_dtype = torch.float32, mixed_precision_transformer = False, save_quantized = False): - from ltx_video.ltxv import LTXV - - ltxv_model = LTXV( - model_filepath = model_filename, - text_encoder_filepath = get_ltxv_text_encoder_filename(text_encoder_quantization), - model_type = model_type, - base_model_type = base_model_type, - model_def = model_def, - dtype = dtype, - # quantizeTransformer = quantizeTransformer, - VAE_dtype = VAE_dtype, - mixed_precision_transformer = mixed_precision_transformer - ) - - pipeline = ltxv_model.pipeline - pipe = {"transformer" : pipeline.video_pipeline.transformer, "vae" : pipeline.vae, "text_encoder" : pipeline.video_pipeline.text_encoder, "latent_upsampler" : pipeline.latent_upsampler} - - return ltxv_model, pipe - - -def load_flux_model(model_filename, model_type, base_model_type, model_def, quantizeTransformer = False, dtype = torch.bfloat16, VAE_dtype = torch.float32, mixed_precision_transformer = False, save_quantized = False): - from flux.flux_main import model_factory - - flux_model = model_factory( - checkpoint_dir="ckpts", - model_filename=model_filename, - model_type = model_type, - model_def = model_def, - base_model_type=base_model_type, - text_encoder_filename= get_ltxv_text_encoder_filename(text_encoder_quantization), - quantizeTransformer = quantizeTransformer, - dtype = dtype, - VAE_dtype = VAE_dtype, - mixed_precision_transformer = mixed_precision_transformer, - save_quantized = save_quantized - ) - - pipe = { "transformer": flux_model.model, "vae" : flux_model.vae, "text_encoder" : flux_model.clip, "text_encoder_2" : flux_model.t5} - - return flux_model, pipe - -def load_hunyuan_model(model_filename, model_type = None, base_model_type = None, model_def = None, quantizeTransformer = False, dtype = torch.bfloat16, VAE_dtype = torch.float32, mixed_precision_transformer = False, save_quantized = False): - from hyvideo.hunyuan import HunyuanVideoSampler - - hunyuan_model = HunyuanVideoSampler.from_pretrained( - model_filepath = model_filename, - model_type = model_type, - base_model_type = base_model_type, - text_encoder_filepath = get_hunyuan_text_encoder_filename(text_encoder_quantization), - dtype = dtype, - quantizeTransformer = quantizeTransformer, - VAE_dtype = VAE_dtype, - mixed_precision_transformer = mixed_precision_transformer, - save_quantized = save_quantized - ) - - pipe = { "transformer" : hunyuan_model.model, "text_encoder" : hunyuan_model.text_encoder, "text_encoder_2" : hunyuan_model.text_encoder_2, "vae" : hunyuan_model.vae } - - if hunyuan_model.wav2vec != None: - pipe["wav2vec"] = hunyuan_model.wav2vec - - - # if hunyuan_model.align_instance != None: - # pipe["align_instance"] = hunyuan_model.align_instance.facedet.model - - - from hyvideo.modules.models import get_linear_split_map - - split_linear_modules_map = get_linear_split_map() - hunyuan_model.model.split_linear_modules_map = split_linear_modules_map - offload.split_linear_modules(hunyuan_model.model, split_linear_modules_map ) - - - return hunyuan_model, pipe - def get_transformer_model(model, submodel_no = 1): if submodel_no > 1: model_key = f"model{submodel_no}" @@ -2848,17 +2572,20 @@ def load_models(model_type): preload = server_config.get("preload_in_VRAM", 0) model_file_list = [model_filename] model_type_list = [model_type] + module_type_list = [None] model_submodel_no_list = [1] if model_filename2 != None: model_file_list += [model_filename2] model_type_list += [model_type] + module_type_list += [None] model_submodel_no_list += [2] for module_type in modules: - model_file_list.append(get_model_filename(module_type, transformer_quantization, transformer_dtype, is_module= True)) - model_type_list.append(module_type) + model_file_list.append(get_model_filename(model_type, transformer_quantization, transformer_dtype, module_type= module_type)) + model_type_list.append(model_type) + module_type_list.append(module_type) model_submodel_no_list.append(0) - for filename, file_model_type, submodel_no in zip(model_file_list, model_type_list, model_submodel_no_list): - download_models(filename, file_model_type, submodel_no) + for filename, file_model_type, file_module_type, submodel_no in zip(model_file_list, model_type_list, module_type_list, model_submodel_no_list): + download_models(filename, file_model_type, file_module_type, submodel_no) VAE_dtype = torch.float16 if server_config.get("vae_precision","16") == "16" else torch.float mixed_precision_transformer = server_config.get("mixed_precision","0") == "1" transformer_type = None @@ -2868,16 +2595,10 @@ def load_models(model_type): else: print(f"Loading Module '{filename}' ...") - if model_family == "wan" : - wan_model, pipe = load_wan_model(model_file_list, model_type, base_model_type, model_def, quantizeTransformer = quantizeTransformer, dtype = transformer_dtype, VAE_dtype = VAE_dtype, mixed_precision_transformer = mixed_precision_transformer, save_quantized = save_quantized) - elif model_family == "ltxv": - wan_model, pipe = load_ltxv_model(model_file_list, model_type, base_model_type, model_def, quantizeTransformer = quantizeTransformer, dtype = transformer_dtype, VAE_dtype = VAE_dtype, mixed_precision_transformer = mixed_precision_transformer, save_quantized = save_quantized) - elif model_family == "flux": - wan_model, pipe = load_flux_model(model_file_list, model_type, base_model_type, model_def, quantizeTransformer = quantizeTransformer, dtype = transformer_dtype, VAE_dtype = VAE_dtype, mixed_precision_transformer = mixed_precision_transformer, save_quantized = save_quantized) - elif model_family == "hunyuan": - wan_model, pipe = load_hunyuan_model(model_file_list, model_type, base_model_type, model_def, quantizeTransformer = quantizeTransformer, dtype = transformer_dtype, VAE_dtype = VAE_dtype, mixed_precision_transformer = mixed_precision_transformer, save_quantized = save_quantized) - else: - raise Exception(f"Model '{model_filename}' not supported.") + wan_model, pipe = model_types_handlers[base_model_type].load_model( + model_file_list, model_type, base_model_type, model_def, quantizeTransformer = quantizeTransformer, text_encoder_quantization = text_encoder_quantization, + dtype = transformer_dtype, VAE_dtype = VAE_dtype, mixed_precision_transformer = mixed_precision_transformer, save_quantized = save_quantized) + kwargs = { "extraModelsToQuantize": None } loras_transformer = ["transformer"] if profile in (2, 4, 5): @@ -3387,6 +3108,7 @@ def select_video(state, input_file_list, event_data: gr.EventData): + [ v for s,v in map_audio_prompt.items() if all_letters(video_audio_prompt_type,s)] video_model_type = configs.get("model_type", "t2v") model_family = get_model_family(video_model_type) + model_def = get_model_def(video_model_type) video_other_prompts = ", ".join(video_other_prompts) video_resolution = configs.get("resolution", "") + f" (real: {width}x{height})" video_length = configs.get("video_length", 0) @@ -3405,8 +3127,8 @@ def select_video(state, input_file_list, event_data: gr.EventData): video_guidance_scale = configs.get("guidance_scale", None) video_guidance2_scale = configs.get("guidance2_scale", None) video_switch_threshold = configs.get("switch_threshold", 0) - video_embedded_guidance_scale = configs.get("embedded_guidance_scale ", None) - if model_family in ["hunyuan", "flux"]: + video_embedded_guidance_scale = configs.get("embedded_guidance_scale", None) + if model_def.get("embedded_guidance", False): video_guidance_scale = video_embedded_guidance_scale video_guidance_label = "Embedded Guidance Scale" else: @@ -3506,7 +3228,7 @@ def convert_image(image): return cast(Image, ImageOps.exif_transpose(image)) def get_resampled_video(video_in, start_frame, max_frames, target_fps, bridge='torch'): - from wan.utils.utils import resample + from shared.utils.utils import resample import decord decord.bridge.set_bridge(bridge) @@ -3610,7 +3332,7 @@ def process_images_multithread(image_processor, items, process_type, wrap_in_lis return results def preprocess_video_with_mask(input_video_path, input_mask_path, height, width, max_frames, start_frame=0, fit_canvas = False, target_fps = 16, block_size= 16, expand_scale = 2, process_type = "inpaint", process_type2 = None, to_bbox = False, RGB_Mask = False, negate_mask = False, process_outside_mask = None, inpaint_color = 127, outpainting_dims = None, proc_no = 1): - from wan.utils.utils import calculate_new_dimensions, get_outpainting_frame_location, get_outpainting_full_area_dimensions + from shared.utils.utils import calculate_new_dimensions, get_outpainting_frame_location, get_outpainting_full_area_dimensions def mask_to_xyxy_box(mask): rows, cols = np.where(mask == 255) @@ -3895,7 +3617,7 @@ def perform_temporal_upsampling(sample, previous_last_frame, temporal_upsampling def perform_spatial_upsampling(sample, spatial_upsampling): - from wan.utils.utils import resize_lanczos + from shared.utils.utils import resize_lanczos if spatial_upsampling == "lanczos1.5": scale = 1.5 else: @@ -3990,7 +3712,7 @@ def edit_video( seed = set_seed(seed) - from wan.utils.utils import get_video_info + from shared.utils.utils import get_video_info fps, width, height, frames_count = get_video_info(video_source) frames_count = min(frames_count, max_source_video_frames) sample = None @@ -4223,10 +3945,12 @@ def generate_video( temp_filenames_list.append(video_mask) image_mask = None + base_model_type = get_base_model_type(model_type) fit_canvas = server_config.get("fit_canvas", 0) + model_handler = get_model_handler(base_model_type) + block_size = model_handler.get_vae_block_size(base_model_type) if hasattr(model_handler, "get_vae_block_size") else 16 - if "P" in preload_model_policy and not "U" in preload_model_policy: while wan_model == None: time.sleep(1) @@ -4261,12 +3985,14 @@ def generate_video( offload.shared_state["_attention"] = attn device_mem_capacity = torch.cuda.get_device_properties(0).total_memory / 1048576 - VAE_tile_size = wan_model.vae.get_VAE_tile_size(vae_config, device_mem_capacity, server_config.get("vae_precision", "16") == "32") + if hasattr(wan_model.vae, "get_VAE_tile_size"): + VAE_tile_size = wan_model.vae.get_VAE_tile_size(vae_config, device_mem_capacity, server_config.get("vae_precision", "16") == "32") + else: + VAE_tile_size = None trans = get_transformer_model(wan_model) trans2 = get_transformer_model(wan_model, 2) audio_sampling_rate = 16000 - base_model_type = get_base_model_type(model_type) prompts = prompt.split("\n") prompts = [part for part in prompts if len(prompt)>0] @@ -4329,7 +4055,7 @@ def generate_video( flux = base_model_type in ["flux"] if "B" in audio_prompt_type or "X" in audio_prompt_type: - from wan.multitalk.multitalk import parse_speakers_locations + from models.wan.multitalk.multitalk import parse_speakers_locations speakers_bboxes, error = parse_speakers_locations(speakers_locations) else: speakers_bboxes = None @@ -4373,7 +4099,7 @@ def generate_video( for i, pos in enumerate(frames_positions_list): frames_to_inject[pos] = image_refs[i] if video_guide == None and video_source == None and not "L" in image_prompt_type and (nb_frames_positions > 0 or "K" in video_prompt_type) : - from wan.utils.utils import get_outpainting_full_area_dimensions + from shared.utils.utils import get_outpainting_full_area_dimensions w, h = image_refs[0].size if outpainting_dims != None: h, w = get_outpainting_full_area_dimensions(h,w, outpainting_dims) @@ -4384,7 +4110,7 @@ def generate_video( if remove_background_images_ref > 0: send_cmd("progress", [0, get_latest_status(state, "Removing Images References Background")]) os.environ["U2NET_HOME"] = os.path.join(os.getcwd(), "ckpts", "rembg") - from wan.utils.utils import resize_and_remove_background + from shared.utils.utils import resize_and_remove_background image_refs[nb_frames_positions:] = resize_and_remove_background(image_refs[nb_frames_positions:] , width, height, remove_background_images_ref > 0, any_background_ref, fit_into_canvas= not (vace or hunyuan_avatar or flux) ) # no fit for vace ref images as it is done later update_task_thumbnails(task, locals()) send_cmd("output") @@ -4441,7 +4167,7 @@ def generate_video( audio_scale = None audio_context_lens = None if (fantasy or multitalk or hunyuan_avatar or hunyuan_custom_audio) and audio_guide != None: - from wan.fantasytalking.infer import parse_audio + from models.wan.fantasytalking.infer import parse_audio import librosa duration = librosa.get_duration(path=audio_guide) combination_type = "add" @@ -4465,7 +4191,7 @@ def generate_video( # audio_proj_split_full, audio_context_lens_full = parse_audio(audio_guide, num_frames= max_source_video_frames, fps= fps, padded_frames_for_embeddings= (reuse_frames if reset_control_aligment else 0), device= processing_device ) audio_scale = 1.0 elif multitalk: - from wan.multitalk.multitalk import get_full_audio_embeddings + from models.wan.multitalk.multitalk import get_full_audio_embeddings # pad audio_proj_full if aligned to beginning of window to simulate source window overlap audio_proj_full, output_new_audio_data = get_full_audio_embeddings(audio_guide1 = audio_guide, audio_guide2= audio_guide2, combination_type= combination_type , num_frames= max_source_video_frames, sr= audio_sampling_rate, fps =fps, padded_frames_for_embeddings = (reuse_frames if reset_control_aligment else 0)) if output_new_audio_filepath is not None: output_new_audio_data = None @@ -4548,7 +4274,7 @@ def generate_video( if len(original_prompts) == 0 and not "T" in prompt_enhancer: pass else: - from wan.utils.utils import seed_everything + from shared.utils.utils import seed_everything seed_everything(seed) # for i, original_prompt in enumerate(original_prompts): prompts = generate_cinematic_prompt( @@ -4607,9 +4333,9 @@ def generate_video( image_end_tensor = torch.from_numpy(np.array(image_end_tensor).astype(np.float32)).div_(127.5).sub_(1.).movedim(-1, 0) else: if "L" in image_prompt_type: - from wan.utils.utils import get_video_frame + from shared.utils.utils import get_video_frame refresh_preview["video_source"] = get_video_frame(video_source, 0) - prefix_video = preprocess_video(width=width, height=height,video_in=video_source, max_frames= parsed_keep_frames_video_source , start_frame = 0, fit_canvas= sample_fit_canvas, target_fps = fps, block_size = 32 if ltxv else 16) + prefix_video = preprocess_video(width=width, height=height,video_in=video_source, max_frames= parsed_keep_frames_video_source , start_frame = 0, fit_canvas= sample_fit_canvas, target_fps = fps, block_size = block_size ) prefix_video = prefix_video.permute(3, 0, 1, 2) prefix_video = prefix_video.float().div_(127.5).sub_(1.) # c, f, h, w pre_video_guide = prefix_video[:, -reuse_frames:] @@ -4628,7 +4354,7 @@ def generate_video( if fantasy: audio_proj_split , audio_context_lens = parse_audio(audio_guide, start_frame = aligned_window_start_frame, num_frames= current_video_length, fps= fps, device= processing_device ) if multitalk: - from wan.multitalk.multitalk import get_window_audio_embeddings + from models.wan.multitalk.multitalk import get_window_audio_embeddings # special treatment for start frame pos when alignement to first frame requested as otherwise the start frame number will be negative due to overlapped frames (has been previously compensated later with padding) audio_proj_split = get_window_audio_embeddings(audio_proj_full, audio_start_idx= aligned_window_start_frame + (source_video_overlap_frames_count if reset_control_aligment else 0 ), clip_length = current_video_length) @@ -4643,7 +4369,7 @@ def generate_video( status_info = "Extracting " + processes_names[preprocess_type] send_cmd("progress", [0, get_latest_status(state, status_info)]) # start one frame ealier to faciliate latents merging later - src_video, _ = preprocess_video_with_mask(video_guide, video_mask, height=image_size[0], width = image_size[1], max_frames= len(keep_frames_parsed) + (0 if aligned_guide_start_frame == 0 else 1), start_frame = aligned_guide_start_frame - (0 if aligned_guide_start_frame == 0 else 1), fit_canvas = sample_fit_canvas, target_fps = fps, process_type = preprocess_type, inpaint_color = 0, proc_no =1, negate_mask = "N" in video_prompt_type, process_outside_mask = "inpaint" if "X" in video_prompt_type else "identity", block_size =32 ) + src_video, _ = preprocess_video_with_mask(video_guide, video_mask, height=image_size[0], width = image_size[1], max_frames= len(keep_frames_parsed) + (0 if aligned_guide_start_frame == 0 else 1), start_frame = aligned_guide_start_frame - (0 if aligned_guide_start_frame == 0 else 1), fit_canvas = sample_fit_canvas, target_fps = fps, process_type = preprocess_type, inpaint_color = 0, proc_no =1, negate_mask = "N" in video_prompt_type, process_outside_mask = "inpaint" if "X" in video_prompt_type else "identity", block_size =block_size ) if src_video != None: src_video = src_video[ :(len(src_video)-1)// latent_size * latent_size +1 ] refresh_preview["video_guide"] = Image.fromarray(src_video[0].cpu().numpy()) @@ -4938,7 +4664,7 @@ def generate_video( time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%Hh%Mm%Ss") save_prompt = original_prompts[0] - from wan.utils.utils import truncate_for_filesystem + from shared.utils.utils import truncate_for_filesystem extension = "jpg" if is_image else "mp4" if os.name == 'nt': @@ -5064,222 +4790,16 @@ def prepare_generate_video(state): else: return gr.Button(visible= False), gr.Button(visible= True), gr.Column(visible= True), gr.update(visible= False) -def generate_preview(latents): + +def generate_preview(model_type, latents): import einops - # thanks Comfyui for the rgb factors - model_family = get_model_family(transformer_type) - if model_family == "wan": - latent_channels = 16 - latent_dimensions = 3 - latent_rgb_factors = [ - [-0.1299, -0.1692, 0.2932], - [ 0.0671, 0.0406, 0.0442], - [ 0.3568, 0.2548, 0.1747], - [ 0.0372, 0.2344, 0.1420], - [ 0.0313, 0.0189, -0.0328], - [ 0.0296, -0.0956, -0.0665], - [-0.3477, -0.4059, -0.2925], - [ 0.0166, 0.1902, 0.1975], - [-0.0412, 0.0267, -0.1364], - [-0.1293, 0.0740, 0.1636], - [ 0.0680, 0.3019, 0.1128], - [ 0.0032, 0.0581, 0.0639], - [-0.1251, 0.0927, 0.1699], - [ 0.0060, -0.0633, 0.0005], - [ 0.3477, 0.2275, 0.2950], - [ 0.1984, 0.0913, 0.1861] - ] - - # credits for the rgb factors to ComfyUI ? - - latent_rgb_factors_bias = [-0.1835, -0.0868, -0.3360] - - # latent_rgb_factors_bias = [0.0259, -0.0192, -0.0761] - elif model_family =="flux": - scale_factor = 0.3611 - shift_factor = 0.1159 - latent_rgb_factors =[ - [-0.0346, 0.0244, 0.0681], - [ 0.0034, 0.0210, 0.0687], - [ 0.0275, -0.0668, -0.0433], - [-0.0174, 0.0160, 0.0617], - [ 0.0859, 0.0721, 0.0329], - [ 0.0004, 0.0383, 0.0115], - [ 0.0405, 0.0861, 0.0915], - [-0.0236, -0.0185, -0.0259], - [-0.0245, 0.0250, 0.1180], - [ 0.1008, 0.0755, -0.0421], - [-0.0515, 0.0201, 0.0011], - [ 0.0428, -0.0012, -0.0036], - [ 0.0817, 0.0765, 0.0749], - [-0.1264, -0.0522, -0.1103], - [-0.0280, -0.0881, -0.0499], - [-0.1262, -0.0982, -0.0778] - ] - latent_rgb_factors_bias = [-0.0329, -0.0718, -0.0851] - - elif model_family == "ltxv": - latent_channels = 128 - latent_dimensions = 3 - - latent_rgb_factors = [ - [ 1.1202e-02, -6.3815e-04, -1.0021e-02], - [ 8.6031e-02, 6.5813e-02, 9.5409e-04], - [-1.2576e-02, -7.5734e-03, -4.0528e-03], - [ 9.4063e-03, -2.1688e-03, 2.6093e-03], - [ 3.7636e-03, 1.2765e-02, 9.1548e-03], - [ 2.1024e-02, -5.2973e-03, 3.4373e-03], - [-8.8896e-03, -1.9703e-02, -1.8761e-02], - [-1.3160e-02, -1.0523e-02, 1.9709e-03], - [-1.5152e-03, -6.9891e-03, -7.5810e-03], - [-1.7247e-03, 4.6560e-04, -3.3839e-03], - [ 1.3617e-02, 4.7077e-03, -2.0045e-03], - [ 1.0256e-02, 7.7318e-03, 1.3948e-02], - [-1.6108e-02, -6.2151e-03, 1.1561e-03], - [ 7.3407e-03, 1.5628e-02, 4.4865e-04], - [ 9.5357e-04, -2.9518e-03, -1.4760e-02], - [ 1.9143e-02, 1.0868e-02, 1.2264e-02], - [ 4.4575e-03, 3.6682e-05, -6.8508e-03], - [-4.5681e-04, 3.2570e-03, 7.7929e-03], - [ 3.3902e-02, 3.3405e-02, 3.7454e-02], - [-2.3001e-02, -2.4877e-03, -3.1033e-03], - [ 5.0265e-02, 3.8841e-02, 3.3539e-02], - [-4.1018e-03, -1.1095e-03, 1.5859e-03], - [-1.2689e-01, -1.3107e-01, -2.1005e-01], - [ 2.6276e-02, 1.4189e-02, -3.5963e-03], - [-4.8679e-03, 8.8486e-03, 7.8029e-03], - [-1.6610e-03, -4.8597e-03, -5.2060e-03], - [-2.1010e-03, 2.3610e-03, 9.3796e-03], - [-2.2482e-02, -2.1305e-02, -1.5087e-02], - [-1.5753e-02, -1.0646e-02, -6.5083e-03], - [-4.6975e-03, 5.0288e-03, -6.7390e-03], - [ 1.1951e-02, 2.0712e-02, 1.6191e-02], - [-6.3704e-03, -8.4827e-03, -9.5483e-03], - [ 7.2610e-03, -9.9326e-03, -2.2978e-02], - [-9.1904e-04, 6.2882e-03, 9.5720e-03], - [-3.7178e-02, -3.7123e-02, -5.6713e-02], - [-1.3373e-01, -1.0720e-01, -5.3801e-02], - [-5.3702e-03, 8.1256e-03, 8.8397e-03], - [-1.5247e-01, -2.1437e-01, -2.1843e-01], - [ 3.1441e-02, 7.0335e-03, -9.7541e-03], - [ 2.1528e-03, -8.9817e-03, -2.1023e-02], - [ 3.8461e-03, -5.8957e-03, -1.5014e-02], - [-4.3470e-03, -1.2940e-02, -1.5972e-02], - [-5.4781e-03, -1.0842e-02, -3.0204e-03], - [-6.5347e-03, 3.0806e-03, -1.0163e-02], - [-5.0414e-03, -7.1503e-03, -8.9686e-04], - [-8.5851e-03, -2.4351e-03, 1.0674e-03], - [-9.0016e-03, -9.6493e-03, 1.5692e-03], - [ 5.0914e-03, 1.2099e-02, 1.9968e-02], - [ 1.3758e-02, 1.1669e-02, 8.1958e-03], - [-1.0518e-02, -1.1575e-02, -4.1307e-03], - [-2.8410e-02, -3.1266e-02, -2.2149e-02], - [ 2.9336e-03, 3.6511e-02, 1.8717e-02], - [-1.6703e-02, -1.6696e-02, -4.4529e-03], - [ 4.8818e-02, 4.0063e-02, 8.7410e-03], - [-1.5066e-02, -5.7328e-04, 2.9785e-03], - [-1.7613e-02, -8.1034e-03, 1.3086e-02], - [-9.2633e-03, 1.0803e-02, -6.3489e-03], - [ 3.0851e-03, 4.7750e-04, 1.2347e-02], - [-2.2785e-02, -2.3043e-02, -2.6005e-02], - [-2.4787e-02, -1.5389e-02, -2.2104e-02], - [-2.3572e-02, 1.0544e-03, 1.2361e-02], - [-7.8915e-03, -1.2271e-03, -6.0968e-03], - [-1.1478e-02, -1.2543e-03, 6.2679e-03], - [-5.4229e-02, 2.6644e-02, 6.3394e-03], - [ 4.4216e-03, -7.3338e-03, -1.0464e-02], - [-4.5013e-03, 1.6082e-03, 1.4420e-02], - [ 1.3673e-02, 8.8877e-03, 4.1253e-03], - [-1.0145e-02, 9.0072e-03, 1.5695e-02], - [-5.6234e-03, 1.1847e-03, 8.1261e-03], - [-3.7171e-03, -5.3538e-03, 1.2590e-03], - [ 2.9476e-02, 2.1424e-02, 3.0424e-02], - [-3.4925e-02, -2.4340e-02, -2.5316e-02], - [-3.4127e-02, -2.2406e-02, -1.0589e-02], - [-1.7342e-02, -1.3249e-02, -1.0719e-02], - [-2.1478e-03, -8.6051e-03, -2.9878e-03], - [ 1.2089e-03, -4.2391e-03, -6.8569e-03], - [ 9.0411e-04, -6.6886e-03, -6.7547e-05], - [ 1.6048e-02, -1.0057e-02, -2.8929e-02], - [ 1.2290e-03, 1.0163e-02, 1.8861e-02], - [ 1.7264e-02, 2.7257e-04, 1.3785e-02], - [-1.3482e-02, -3.6427e-03, 6.7481e-04], - [ 4.6782e-03, -5.2423e-03, 2.4467e-03], - [-5.9113e-03, -6.2244e-03, -1.8162e-03], - [ 1.5496e-02, 1.4582e-02, 1.9514e-03], - [ 7.4958e-03, 1.5886e-03, -8.2305e-03], - [ 1.9086e-02, 1.6360e-03, -3.9674e-03], - [-5.7021e-03, -2.7307e-03, -4.1066e-03], - [ 1.7450e-03, 1.4602e-02, 2.5794e-02], - [-8.2788e-04, 2.2902e-03, 4.5161e-03], - [ 1.1632e-02, 8.9193e-03, -7.2813e-03], - [ 7.5721e-03, 2.6784e-03, 1.1393e-02], - [ 5.1939e-03, 3.6903e-03, 1.4049e-02], - [-1.8383e-02, -2.2529e-02, -2.4477e-02], - [ 5.8842e-04, -5.7874e-03, -1.4770e-02], - [-1.6125e-02, -8.6101e-03, -1.4533e-02], - [ 2.0540e-02, 2.0729e-02, 6.4338e-03], - [ 3.3587e-03, -1.1226e-02, -1.6444e-02], - [-1.4742e-03, -1.0489e-02, 1.7097e-03], - [ 2.8130e-02, 2.3546e-02, 3.2791e-02], - [-1.8532e-02, -1.2842e-02, -8.7756e-03], - [-8.0533e-03, -1.0771e-02, -1.7536e-02], - [-3.9009e-03, 1.6150e-02, 3.3359e-02], - [-7.4554e-03, -1.4154e-02, -6.1910e-03], - [ 3.4734e-03, -1.1370e-02, -1.0581e-02], - [ 1.1476e-02, 3.9281e-03, 2.8231e-03], - [ 7.1639e-03, -1.4741e-03, -3.8066e-03], - [ 2.2250e-03, -8.7552e-03, -9.5719e-03], - [ 2.4146e-02, 2.1696e-02, 2.8056e-02], - [-5.4365e-03, -2.4291e-02, -1.7802e-02], - [ 7.4263e-03, 1.0510e-02, 1.2705e-02], - [ 6.2669e-03, 6.2658e-03, 1.9211e-02], - [ 1.6378e-02, 9.4933e-03, 6.6971e-03], - [ 1.7173e-02, 2.3601e-02, 2.3296e-02], - [-1.4568e-02, -9.8279e-03, -1.1556e-02], - [ 1.4431e-02, 1.4430e-02, 6.6362e-03], - [-6.8230e-03, 1.8863e-02, 1.4555e-02], - [ 6.1156e-03, 3.4700e-03, -2.6662e-03], - [-2.6983e-03, -5.9402e-03, -9.2276e-03], - [ 1.0235e-02, 7.4173e-03, -7.6243e-03], - [-1.3255e-02, 1.9322e-02, -9.2153e-04], - [ 2.4222e-03, -4.8039e-03, -1.5759e-02], - [ 2.6244e-02, 2.5951e-02, 2.0249e-02], - [ 1.5711e-02, 1.8498e-02, 2.7407e-03], - [-2.1714e-03, 4.7214e-03, -2.2443e-02], - [-7.4747e-03, 7.4166e-03, 1.4430e-02], - [-8.3906e-03, -7.9776e-03, 9.7927e-03], - [ 3.8321e-02, 9.6622e-03, -1.9268e-02], - [-1.4605e-02, -6.7032e-03, 3.9675e-03] - ] - latent_rgb_factors_bias = [-0.0571, -0.1657, -0.2512] - - elif model_family == "hunyuan": - latent_channels = 16 - latent_dimensions = 3 - scale_factor = 0.476986 - latent_rgb_factors = [ - [-0.0395, -0.0331, 0.0445], - [ 0.0696, 0.0795, 0.0518], - [ 0.0135, -0.0945, -0.0282], - [ 0.0108, -0.0250, -0.0765], - [-0.0209, 0.0032, 0.0224], - [-0.0804, -0.0254, -0.0639], - [-0.0991, 0.0271, -0.0669], - [-0.0646, -0.0422, -0.0400], - [-0.0696, -0.0595, -0.0894], - [-0.0799, -0.0208, -0.0375], - [ 0.1166, 0.1627, 0.0962], - [ 0.1165, 0.0432, 0.0407], - [-0.2315, -0.1920, -0.1355], - [-0.0270, 0.0401, -0.0821], - [-0.0616, -0.0997, -0.0727], - [ 0.0249, -0.0469, -0.1703] - ] - - latent_rgb_factors_bias = [ 0.0259, -0.0192, -0.0761] + if latents is None: return None + model_handler = get_model_handler(model_type) + if hasattr(model_handler, "get_rgb_factors"): + latent_rgb_factors, latent_rgb_factors_bias = model_handler.get_rgb_factors(model_type) else: - raise Exception("preview not supported") + return None + if latent_rgb_factors is None: return None latents = latents.unsqueeze(0) nb_latents = latents.shape[2] latents_to_preview = 4 @@ -5310,7 +4830,7 @@ def generate_preview(latents): def process_tasks(state): - from wan.utils.thread_utils import AsyncStream, async_run + from shared.utils.thread_utils import AsyncStream, async_run gen = get_gen_info(state) queue = gen.get("queue", []) @@ -5393,7 +4913,7 @@ def process_tasks(state): # progress(*data) elif cmd == "preview": torch.cuda.current_stream().synchronize() - preview= None if data== None else generate_preview(data) + preview= None if data== None else generate_preview(params["model_type"], data) gen["preview"] = preview yield time.time() , gr.Text() else: @@ -5945,8 +5465,8 @@ def prepare_inputs_dict(target, inputs, model_type = None, model_filename = None if not get_model_family(model_type) == "wan" or diffusion_forcing: pop += ["sample_solver"] - if not (test_class_i2v(base_model_type) or diffusion_forcing or ltxv or recammaster or vace): - pop += ["image_prompt_type"] + # if not (test_class_i2v(base_model_type) or diffusion_forcing or ltxv or recammaster or vace): + # pop += ["image_prompt_type"] if any_audio_track(base_model_type) or server_config.get("mmaudio_enabled", 0) == 0: pop += ["MMAudio_setting", "MMAudio_prompt", "MMAudio_neg_prompt"] @@ -5962,8 +5482,8 @@ def prepare_inputs_dict(target, inputs, model_type = None, model_filename = None pop += ["model_mode"] if not vace and not phantom and not hunyuan_video_custom: - unsaved_params = ["keep_frames_video_guide", "video_prompt_type", "remove_background_images_ref", "mask_expand"] - if base_model_type in ["t2v"]: unsaved_params = unsaved_params[2:] + unsaved_params = ["keep_frames_video_guide", "remove_background_images_ref", "mask_expand"] #"video_prompt_type", + if base_model_type in ["t2v"]: unsaved_params = unsaved_params[1:] pop += unsaved_params if not vace: pop += ["frames_positions", "video_guide_outpainting", "control_net_weight", "control_net_weight2", "min_frames_if_references"] @@ -5977,24 +5497,39 @@ def prepare_inputs_dict(target, inputs, model_type = None, model_filename = None if not base_model_type in ["fantasy", "multitalk", "vace_multitalk_14B"]: pop += ["audio_guidance_scale", "speakers_locations"] - if not model_family in ["hunyuan", "flux"] or model_def.get("no_guidance", False): + if not model_def.get("embedded_guidance", False) or model_def.get("no_guidance", False): pop += ["embedded_guidance_scale"] - if not model_family in ["hunyuan", "wan"]: + if not model_def.get("skip_steps_cache", False) : pop += ["skip_steps_cache_type", "skip_steps_multiplier", "skip_steps_start_step_perc"] - if model_def.get("no_guidance", False) or ltxv or model_family in ["hunyuan", "flux"] : + if model_def.get("no_guidance", False) : pop += ["guidance_scale", "guidance2_scale", "switch_threshold", "audio_guidance_scale"] + + if not model_def.get("guidance_max_phases",1) >1: + pop += ["guidance2_scale", "switch_threshold"] + if model_def.get("image_outputs", False) or ltxv: pop += ["flow_shift"] - if model_def.get("no_negative_prompt", False) or model_family in ["flux"]: - pop += ["negative_prompt", "apg_switch", "cfg_star_switch", "cfg_zero_step", ] + if model_def.get("no_negative_prompt", False) : + pop += ["negative_prompt" ] + if not model_def.get("skip_layer_guidance", False): + pop += ["slg_switch", "slg_layers", "slg_start_perc", "slg_end_perc"] + + if not model_def.get("cfg_zero", False): + pop += [ "cfg_zero_step" ] + + if not model_def.get("cfg_star", False): + pop += ["cfg_star_switch" ] + + if not model_def.get("adaptive_projected_guidance", False): + pop += ["apg_switch"] if not model_family == "wan" or diffusion_forcing: - pop +=["NAG_scale", "NAG_tau", "NAG_alpha", "slg_switch", "slg_layers", "slg_start_perc", "slg_end_perc" ] + pop +=["NAG_scale", "NAG_tau", "NAG_alpha" ] for k in pop: if k in inputs: inputs.pop(k) @@ -6692,11 +6227,13 @@ def refresh_video_guide_outpainting_row(video_guide_outpainting_checkbox, video_ return gr.update(visible=video_guide_outpainting_checkbox), video_guide_outpainting custom_resolutions = None -def get_resolution_choices(current_resolution_choice): +def get_resolution_choices(current_resolution_choice, model_resolutions= None): global custom_resolutions resolution_file = "resolutions.json" - if custom_resolutions == None and os.path.isfile(resolution_file) : + if model_resolutions is not None: + resolution_choices = model_resolutions + elif custom_resolutions == None and os.path.isfile(resolution_file) : with open(resolution_file, 'r', encoding='utf-8') as f: try: resolution_choices = json.load(f) @@ -6756,9 +6293,13 @@ def get_resolution_choices(current_resolution_choice): if current_resolution_choice == res: found = True break - if not found: - resolution_choices.append( (current_resolution_choice, current_resolution_choice )) - return resolution_choices + if not found: + if model_resolutions is None: + resolution_choices.append( (current_resolution_choice, current_resolution_choice )) + else: + current_resolution_choice = resolution_choices[0][1] + + return resolution_choices, current_resolution_choice group_thresholds = { "360p": 320 * 640, @@ -6795,7 +6336,10 @@ def group_resolutions(resolutions, selected_resolution): return available_groups, selected_group_resolutions, selected_group def change_resolution_group(state, selected_group): - resolution_choices = get_resolution_choices(None) + model_type = state["model_type"] + model_def = get_model_def(model_type) + model_resolutions = model_def.get("resolutions", None) + resolution_choices, _ = get_resolution_choices(None, model_resolutions) group_resolution_choices = [ resolution for resolution in resolution_choices if categorize_resolution(resolution[1]) == selected_group ] last_resolution_per_group = state["last_resolution_per_group"] @@ -6929,7 +6473,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ltxv = "ltxv" in model_filename lock_inference_steps = model_def.get("lock_inference_steps", False) model_reference_image = model_def.get("reference_image", False) - no_steps_skipping = model_def.get("no_steps_skipping", False) + any_steps_skipping = model_def.get("skip_steps_cache", False) recammaster = base_model_type in ["recam_1.3B"] vace = test_vace_module(base_model_type) phantom = base_model_type in ["phantom_1.3B", "phantom_14B"] @@ -6947,7 +6491,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non sliding_window_enabled = test_any_sliding_window(model_type) multi_prompts_gen_type_value = ui_defaults.get("multi_prompts_gen_type_value",0) prompt_label, wizard_prompt_label = get_prompt_labels(multi_prompts_gen_type_value, image_outputs) - any_video_source = True + any_video_source = False fps = get_model_fps(base_model_type) image_prompt_type_value = "" video_prompt_type_value = "" @@ -6955,6 +6499,8 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non any_end_image = False any_reference_image = False v2i_switch_supported = (vace or t2v) and not image_outputs + ti2v_2_2 = base_model_type in ["ti2v_2_2"] + image_mode_value = ui_defaults.get("image_mode", 1 if image_outputs else 0 ) if not v2i_switch_supported and not image_outputs: image_mode_value = 0 @@ -6969,7 +6515,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non pass - with gr.Column(visible= test_class_i2v(model_type) or hunyuan_i2v or diffusion_forcing or ltxv or recammaster or vace) as image_prompt_column: + with gr.Column(visible= test_class_i2v(model_type) or hunyuan_i2v or diffusion_forcing or ltxv or recammaster or vace or ti2v_2_2) as image_prompt_column: if vace: image_prompt_type_value= ui_defaults.get("image_prompt_type","") image_prompt_type_value = "" if image_prompt_type_value == "S" else image_prompt_type_value @@ -6980,14 +6526,17 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non video_source = gr.Video(label= "Video Source", visible = "V" in image_prompt_type_value, value= ui_defaults.get("video_source", None)) model_mode = gr.Dropdown(visible = False) keep_frames_video_source = gr.Text(value=ui_defaults.get("keep_frames_video_source","") , visible= len(filter_letters(image_prompt_type_value, "VLG"))>0 , scale = 2, label= "Truncate Video beyond this number of resampled Frames (empty=Keep All, negative truncates from End)" ) + any_video_source = True - elif diffusion_forcing or ltxv: + elif diffusion_forcing or ltxv or ti2v_2_2: image_prompt_type_value= ui_defaults.get("image_prompt_type","T") # image_prompt_type = gr.Radio( [("Start Video with Image", "S"),("Start and End Video with Images", "SE"), ("Continue Video", "V"),("Text Prompt Only", "T")], value =image_prompt_type_value, label="Location", show_label= False, visible= True, scale= 3) image_prompt_type_choices = [("Text Prompt Only", "T"),("Start Video with Image", "S")] if ltxv: image_prompt_type_choices += [("Use both a Start and an End Image", "SE")] - image_prompt_type_choices += [("Continue Video", "V")] + if sliding_window_enabled: + any_video_source = True + image_prompt_type_choices += [("Continue Video", "V")] image_prompt_type = gr.Radio( image_prompt_type_choices, value =image_prompt_type_value, label="Location", show_label= False, visible= True , scale= 3) # image_start = gr.Image(label= "Image as a starting point for a new video", type ="pil",value= ui_defaults.get("image_start", None), visible= "S" in image_prompt_type_value ) @@ -6998,7 +6547,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non label="Images as ending points for new videos", type ="pil", #file_types= "image", columns=[3], rows=[1], object_fit="contain", height="auto", selected_index=0, interactive= True, visible="E" in image_prompt_type_value, value= ui_defaults.get("image_end", None)) video_source = gr.Video(label= "Video to Continue", visible= "V" in image_prompt_type_value, value= ui_defaults.get("video_source", None),) - if ltxv: + if not diffusion_forcing: model_mode = gr.Dropdown( choices=[ ], value=None, @@ -7045,8 +6594,9 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non image_prompt_type_choices = [("Start Video with Image", "S")] image_prompt_type_choices += [("Use both a Start and an End Image", "SE")] if not hunyuan_i2v: + any_video_source = True image_prompt_type_choices += [("Continue Video", "V")] - + image_prompt_type = gr.Radio( image_prompt_type_choices, value =image_prompt_type_value, label="Location", show_label= False, visible= not hunyuan_i2v, scale= 3) any_start_image = True any_end_image = True @@ -7061,13 +6611,11 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non video_source = gr.Video(value=None, visible=False) else: video_source = gr.Video(label= "Video to Continue", visible= "V" in image_prompt_type_value, value= ui_defaults.get("video_source", None),) - any_video_source = True else: image_prompt_type = gr.Radio(choices=[("", "")], value="") image_start = gr.Gallery(value=None) image_end = gr.Gallery(value=None) video_source = gr.Video(value=None, visible=False) - any_video_source = False model_mode = gr.Dropdown(value=None, visible=False) keep_frames_video_source = gr.Text(visible=False) @@ -7147,7 +6695,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ], value= filter_letters(video_prompt_type_value, "NA"), visible= "V" in video_prompt_type_value, - label="Area Processed", scale = 2 + label="Area Processed", scale = 2, show_label= True, ) elif ltxv: video_prompt_type_video_mask = gr.Dropdown( @@ -7160,7 +6708,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ], value= filter_letters(video_prompt_type_value, "XNA"), visible= "V" in video_prompt_type_value and not "U" in video_prompt_type_value, - label="Area Processed", scale = 2 + label="Area Processed", scale = 2, show_label= True, ) else: video_prompt_type_video_mask = gr.Dropdown( @@ -7179,7 +6727,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ], value= filter_letters(video_prompt_type_value, "XYZWNA"), visible= "V" in video_prompt_type_value and not "U" in video_prompt_type_value and not hunyuan_video_custom and not ltxv, - label="Area Processed", scale = 2 + label="Area Processed", scale = 2, show_label= True, ) if t2v: video_prompt_type_image_refs = gr.Dropdown(value="", label="Ref Image", choices=[""], visible =False) @@ -7193,7 +6741,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ], value=filter_letters(video_prompt_type_value, "KFI"), visible = True, - label="Reference Images", scale = 2 + label="Reference Images", show_label= True, scale = 2 ) @@ -7327,7 +6875,8 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non else: label = "Max Resolution (Pixels will be reallocated depending on the output width / height ratio)" current_resolution_choice = ui_defaults.get("resolution","832x480") if update_form or last_resolution is None else last_resolution - resolution_choices= get_resolution_choices(current_resolution_choice) + model_resolutions = model_def.get("resolutions", None) + resolution_choices, current_resolution_choice = get_resolution_choices(current_resolution_choice, model_resolutions) available_groups, selected_group_resolutions, selected_group = group_resolutions(resolution_choices, current_resolution_choice) resolution_group = gr.Dropdown( choices = available_groups, @@ -7366,14 +6915,15 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non with gr.Tab("General"): with gr.Column(): seed = gr.Slider(-1, 999999999, value=ui_defaults.get("seed",-1), step=1, label="Seed (-1 for random)") + any_embedded_guidance = model_def.get("embedded_guidance", False) with gr.Row(visible = not ltxv and not (no_guidance and image_outputs)) as guidance_row: - guidance_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("guidance_scale",5), step=0.5, label="Guidance (CFG)", visible=not (hunyuan_t2v or hunyuan_i2v or flux) and not no_guidance) + guidance_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("guidance_scale",5), step=0.5, label="Guidance (CFG)", visible=not (hunyuan_t2v or hunyuan_i2v or any_embedded_guidance) and not no_guidance) audio_guidance_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("audio_guidance_scale", 5 if fantasy else 4), step=0.5, label="Audio Guidance", visible=(fantasy or multitalk) and not no_guidance) - embedded_guidance_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("embedded_guidance", 2.5 if flux else 6.0), step=0.5, label="Embedded Guidance Scale", visible=(hunyuan_t2v or hunyuan_i2v or flux) and not no_guidance) + embedded_guidance_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("embedded_guidance", 2.5 if flux else 6.0), step=0.5, label="Embedded Guidance Scale", visible=(hunyuan_t2v or hunyuan_i2v or any_embedded_guidance) and not no_guidance) flow_shift = gr.Slider(1.0, 25.0, value=ui_defaults.get("flow_shift",3), step=0.1, label="Shift Scale", visible = not image_outputs) - with gr.Row(visible = not ltxv and not (no_guidance and image_outputs)) as guidance_row2: - guidance2_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("guidance2_scale",5), step=0.5, label="Guidance2 (CFG)", visible=not (hunyuan_t2v or hunyuan_i2v or flux) and not no_guidance) - switch_threshold = gr.Slider(0, 1000, value=ui_defaults.get("switch_threshold",0), step=1, label="Guidance / Model Switch Threshold", visible=not (hunyuan_t2v or hunyuan_i2v or flux) and not no_guidance) + with gr.Row(visible = model_def.get("guidance_max_phases",1) >1 and not (no_guidance and image_outputs)) as guidance_row2: + guidance2_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("guidance2_scale",5), step=0.5, label="Guidance2 (CFG)", visible=not (hunyuan_t2v or hunyuan_i2v or any_embedded_guidance) and not no_guidance) + switch_threshold = gr.Slider(0, 1000, value=ui_defaults.get("switch_threshold",0), step=1, label="Guidance / Model Switch Threshold", visible=not (hunyuan_t2v or hunyuan_i2v or any_embedded_guidance) and not no_guidance) with gr.Row(visible = get_model_family(model_type) == "wan" and not diffusion_forcing ) as sample_solver_row: sample_solver = gr.Dropdown( value=ui_defaults.get("sample_solver",""), @@ -7388,7 +6938,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non with gr.Row(visible = vace) as control_net_weights_row: control_net_weight = gr.Slider(0.0, 2.0, value=ui_defaults.get("control_net_weight",1), step=0.1, label="Control Net Weight #1", visible=vace) control_net_weight2 = gr.Slider(0.0, 2.0, value=ui_defaults.get("control_net_weight2",1), step=0.1, label="Control Net Weight #2", visible=vace) - negative_prompt = gr.Textbox(label="Negative Prompt (ignored if no Guidance that is if CFG = 1)", value=ui_defaults.get("negative_prompt", ""), visible = not (hunyuan_t2v or hunyuan_i2v or flux or no_negative_prompt) ) + negative_prompt = gr.Textbox(label="Negative Prompt (ignored if no Guidance that is if CFG = 1)", value=ui_defaults.get("negative_prompt", ""), visible = not (hunyuan_t2v or hunyuan_i2v or no_negative_prompt) ) with gr.Column(visible = vace or t2v or test_class_i2v(model_type)) as NAG_col: gr.Markdown("NAG enforces Negative Prompt even if no Guidance is set (CFG = 1), set NAG Scale to > 1 to enable it") with gr.Row(): @@ -7415,7 +6965,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non label="Activated Loras" ) loras_multipliers = gr.Textbox(label="Loras Multipliers (1.0 by default) separated by Space chars or CR, lines that start with # are ignored", value=launch_multis_str) - with gr.Tab("Steps Skipping", visible = not (ltxv or image_outputs) and not no_steps_skipping) as speed_tab: + with gr.Tab("Steps Skipping", visible = any_steps_skipping) as speed_tab: with gr.Column(): gr.Markdown("Tea Cache and Mag Cache accelerate the Video Generation by skipping intelligently some steps, the more steps are skipped the lower the quality of the video.") gr.Markdown("Steps Skipping consumes also VRAM. It is recommended not to skip at least the first 10% steps.") @@ -7516,9 +7066,13 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non gr.Markdown("Add Custom Soundtrack to Video") audio_source = gr.Audio(value= ui_defaults.get("audio_source", None), type="filepath", label="Soundtrack", show_download_button= True) + any_skip_layer_guidance = model_def.get("skip_layer_guidance", False) + any_cfg_zero = model_def.get("cfg_zero", False) + any_cfg_star = model_def.get("cfg_star", False) + any_apg = model_def.get("adaptive_projected_guidance", False) - with gr.Tab("Quality", visible = not (ltxv and no_negative_prompt or flux)) as quality_tab: - with gr.Column(visible = not (hunyuan_i2v or hunyuan_t2v or hunyuan_video_custom or hunyuan_video_avatar or ltxv) ) as skip_layer_guidance_row: + with gr.Tab("Quality", visible = vace and image_outputs or any_skip_layer_guidance or any_cfg_zero or any_cfg_star or any_apg ) as quality_tab: + with gr.Column(visible = any_skip_layer_guidance ) as skip_layer_guidance_row: gr.Markdown("Skip Layer Guidance (improves video quality, requires guidance > 1)") with gr.Row(): slg_switch = gr.Dropdown( @@ -7544,7 +7098,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non slg_start_perc = gr.Slider(0, 100, value=ui_defaults.get("slg_start_perc",10), step=1, label="Denoising Steps % start") slg_end_perc = gr.Slider(0, 100, value=ui_defaults.get("slg_end_perc",90), step=1, label="Denoising Steps % end") - with gr.Column(visible= not no_negative_prompt and (vace or multitalk or t2v or test_class_i2v(model_type) or ltxv) ) as apg_col: + with gr.Column(visible= any_apg ) as apg_col: gr.Markdown("Correct Progressive Color Saturation during long Video Generations") apg_switch = gr.Dropdown( choices=[ @@ -7557,7 +7111,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non label="Adaptive Projected Guidance (requires Guidance > 1) " ) - with gr.Column(visible = not ltxv) as cfg_free_guidance_col: + with gr.Column(visible = any_cfg_star) as cfg_free_guidance_col: gr.Markdown("Classifier-Free Guidance Zero Star, better adherence to Text Prompt") cfg_star_switch = gr.Dropdown( choices=[ @@ -7570,7 +7124,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non label="Classifier-Free Guidance Star (requires Guidance > 1)" ) with gr.Row(): - cfg_zero_step = gr.Slider(-1, 39, value=ui_defaults.get("cfg_zero_step",-1), step=1, label="CFG Zero below this Layer (Extra Process)", visible = not (hunyuan_i2v or hunyuan_t2v or hunyuan_video_avatar or hunyuan_i2v or hunyuan_video_custom )) + cfg_zero_step = gr.Slider(-1, 39, value=ui_defaults.get("cfg_zero_step",-1), step=1, label="CFG Zero below this Layer (Extra Process)", visible = any_cfg_zero) with gr.Column(visible = vace and image_outputs) as min_frames_if_references_col: gr.Markdown("If using Reference Images, generating a single Frame alone may not be sufficient to preserve Identity") @@ -9002,7 +8556,7 @@ def create_ui(): } """ if server_config.get("display_stats", 0) == 1: - from wan.utils.stats import SystemStatsApp + from shared.utils.stats import SystemStatsApp stats_app = SystemStatsApp() else: stats_app = None