edessa/EMG

1import os2import sys3 4import numpy as np5import torch6from diffusers import FlowMatchEulerDiscreteScheduler7from omegaconf import OmegaConf8from PIL import Image9from transformers import AutoTokenizer10 11current_file_path = os.path.abspath(__file__)12project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]13for project_root in project_roots:14    sys.path.insert(0, project_root) if project_root not in sys.path else None15 16from videox_fun.dist import set_multi_gpus_devices, shard_model17from videox_fun.models import (AutoencoderKLWan, AutoTokenizer, CLIPModel,18                               WanT5EncoderModel, WanTransformer3DModel)19from videox_fun.data.dataset_image_video import process_pose_file20from videox_fun.models.cache_utils import get_teacache_coefficients21from videox_fun.pipeline import WanFunControlPipeline, WanPipeline22from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,23                                               convert_weight_dtype_wrapper,24                                               replace_parameters_by_name)25from videox_fun.utils.lora_utils import merge_lora, unmerge_lora26from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent, get_image_latent,27                                    get_video_to_video_latent,28                                    save_videos_grid)29from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler30from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler31 32# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].33# model_full_load means that the entire model will be moved to the GPU.34# 35# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,36# and the transformer model has been quantized to float8, which can save more GPU memory. 37# 38# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.39# 40# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use, 41# and the transformer model has been quantized to float8, which can save more GPU memory. 42# 43# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use, 44# resulting in slower speeds but saving a large amount of GPU memory.45GPU_memory_mode     = "sequential_cpu_offload"46# Multi GPUs config47# Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used. 48# For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4.49# If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1.50ulysses_degree      = 151ring_degree         = 152# Use FSDP to save more GPU memory in multi gpus.53fsdp_dit            = False54fsdp_text_encoder   = True55# Compile will give a speedup in fixed resolution and need a little GPU memory. 56# The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload.57compile_dit         = False58 59# Support TeaCache.60enable_teacache     = True61# Recommended to be set between 0.05 and 0.30. A larger threshold can cache more steps, speeding up the inference process, 62# but it may cause slight differences between the generated content and the original content.63# # --------------------------------------------------------------------------------------------------- #64# | Model Name          | threshold | Model Name          | threshold | Model Name          | threshold |65# | Wan2.1-T2V-1.3B     | 0.05~0.10 | Wan2.1-T2V-14B      | 0.10~0.15 | Wan2.1-I2V-14B-720P | 0.20~0.30 |66# | Wan2.1-I2V-14B-480P | 0.20~0.25 | Wan2.1-Fun-*-1.3B-* | 0.05~0.10 | Wan2.1-Fun-*-14B-*  | 0.20~0.30 |67# # --------------------------------------------------------------------------------------------------- #68teacache_threshold  = 0.1069# The number of steps to skip TeaCache at the beginning of the inference process, which can70# reduce the impact of TeaCache on generated video quality.71num_skip_start_steps = 572# Whether to offload TeaCache tensors to cpu to save a little bit of GPU memory.73teacache_offload    = False74 75# Skip some cfg steps in inference76# Recommended to be set between 0.00 and 0.2577cfg_skip_ratio      = 078 79# Riflex config80enable_riflex       = False81# Index of intrinsic frequency82riflex_k            = 683 84# Config and model path85config_path         = "config/wan2.1/wan_civitai.yaml"86# model path87model_name          = "models/Diffusion_Transformer/Wan2.1-Fun-V1.1-1.3B-Control-Camera"88 89# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++"90sampler_name        = "Flow"91# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics. 92# Used when the sampler is in "Flow_Unipc", "Flow_DPM++".93# If you want to generate a 480p video, it is recommended to set the shift value to 3.0.94# If you want to generate a 720p video, it is recommended to set the shift value to 5.0.95shift               = 3 96 97# Load pretrained model if need98transformer_path    = None99vae_path            = None100lora_path           = None101 102# Other params103sample_size         = [480, 832]104video_length        = 81105fps                 = 16106 107# Use torch.float16 if GPU does not support torch.bfloat16108# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16109weight_dtype            = torch.bfloat16110control_video           = None111control_camera_txt      = "asset/Pan_Left.txt"112start_image             = "asset/7.png"113ref_image               = None114 115# 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性116# 在neg prompt中添加"安静，固定"等词语可以增加动态性。117prompt                  = "一个小女孩正在户外玩耍。她穿着一件蓝色的短袖上衣和粉色的短裤，头发扎成一个可爱的辫子。她的脚上没有穿鞋，显得非常自然和随意。她正用一把红色的小铲子在泥土里挖土，似乎在进行某种有趣的活动，可能是种花或是挖掘宝藏。地上有一根长长的水管，可能是用来浇水的。背景是一片草地和一些绿色植物，阳光明媚，整个场景充满了童趣和生机。小女孩专注的表情和认真的动作让人感受到她的快乐和好奇心。"118negative_prompt         = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"119 120# Using longer neg prompt such as "Blurring, mutation, deformation, distortion, dark and solid, comics, text subtitles, line art." can increase stability121# Adding words such as "quiet, solid" to the neg prompt can increase dynamism.122# prompt                  = "A young woman with beautiful, clear eyes and blonde hair stands in the forest, wearing a white dress and a crown. Her expression is serene, reminiscent of a movie star, with fair and youthful skin. Her brown long hair flows in the wind. The video quality is very high, with a clear view. High quality, masterpiece, best quality, high resolution, ultra-fine, fantastical."123# negative_prompt         = "Twisted body, limb deformities, text captions, comic, static, ugly, error, messy code."124guidance_scale          = 6.0125seed                    = 43126num_inference_steps     = 50127lora_weight             = 0.55128save_path               = "samples/wan-videos-fun-control"129 130device = set_multi_gpus_devices(ulysses_degree, ring_degree)131config = OmegaConf.load(config_path)132 133transformer = WanTransformer3DModel.from_pretrained(134    os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),135    transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),136    low_cpu_mem_usage=True,137    torch_dtype=weight_dtype,138)139 140if transformer_path is not None:141    print(f"From checkpoint: {transformer_path}")142    if transformer_path.endswith("safetensors"):143        from safetensors.torch import load_file, safe_open144        state_dict = load_file(transformer_path)145    else:146        state_dict = torch.load(transformer_path, map_location="cpu")147    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict148 149    m, u = transformer.load_state_dict(state_dict, strict=False)150    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")151 152# Get Vae153vae = AutoencoderKLWan.from_pretrained(154    os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')),155    additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),156).to(weight_dtype)157 158if vae_path is not None:159    print(f"From checkpoint: {vae_path}")160    if vae_path.endswith("safetensors"):161        from safetensors.torch import load_file, safe_open162        state_dict = load_file(vae_path)163    else:164        state_dict = torch.load(vae_path, map_location="cpu")165    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict166 167    m, u = vae.load_state_dict(state_dict, strict=False)168    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")169 170# Get Tokenizer171tokenizer = AutoTokenizer.from_pretrained(172    os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),173)174 175# Get Text encoder176text_encoder = WanT5EncoderModel.from_pretrained(177    os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),178    additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),179    low_cpu_mem_usage=True,180    torch_dtype=weight_dtype,181)182text_encoder = text_encoder.eval()183 184# Get Clip Image Encoder185clip_image_encoder = CLIPModel.from_pretrained(186    os.path.join(model_name, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),187).to(weight_dtype)188clip_image_encoder = clip_image_encoder.eval()189 190# Get Scheduler191Chosen_Scheduler = scheduler_dict = {192    "Flow": FlowMatchEulerDiscreteScheduler,193    "Flow_Unipc": FlowUniPCMultistepScheduler,194    "Flow_DPM++": FlowDPMSolverMultistepScheduler,195}[sampler_name]196if sampler_name == "Flow_Unipc" or sampler_name == "Flow_DPM++":197    config['scheduler_kwargs']['shift'] = 1198scheduler = Chosen_Scheduler(199    **filter_kwargs(Chosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))200)201 202# Get Pipeline203pipeline = WanFunControlPipeline(204    transformer=transformer,205    vae=vae,206    tokenizer=tokenizer,207    text_encoder=text_encoder,208    scheduler=scheduler,209    clip_image_encoder=clip_image_encoder210)211if ulysses_degree > 1 or ring_degree > 1:212    from functools import partial213    transformer.enable_multi_gpus_inference()214    if fsdp_dit:215        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)216        pipeline.transformer = shard_fn(pipeline.transformer)217        print("Add FSDP DIT")218    if fsdp_text_encoder:219        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)220        pipeline.text_encoder = shard_fn(pipeline.text_encoder)221        print("Add FSDP TEXT ENCODER")222 223if compile_dit:224    for i in range(len(pipeline.transformer.blocks)):225        pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])226    print("Add Compile")227 228if GPU_memory_mode == "sequential_cpu_offload":229    replace_parameters_by_name(transformer, ["modulation",], device=device)230    transformer.freqs = transformer.freqs.to(device=device)231    pipeline.enable_sequential_cpu_offload(device=device)232elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":233    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)234    convert_weight_dtype_wrapper(transformer, weight_dtype)235    pipeline.enable_model_cpu_offload(device=device)236elif GPU_memory_mode == "model_cpu_offload":237    pipeline.enable_model_cpu_offload(device=device)238elif GPU_memory_mode == "model_full_load_and_qfloat8":239    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)240    convert_weight_dtype_wrapper(transformer, weight_dtype)241    pipeline.to(device=device)242else:243    pipeline.to(device=device)244 245coefficients = get_teacache_coefficients(model_name) if enable_teacache else None246if coefficients is not None:247    print(f"Enable TeaCache with threshold {teacache_threshold} and skip the first {num_skip_start_steps} steps.")248    pipeline.transformer.enable_teacache(249        coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload250    )251 252if cfg_skip_ratio is not None:253    print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.")254    pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps)255 256generator = torch.Generator(device=device).manual_seed(seed)257 258if lora_path is not None:259    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)260 261with torch.no_grad():262    video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1263    latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1264 265    if enable_riflex:266        pipeline.transformer.enable_riflex(k = riflex_k, L_test = latent_frames)267    268    if ref_image is not None:269        clip_image = Image.open(ref_image).convert("RGB")270    elif start_image is not None:271        clip_image = Image.open(start_image).convert("RGB")272    else:273        clip_image = None274    275    if ref_image is not None:276        ref_image = get_image_latent(ref_image, sample_size=sample_size)277    278    if start_image is not None:279        start_image = get_image_latent(start_image, sample_size=sample_size)280 281    if control_camera_txt is not None:282        input_video, input_video_mask = None, None283        control_camera_video = process_pose_file(control_camera_txt, sample_size[1], sample_size[0])284        control_camera_video = control_camera_video[:video_length].permute([3, 0, 1, 2]).unsqueeze(0)285    else:286        input_video, input_video_mask, _, _ = get_video_to_video_latent(control_video, video_length=video_length, sample_size=sample_size, fps=fps, ref_image=None)287        control_camera_video = None288 289    sample = pipeline(290        prompt, 291        num_frames = video_length,292        negative_prompt = negative_prompt,293        height      = sample_size[0],294        width       = sample_size[1],295        generator   = generator,296        guidance_scale = guidance_scale,297        num_inference_steps = num_inference_steps,298 299        control_video = input_video,300        control_camera_video = control_camera_video,301        ref_image = ref_image,302        start_image = start_image,303        clip_image = clip_image,304        shift = shift,305    ).videos306 307if lora_path is not None:308    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)309 310def save_results():311    if not os.path.exists(save_path):312        os.makedirs(save_path, exist_ok=True)313 314    index = len([path for path in os.listdir(save_path)]) + 1315    prefix = str(index).zfill(8)316    if video_length == 1:317        video_path = os.path.join(save_path, prefix + ".png")318 319        image = sample[0, :, 0]320        image = image.transpose(0, 1).transpose(1, 2)321        image = (image * 255).numpy().astype(np.uint8)322        image = Image.fromarray(image)323        image.save(video_path)324    else:325        video_path = os.path.join(save_path, prefix + ".mp4")326        save_videos_grid(sample, video_path, fps=fps)327 328if ulysses_degree * ring_degree > 1:329    import torch.distributed as dist330    if dist.get_rank() == 0:331        save_results()332else:333    save_results()