Added Vac Contronet support

gradio_server.py

@@ -14,7 +14,7 @@ import gradio as gr
 import random
 import json
 import wan
-from wan.configs import MAX_AREA_CONFIGS, WAN_CONFIGS, SUPPORTED_SIZES
+from wan.configs import MAX_AREA_CONFIGS, WAN_CONFIGS, SUPPORTED_SIZES, VACE_SIZE_CONFIGS
 from wan.utils.utils import cache_video
 from wan.modules.attention import get_attention_modes, get_supported_attention_modes
 import torch
@@ -55,6 +55,11 @@ def format_time(seconds):
 def pil_to_base64_uri(pil_image, format="png", quality=75):
     if pil_image is None:
         return None
+
+    if isinstance(pil_image, str):
+        from wan.utils.utils import get_video_frame
+        pil_image = get_video_frame(pil_image, 0)
+
     buffer = io.BytesIO()
     try:
         img_to_save = pil_image
@@ -93,10 +98,11 @@ def process_prompt_and_add_tasks(
     loras_choices,
     loras_mult_choices,
     image_prompt_type,
-    image_to_continue,
-    image_to_end,
-    video_to_continue,
+    image_source1,
+    image_source2,
+    image_source3,
     max_frames,
+    remove_background_image_ref,
     temporal_upsampling,
     spatial_upsampling,    
     RIFLEx_setting,
@@ -127,9 +133,9 @@ def process_prompt_and_add_tasks(
         return
 
     file_model_needed = model_needed(image2video)
+    width, height = resolution.split("x")
+    width, height = int(width), int(height)
     if image2video:
-        width, height = resolution.split("x")
-        width, height = int(width), int(height)
 
         if "480p" in  file_model_needed and not "Fun" in file_model_needed and width * height > 848*480:
             gr.Info("You must use the 720P image to video model to generate videos with a resolution equivalent to 720P")
@@ -143,74 +149,94 @@ def process_prompt_and_add_tasks(
         gr.Info("You must use the 14B model to generate videos with a resolution equivalent to 720P")
         return
     
-    if image2video:
-        if image_to_continue == None or isinstance(image_to_continue, list) and len(image_to_continue) == 0:
+    if not image2video:
+        if "Vace" in file_model_needed and "1.3B" in file_model_needed :                
+            resolution_reformated = str(height) + "*" + str(width) 
+            if not resolution_reformated in VACE_SIZE_CONFIGS:
+                res = VACE_SIZE_CONFIGS.keys().join(" and ")
+                gr.Info(f"Video Resolution for Vace model is not supported. Only {res} resolutions are allowed.")
+                return
+
+        if not "I" in image_prompt_type:
+            image_source1 = None
+        if not "V" in image_prompt_type:
+            image_source2 = None
+        if not "M" in image_prompt_type:
+            image_source3 = None
+
+        if isinstance(image_source1, list):
+            image_source1 = [ convert_image(tup[0]) for tup in image_source1 ]
+
+            from wan.utils.utils import resize_and_remove_background
+            image_source1 = resize_and_remove_background(image_source1, width, height, remove_background_image_ref ==1)
+        
+        image_source1 = [ image_source1 ] * len(prompts)
+        image_source2 = [ image_source2 ] * len(prompts)
+        image_source3 = [ image_source3 ] * len(prompts)
+
+    else:
+        if image_source1 == None or isinstance(image_source1, list) and len(image_source1) == 0:
             return
         if image_prompt_type == 0:
-            image_to_end = None
-        if isinstance(image_to_continue, list):
-            image_to_continue = [ convert_image(tup[0]) for tup in image_to_continue ]
+            image_source2 = None
+        if isinstance(image_source1, list):
+            image_source1 = [ convert_image(tup[0]) for tup in image_source1 ]
         else:
-            image_to_continue = [convert_image(image_to_continue)]
-        if image_to_end != None:
-            if isinstance(image_to_end , list):
-                image_to_end  = [ convert_image(tup[0]) for tup in image_to_end  ]
+            image_source1 = [convert_image(image_source1)]
+        if image_source2 != None:
+            if isinstance(image_source2 , list):
+                image_source2  = [ convert_image(tup[0]) for tup in image_source2  ]
             else:
-                image_to_end  = [convert_image(image_to_end) ]
-            if len(image_to_continue) != len(image_to_end):
+                image_source2  = [convert_image(image_source2) ]
+            if len(image_source1) != len(image_source2):
                 gr.Info("The number of start and end images should be the same ")
                 return 
                             
         if multi_images_gen_type == 0:
             new_prompts = []
-            new_image_to_continue = []
-            new_image_to_end = []
-            for i in range(len(prompts) * len(image_to_continue) ):
+            new_image_source1 = []
+            new_image_source2 = []
+            for i in range(len(prompts) * len(image_source1) ):
                 new_prompts.append(  prompts[ i % len(prompts)] )
-                new_image_to_continue.append(image_to_continue[i // len(prompts)] )
-                if image_to_end != None:
-                    new_image_to_end.append(image_to_end[i // len(prompts)] )
+                new_image_source1.append(image_source1[i // len(prompts)] )
+                if image_source2 != None:
+                    new_image_source2.append(image_source2[i // len(prompts)] )
             prompts = new_prompts
-            image_to_continue = new_image_to_continue 
-            if image_to_end != None:
-                image_to_end = new_image_to_end 
+            image_source1 = new_image_source1 
+            if image_source2 != None:
+                image_source2 = new_image_source2 
         else:
-            if len(prompts) >= len(image_to_continue):
-                if len(prompts) % len(image_to_continue) !=0:
+            if len(prompts) >= len(image_source1):
+                if len(prompts) % len(image_source1) !=0:
                     raise gr.Error("If there are more text prompts than input images the number of text prompts should be dividable by the number of images")
-                rep = len(prompts) // len(image_to_continue)
-                new_image_to_continue = []
-                new_image_to_end = []
+                rep = len(prompts) // len(image_source1)
+                new_image_source1 = []
+                new_image_source2 = []
                 for i, _ in enumerate(prompts):
-                    new_image_to_continue.append(image_to_continue[i//rep] )
-                    if image_to_end != None:
-                        new_image_to_end.append(image_to_end[i//rep] )
-                image_to_continue = new_image_to_continue 
-                if image_to_end != None:
-                    image_to_end = new_image_to_end 
+                    new_image_source1.append(image_source1[i//rep] )
+                    if image_source2 != None:
+                        new_image_source2.append(image_source2[i//rep] )
+                image_source1 = new_image_source1 
+                if image_source2 != None:
+                    image_source2 = new_image_source2 
             else: 
-                if len(image_to_continue) % len(prompts)  !=0:
+                if len(image_source1) % len(prompts)  !=0:
                     raise gr.Error("If there are more input images than text prompts the number of images should be dividable by the number of text prompts")
-                rep = len(image_to_continue) // len(prompts)  
+                rep = len(image_source1) // len(prompts)  
                 new_prompts = []
-                for i, _ in enumerate(image_to_continue):
+                for i, _ in enumerate(image_source1):
                     new_prompts.append(  prompts[ i//rep] )
                 prompts = new_prompts
 
-        #     elif video_to_continue != None and len(video_to_continue) >0 :
-        #         input_image_or_video_path = video_to_continue
-        #         # pipeline.num_input_frames = max_frames
-        #         # pipeline.max_frames = max_frames
-        #     else:
-        #         return
-        # else:
-        #     input_image_or_video_path = None
-    if image_to_continue == None:
-        image_to_continue = [None] * len(prompts)
-    if image_to_end == None:
-        image_to_end = [None] * len(prompts)
-
-    for single_prompt, image_start, image_end in zip(prompts, image_to_continue, image_to_end) :
+  
+        if image_source1 == None:
+            image_source1 = [None] * len(prompts)
+        if image_source2 == None:
+            image_source2 = [None] * len(prompts)
+        if image_source3 == None:
+            image_source3 = [None] * len(prompts)
+
+    for single_prompt, image_source1, image_source2, image_source3 in zip(prompts, image_source1, image_source2, image_source3) :
         kwargs = {
             "prompt" : single_prompt,
             "negative_prompt" : negative_prompt,
@@ -228,10 +254,11 @@ def process_prompt_and_add_tasks(
             "loras_choices" : loras_choices,
             "loras_mult_choices" : loras_mult_choices,
             "image_prompt_type" : image_prompt_type,
-            "image_to_continue": image_start,
-            "image_to_end" : image_end,
-            "video_to_continue" : video_to_continue ,
+            "image_source1": image_source1,
+            "image_source2" : image_source2,
+            "image_source3" : image_source3 ,
             "max_frames" : max_frames,
+            "remove_background_image_ref" : remove_background_image_ref,
             "temporal_upsampling" : temporal_upsampling,
             "spatial_upsampling" : spatial_upsampling,
             "RIFLEx_setting" : RIFLEx_setting,
@@ -262,8 +289,9 @@ def add_video_task(**kwargs):
     queue = gen["queue"]
     task_id += 1
     current_task_id = task_id
-    start_image_data = kwargs["image_to_continue"]
-    end_image_data = kwargs["image_to_end"]
+    start_image_data = kwargs["image_source1"]
+    start_image_data = [start_image_data] if not isinstance(start_image_data, list) else start_image_data
+    end_image_data = kwargs["image_source2"]
 
     queue.append({
         "id": current_task_id,
@@ -275,7 +303,7 @@ def add_video_task(**kwargs):
         "prompt": kwargs["prompt"],
         "start_image_data": start_image_data,
         "end_image_data": end_image_data,
-        "start_image_data_base64":  pil_to_base64_uri(start_image_data, format="jpeg", quality=70),
+        "start_image_data_base64": [pil_to_base64_uri(img, format="jpeg", quality=70) for img in start_image_data],
         "end_image_data_base64": pil_to_base64_uri(end_image_data, format="jpeg", quality=70)
     })
     return update_queue_data(queue)
@@ -342,6 +370,7 @@ def get_queue_table(queue):
         full_prompt = item['prompt'].replace('"', '"')
         prompt_cell = f'<span title="{full_prompt}">{truncated_prompt}</span>'
         start_img_uri =item.get('start_image_data_base64')
+        start_img_uri = start_img_uri[0] if start_img_uri !=None else None
         end_img_uri = item.get('end_image_data_base64')
         thumbnail_size = "50px"
         num_steps = item.get('steps')
@@ -694,6 +723,9 @@ attention_modes_installed = get_attention_modes()
 attention_modes_supported = get_supported_attention_modes()
 args = _parse_args()
 args.flow_reverse = True
+processing_device = args.gpu
+if len(processing_device) == 0:
+    processing_device ="cuda"
 # torch.backends.cuda.matmul.allow_fp16_accumulation = True
 lock_ui_attention = False
 lock_ui_transformer = False
@@ -706,7 +738,7 @@ quantizeTransformer = args.quantize_transformer
 check_loras = args.check_loras ==1
 advanced = args.advanced
 
-transformer_choices_t2v=["ckpts/wan2.1_text2video_1.3B_bf16.safetensors", "ckpts/wan2.1_text2video_14B_bf16.safetensors", "ckpts/wan2.1_text2video_14B_quanto_int8.safetensors"]   
+transformer_choices_t2v=["ckpts/wan2.1_text2video_1.3B_bf16.safetensors", "ckpts/wan2.1_text2video_14B_bf16.safetensors", "ckpts/wan2.1_text2video_14B_quanto_int8.safetensors", "ckpts/wan2.1_Vace_1.3B_preview_bf16.safetensors"]   
 transformer_choices_i2v=["ckpts/wan2.1_image2video_480p_14B_bf16.safetensors", "ckpts/wan2.1_image2video_480p_14B_quanto_int8.safetensors", "ckpts/wan2.1_image2video_720p_14B_bf16.safetensors", "ckpts/wan2.1_image2video_720p_14B_quanto_int8.safetensors", "ckpts/wan2.1_Fun_InP_1.3B_bf16.safetensors", "ckpts/wan2.1_Fun_InP_14B_bf16.safetensors", "ckpts/wan2.1_Fun_InP_14B_quanto_int8.safetensors", ]
 text_encoder_choices = ["ckpts/models_t5_umt5-xxl-enc-bf16.safetensors", "ckpts/models_t5_umt5-xxl-enc-quanto_int8.safetensors"]
 
@@ -750,7 +782,7 @@ def get_default_settings(filename, i2v):
             "prompts": get_default_prompt(i2v),
             "resolution": "832x480",
             "video_length": 81,
-            "image_prompt_type" : 0, 
+            "image_prompt_type" : 0 if i2v else "", 
             "num_inference_steps": 30,
             "seed": -1,
             "repeat_generation": 1,
@@ -1149,6 +1181,9 @@ def get_model_name(model_filename):
     if "Fun" in model_filename:
         model_name = "Fun InP image2video"
         model_name += " 14B" if "14B" in model_filename else " 1.3B"
+    elif "Vace" in model_filename:
+        model_name = "Vace ControlNet text2video"
+        model_name += " 14B" if "14B" in model_filename else " 1.3B"
     elif "image" in model_filename:
         model_name = "Wan2.1 image2video"
         model_name += " 720p" if "720p" in model_filename else " 480p"
@@ -1353,22 +1388,22 @@ def refresh_gallery(state, msg):
         end_img_md = ""
         prompt =  task["prompt"]
 
-        if task.get('image2video'):
-            start_img_uri = task.get('start_image_data_base64')
-            end_img_uri = task.get('end_image_data_base64')
-            thumbnail_size = "100px"
-            if start_img_uri:
-                start_img_md = f'<img src="{start_img_uri}" alt="Start" style="max-width:{thumbnail_size}; max-height:{thumbnail_size}; display: block; margin: auto; object-fit: contain;" />'
-            if end_img_uri:
-                end_img_md = f'<img src="{end_img_uri}" alt="End" style="max-width:{thumbnail_size}; max-height:{thumbnail_size}; display: block; margin: auto; object-fit: contain;" />'
+        start_img_uri = task.get('start_image_data_base64')
+        start_img_uri = start_img_uri[0] if start_img_uri !=None else None
+        end_img_uri = task.get('end_image_data_base64')
+        thumbnail_size = "100px"
+        if start_img_uri:
+            start_img_md = f'<img src="{start_img_uri}" alt="Start" style="max-width:{thumbnail_size}; max-height:{thumbnail_size}; display: block; margin: auto; object-fit: contain;" />'
+        if end_img_uri:
+            end_img_md = f'<img src="{end_img_uri}" alt="End" style="max-width:{thumbnail_size}; max-height:{thumbnail_size}; display: block; margin: auto; object-fit: contain;" />'
         
         label = f"Prompt of Video being Generated"            
  
         html = "<STYLE> #PINFO, #PINFO  th, #PINFO td {border: 1px solid #CCCCCC;background-color:#FFFFFF;}</STYLE><TABLE WIDTH=100% ID=PINFO ><TR><TD width=100%>" + prompt + "</TD>"
         if start_img_md != "":
             html += "<TD>" + start_img_md +  "</TD>"
-            if end_img_md != "":
-                html += "<TD>" + end_img_md +  "</TD>" 
+        if end_img_md != "":
+            html += "<TD>" + end_img_md +  "</TD>" 
 
         html += "</TR></TABLE>" 
         html_output = gr.HTML(html, visible= True)
@@ -1419,24 +1454,26 @@ def expand_slist(slist, num_inference_steps ):
         new_slist.append(slist[ int(pos)])
         pos += inc
     return new_slist
-
 def convert_image(image):
-    from PIL import ExifTags
-
-    image = image.convert('RGB')
-    for orientation in ExifTags.TAGS.keys():
-        if ExifTags.TAGS[orientation]=='Orientation':
-            break            
-    exif = image.getexif()
-    if not orientation in exif:
-        return image
-    if exif[orientation] == 3:
-        image=image.rotate(180, expand=True)
-    elif exif[orientation] == 6:
-        image=image.rotate(270, expand=True)
-    elif exif[orientation] == 8:
-        image=image.rotate(90, expand=True)
-    return image
+
+    from PIL import ExifTags, ImageOps
+    from typing import cast
+
+    return cast(Image, ImageOps.exif_transpose(image))
+    # image = image.convert('RGB')
+    # for orientation in ExifTags.TAGS.keys():
+    #     if ExifTags.TAGS[orientation]=='Orientation':
+    #         break            
+    # exif = image.getexif()
+    #     return image
+    # if not orientation in exif:
+    # if exif[orientation] == 3:
+    #     image=image.rotate(180, expand=True)
+    # elif exif[orientation] == 6:
+    #     image=image.rotate(270, expand=True)
+    # elif exif[orientation] == 8:
+    #     image=image.rotate(90, expand=True)
+    # return image
 
 def generate_video(
     task_id,
@@ -1457,10 +1494,11 @@ def generate_video(
     loras_choices,
     loras_mult_choices,
     image_prompt_type,
-    image_to_continue,
-    image_to_end,
-    video_to_continue,
+    image_source1,
+    image_source2,
+    image_source3,
     max_frames,
+    remove_background_image_ref,
     temporal_upsampling,
     spatial_upsampling,
     RIFLEx_setting,
@@ -1507,7 +1545,6 @@ def generate_video(
         gr.Info(f"You have selected attention mode '{attention_mode}'. However it is not installed or supported on your system. You should either install it or switch to the default 'sdpa' attention.")
         return
 
-
     
     if not image2video:
         width, height = resolution.split("x")
@@ -1586,7 +1623,7 @@ def generate_video(
  
     enable_RIFLEx = RIFLEx_setting == 0 and video_length > (6* 16) or RIFLEx_setting == 1
     # VAE Tiling
-    device_mem_capacity = torch.cuda.get_device_properties(0).total_memory / 1048576
+    device_mem_capacity = torch.cuda.get_device_properties(None).total_memory / 1048576
 
     joint_pass = boost ==1 #and profile != 1 and profile != 3  
    # TeaCache   
@@ -1615,6 +1652,17 @@ def generate_video(
             else:
                     raise gr.Error("Teacache not supported for this model")
 
+    if "Vace" in model_filename:
+        resolution_reformated = str(height) + "*" + str(width) 
+        src_video, src_mask, src_ref_images = wan_model.prepare_source([image_source2],
+                                                                [image_source3],
+                                                                [image_source1],
+                                                                video_length, VACE_SIZE_CONFIGS[resolution_reformated], "cpu",
+                                                                trim_video=max_frames)
+    else:
+        src_video, src_mask, src_ref_images = None, None, None
+
+
     import random
     if seed == None or seed <0:
         seed = random.randint(0, 999999999)
@@ -1673,8 +1721,8 @@ def generate_video(
             if image2video:
                 samples = wan_model.generate(
                     prompt,
-                    image_to_continue,  
-                    image_to_end if image_to_end != None else None,  
+                    image_source1,  
+                    image_source2 if image_source2 != None else None,  
                     frame_num=(video_length // 4)* 4 + 1,
                     max_area=MAX_AREA_CONFIGS[resolution], 
                     shift=flow_shift,
@@ -1697,6 +1745,9 @@ def generate_video(
             else:
                 samples = wan_model.generate(
                     prompt,
+                    input_frames = src_video,
+                    input_ref_images=  src_ref_images,
+                    input_masks = src_mask,
                     frame_num=(video_length // 4)* 4 + 1,
                     size=(width, height),
                     shift=flow_shift,
@@ -1745,7 +1796,7 @@ def generate_video(
                 new_error = "The generation of the video has encountered an error: it is likely that you have unsufficient VRAM and you should therefore reduce the video resolution or its number of frames."
             else:
                 new_error =  gr.Error(f"The generation of the video has encountered an error, please check your terminal for more information. '{s}'")
-            tb = traceback.format_exc().split('\n')[:-2] 
+            tb = traceback.format_exc().split('\n')[:-1] 
             print('\n'.join(tb))
             raise gr.Error(new_error, print_exception= False)
 
@@ -1799,7 +1850,7 @@ def generate_video(
             
             if exp > 0: 
                 from rife.inference import temporal_interpolation
-                sample = temporal_interpolation( os.path.join("ckpts", "flownet.pkl"), sample, exp, device="cuda")
+                sample = temporal_interpolation( os.path.join("ckpts", "flownet.pkl"), sample, exp, device=processing_device)
                 fps = fps * 2**exp
 
             if len(spatial_upsampling) > 0:
@@ -1831,8 +1882,7 @@ def generate_video(
                 normalize=True,
                 value_range=(-1, 1))
     
-
-            configs = get_settings_dict(state, image2video, prompt, 0 if image_to_end == None else 1 , video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, loras_choices, 
+            configs = get_settings_dict(state, image2video, True, prompt, image_prompt_type, max_frames , remove_background_image_ref, video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, loras_choices, 
                   loras_mult_choices, tea_cache , tea_cache_start_step_perc, temporal_upsampling, spatial_upsampling, RIFLEx_setting, slg_switch, slg_layers, slg_start, slg_end, cfg_star_switch, cfg_zero_step)
 
             metadata_choice = server_config.get("metadata_choice","metadata")
@@ -2294,7 +2344,7 @@ def switch_advanced(state, new_advanced, lset_name):
         return  gr.Row(visible=new_advanced), gr.Row(visible=True), gr.Button(visible=True), gr.Row(visible= False), gr.Dropdown(choices=lset_choices, value= lset_name)
 
 
-def get_settings_dict(state, i2v, prompt, image_prompt_type, video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, loras_choices, 
+def get_settings_dict(state, i2v, image_metadata, prompt, image_prompt_type, max_frames, remove_background_image_ref, video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, loras_choices, 
                       loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, temporal_upsampling, spatial_upsampling, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc, cfg_star_switch, cfg_zero_step):
 
     loras = state["loras"]
@@ -2330,18 +2380,22 @@ def get_settings_dict(state, i2v, prompt, image_prompt_type, video_length, resol
         ui_settings["type"] = "Wan2.1GP by DeepBeepMeep - image2video"
         ui_settings["image_prompt_type"] = image_prompt_type
     else: 
+        if "Vace" in transformer_filename_t2v or not image_metadata:
+            ui_settings["image_prompt_type"] = image_prompt_type
+            ui_settings["max_frames"] = max_frames
+            ui_settings["remove_background_image_ref"] = remove_background_image_ref
         ui_settings["type"] = "Wan2.1GP by DeepBeepMeep - text2video"
 
     return ui_settings
 
-def save_settings(state, prompt, image_prompt_type, video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, loras_choices, 
+def save_settings(state, prompt, image_prompt_type, max_frames, remove_background_image_ref, video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, loras_choices, 
                       loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, temporal_upsampling, spatial_upsampling, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc, cfg_star_switch, cfg_zero_step):
 
     if state.get("validate_success",0) != 1:
         return
 
     image2video = state["image2video"]
-    ui_defaults = get_settings_dict(state, image2video, prompt, image_prompt_type, video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, loras_choices, 
+    ui_defaults = get_settings_dict(state, image2video, False, prompt, image_prompt_type, max_frames, remove_background_image_ref, video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, loras_choices, 
                       loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, temporal_upsampling, spatial_upsampling, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc, cfg_star_switch, cfg_zero_step)
 
     defaults_filename = get_settings_file_name(image2video)
@@ -2379,6 +2433,25 @@ def download_loras():
         writer.write(f"Loras downloaded on the {dt} at {time.time()} on the {time.time()}")
     return
 
+def refresh_i2v_image_prompt_type_radio(state, image_prompt_type_radio):
+    if args.multiple_images:
+        return gr.Gallery(visible = (image_prompt_type_radio == 1)  )
+    else:
+        return gr.Image(visible = (image_prompt_type_radio == 1)  )
+
+def refresh_t2v_image_prompt_type_radio(state, image_prompt_type_radio):
+    vace_model = "Vace" in state["image_input_type_model"] and not state["image2video"]    
+    return gr.Column(visible= vace_model), gr.Radio(value= image_prompt_type_radio), gr.Gallery(visible = "I" in image_prompt_type_radio), gr.Video(visible= "V" in image_prompt_type_radio),gr.Video(visible= "M" in image_prompt_type_radio ), gr.Text(visible= "V" in image_prompt_type_radio) , gr.Checkbox(visible= "I" in image_prompt_type_radio)
+
+def check_refresh_input_type(state):
+    if not state["image2video"]:
+        model_file_name = state["image_input_type_model"]
+        model_file_needed= model_needed(False)
+        if model_file_name != model_file_needed:
+            state["image_input_type_model"] = model_file_needed
+            return gr.Text(value= str(time.time()))
+    return gr.Text()
+
 def generate_video_tab(image2video=False):
     filename = transformer_filename_i2v if image2video else transformer_filename_t2v
     ui_defaults=  get_default_settings(filename, image2video)
@@ -2387,6 +2460,7 @@ def generate_video_tab(image2video=False):
 
     state_dict["advanced"] = advanced
     state_dict["loras_model"] = filename
+    state_dict["image_input_type_model"] = filename
     state_dict["image2video"] = image2video
     gen = dict()
     gen["queue"] = []
@@ -2461,31 +2535,51 @@ def generate_video_tab(image2video=False):
                         save_lset_btn = gr.Button("Save", size="sm", min_width= 1)
                         delete_lset_btn = gr.Button("Delete", size="sm", min_width= 1)
                         cancel_lset_btn = gr.Button("Don't do it !", size="sm", min_width= 1 , visible=False)  
-            video_to_continue = gr.Video(label= "Video to continue", visible= image2video and False) #######
-            image_prompt_type= ui_defaults.get("image_prompt_type",0)
-            image_prompt_type_radio = gr.Radio( [("Use only a Start Image", 0),("Use both a Start and an End Image", 1)], value =image_prompt_type, label="Location", show_label= False, scale= 3, visible=image2video)
-
-            if args.multiple_images:  
-                image_to_continue = gr.Gallery(
-                        label="Images as starting points for new videos", type ="pil", #file_types= "image", 
-                        columns=[3], rows=[1], object_fit="contain", height="auto", selected_index=0, interactive= True, visible=image2video)
-            else:
-                image_to_continue = gr.Image(label= "Image as a starting point for a new video", type ="pil", visible=image2video)
 
-            if args.multiple_images:  
-                image_to_end  = gr.Gallery(
-                        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=image_prompt_type==1)
-            else:
-                image_to_end = gr.Image(label= "Last Image for a new video", type ="pil", visible=image_prompt_type==1)
+            state = gr.State(state_dict)     
+            vace_model = "Vace" in filename and not image2video
+            trigger_refresh_input_type = gr.Text(interactive= False, visible= False)
+            with gr.Column(visible= image2video or vace_model) as image_prompt_column: 
+                if image2video:
+                    image_source3 = gr.Video(label= "Placeholder", visible= image2video and False)
+
+                    image_prompt_type= ui_defaults.get("image_prompt_type",0)
+                    image_prompt_type_radio = gr.Radio( [("Use only a Start Image", 0),("Use both a Start and an End Image", 1)], value =image_prompt_type, label="Location", show_label= False, scale= 3)
+
+                    if args.multiple_images:  
+                        image_source1 = gr.Gallery(
+                                label="Images as starting points for new videos", type ="pil", #file_types= "image", 
+                                columns=[3], rows=[1], object_fit="contain", height="auto", selected_index=0, interactive= True)
+                    else:
+                        image_source1 = gr.Image(label= "Image as a starting point for a new video", type ="pil")
 
-            def switch_image_prompt_type_radio(image_prompt_type_radio):
-                if args.multiple_images:
-                    return gr.Gallery(visible = (image_prompt_type_radio == 1)  )
+                    if args.multiple_images:  
+                        image_source2  = gr.Gallery(
+                                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=image_prompt_type==1)
+                    else:
+                        image_source2 = gr.Image(label= "Last Image for a new video", type ="pil", visible=image_prompt_type==1)
+
+
+                    image_prompt_type_radio.change(fn=refresh_i2v_image_prompt_type_radio, inputs=[state, image_prompt_type_radio], outputs=[image_source2]) 
+                    max_frames = gr.Slider(1, 100,step=1, visible = False) 
+                    remove_background_image_ref = gr.Text(visible = False)
                 else:
-                    return gr.Image(visible = (image_prompt_type_radio == 1)  )
+                    image_prompt_type= ui_defaults.get("image_prompt_type","I")
+                    image_prompt_type_radio = gr.Radio( [("Use Images Ref", "I"),("a Video", "V"), ("Images + a Video", "IV"), ("Video + Video Mask", "VM"), ("Images + Video + Mask", "IVM")], value =image_prompt_type, label="Location", show_label= False, scale= 3, visible = vace_model)
+                    image_source1 = gr.Gallery(
+                            label="Reference Images of Faces and / or Object to be found in the Video", type ="pil",  
+                            columns=[3], rows=[1], object_fit="contain", height="auto", selected_index=0, interactive= True, visible= "I" in image_prompt_type )
 
-            image_prompt_type_radio.change(fn=switch_image_prompt_type_radio, inputs=[image_prompt_type_radio], outputs=[image_to_end]) 
+                    image_source2 = gr.Video(label= "Reference Video", visible= "V" in image_prompt_type )
+                    with gr.Row():
+                        max_frames = gr.Slider(0, 100, value=ui_defaults.get("max_frames",0), step=1, label="Nb of frames in Reference Video to use in Video (0 for as many as possible)", visible= "V" in image_prompt_type, scale = 2 ) 
+                        remove_background_image_ref = gr.Checkbox(value=ui_defaults.get("remove_background_image_ref",1), label= "Remove Images Ref. Background", visible= "I" in image_prompt_type, scale =1 ) 
+
+                    image_source3 = gr.Video(label= "Video Mask (white pixels = Mask)", visible= "M" in image_prompt_type ) 
+
+    
+                    gr.on(triggers=[image_prompt_type_radio.change, trigger_refresh_input_type.change],  fn=refresh_t2v_image_prompt_type_radio, inputs=[state, image_prompt_type_radio], outputs=[image_prompt_column, image_prompt_type_radio, image_source1, image_source2, image_source3, max_frames, remove_background_image_ref]) 
 
 
             advanced_prompt = advanced
@@ -2518,7 +2612,6 @@ def generate_video_tab(image2video=False):
                 wizard_prompt = gr.Textbox(visible = not advanced_prompt, label="Prompts (each new line of prompt will generate a new video, # lines = comments)", value=default_wizard_prompt, lines=3)
                 wizard_prompt_activated_var = gr.Text(wizard_prompt_activated, visible= False)
                 wizard_variables_var = gr.Text(wizard_variables, visible = False)
-            state = gr.State(state_dict)     
             with gr.Row():
                 if image2video:
                     resolution = gr.Dropdown(
@@ -2555,8 +2648,6 @@ def generate_video_tab(image2video=False):
                     video_length = gr.Slider(5, 193, value=ui_defaults["video_length"], step=4, label="Number of frames (16 = 1s)")
                 with gr.Column():
                     num_inference_steps = gr.Slider(1, 100, value=ui_defaults["num_inference_steps"], step=1, label="Number of Inference Steps")
-            with gr.Row():
-                max_frames = gr.Slider(1, 100, value=9, step=1, label="Number of input frames to use for Video2World prediction", visible=image2video and False) #########
             show_advanced = gr.Checkbox(label="Advanced Mode", value=advanced)
             with gr.Row(visible=advanced) as advanced_row:
                 with gr.Column():
@@ -2605,7 +2696,7 @@ def generate_video_tab(image2video=False):
                         tea_cache_start_step_perc = gr.Slider(0, 100, value=ui_defaults["tea_cache_start_step_perc"], step=1, label="Tea Cache starting moment in % of generation") 
 
                     with gr.Row():
-                        gr.Markdown("<B>Upsampling</B>")
+                        gr.Markdown("<B>Upsampling - postprocessing that may improve fluidity and the size of the video</B>")
                     with gr.Row():
                         temporal_upsampling_choice = gr.Dropdown(
                             choices=[
@@ -2687,9 +2778,10 @@ def generate_video_tab(image2video=False):
             show_advanced.change(fn=switch_advanced, inputs=[state, show_advanced, lset_name], outputs=[advanced_row, preset_buttons_rows, refresh_lora_btn, refresh2_row ,lset_name ]).then(
                 fn=switch_prompt_type, inputs = [state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars], outputs = [wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, prompt_column_advanced, prompt_column_wizard, prompt_column_wizard_vars, *prompt_vars])
         with gr.Column():
-            gen_status = gr.Text(label="Status", interactive= False)
-            full_sync = gr.Text(label="Status", interactive= False, visible= False)
-            light_sync = gr.Text(label="Status", interactive= False, visible= False)
+            gen_status = gr.Text(interactive= False)
+            full_sync = gr.Text(interactive= False, visible= False)
+            light_sync = gr.Text(interactive= False, visible= False)
+
             gen_progress_html = gr.HTML(
                 label="Status",
                 value="Idle",
@@ -2709,8 +2801,8 @@ def generate_video_tab(image2video=False):
                     abort_btn = gr.Button("Abort")
 
                 queue_df = gr.DataFrame(
-                    headers=["Qty","Prompt", "Length","Steps","Start", "End", "", "", ""],
-                    datatype=[ "str","markdown","str", "markdown", "markdown", "markdown",    "str", "str", "str"],
+                    headers=["Qty","Prompt", "Length","Steps","", "", "", "", ""],
+                    datatype=[ "str","markdown","str", "markdown", "markdown", "markdown", "str", "str", "str"],
                     column_widths= ["50","", "65","55", "60", "60", "30", "30", "35"],
                     interactive=False,
                     col_count=(9, "fixed"),
@@ -2792,7 +2884,7 @@ def generate_video_tab(image2video=False):
                     show_progress="hidden"
                 )
         save_settings_btn.click( fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var,  prompt, wizard_prompt, *prompt_vars] , outputs= [prompt]).then(
-            save_settings, inputs = [state, prompt, image_prompt_type_radio, video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, 
+            save_settings, inputs = [state, prompt, image_prompt_type_radio, max_frames, remove_background_image_ref,  video_length, resolution, num_inference_steps, seed, repeat_generation, multi_images_gen_type, guidance_scale, flow_shift, negative_prompt, 
                                                          loras_choices, loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, temporal_upsampling_choice, spatial_upsampling_choice, RIFLEx_setting, slg_switch, slg_layers,
                                                          slg_start_perc, slg_end_perc, cfg_star_switch, cfg_zero_step  ], outputs = [])
         save_lset_btn.click(validate_save_lset, inputs=[lset_name], outputs=[apply_lset_btn, refresh_lora_btn, delete_lset_btn, save_lset_btn,confirm_save_lset_btn, cancel_lset_btn, save_lset_prompt_drop])
@@ -2808,21 +2900,30 @@ def generate_video_tab(image2video=False):
         refresh_lora_btn2.click(refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices])
         output.select(select_video, state, None )
 
+
+            
         gen_status.change(refresh_gallery,
             inputs = [state, gen_status], 
             outputs = [output, gen_info, generate_btn, add_to_queue_btn, current_gen_column,  queue_df, abort_btn])
         
-        full_sync.change(refresh_gallery,
+        full_sync.change(fn= check_refresh_input_type,
+            inputs= [state],
+            outputs= [trigger_refresh_input_type]
+        ).then(fn=refresh_gallery,
             inputs = [state, gen_status],
             outputs = [output, gen_info, generate_btn, add_to_queue_btn, current_gen_column,  queue_df, abort_btn] 
-        ).then( fn=wait_tasks_done,
+        ).then(fn=wait_tasks_done,
             inputs= [state], 
             outputs =[gen_status],
         ).then(finalize_generation,
             inputs= [state], 
             outputs= [output, abort_btn, generate_btn, add_to_queue_btn, current_gen_column, gen_info]
         )
-        light_sync.change(refresh_gallery,
+
+        light_sync.change(fn= check_refresh_input_type,
+            inputs= [state],
+            outputs= [trigger_refresh_input_type]
+        ).then(fn=refresh_gallery,
             inputs = [state, gen_status],
             outputs = [output, gen_info, generate_btn, add_to_queue_btn, current_gen_column,  queue_df, abort_btn] 
         )
@@ -2848,10 +2949,11 @@ def generate_video_tab(image2video=False):
             loras_choices,
             loras_mult_choices,
             image_prompt_type_radio,
-            image_to_continue,
-            image_to_end,
-            video_to_continue,
+            image_source1,
+            image_source2,
+            image_source3,
             max_frames,
+            remove_background_image_ref,
             temporal_upsampling_choice,
             spatial_upsampling_choice,
             RIFLEx_setting,
@@ -2902,7 +3004,7 @@ def generate_video_tab(image2video=False):
         )
     return loras_column, loras_choices, presets_column, lset_name, header, light_sync, full_sync, state 
 
-def generate_doxnload_tab(presets_column, loras_column, lset_name,loras_choices, state):
+def generate_download_tab(presets_column, loras_column, lset_name,loras_choices, state):
     with gr.Row():
         with gr.Row(scale =2):
             gr.Markdown("<I>Wan2GP's Lora Festival ! Press the following button to download i2v <B>Remade</B> Loras collection (and bonuses Loras).")
@@ -2928,6 +3030,7 @@ def generate_configuration_tab():
                 ("WAN 2.1 1.3B Text to Video 16 bits (recommended)- the small model for fast generations with low VRAM requirements", 0),
                 ("WAN 2.1 14B Text to Video 16 bits - the default engine in its original glory, offers a slightly better image quality but slower and requires more RAM", 1),
                 ("WAN 2.1 14B Text to Video quantized to 8 bits (recommended) - the default engine but quantized", 2),
+                ("WAN 2.1 VACE 1.3B Text to Video / Control Net - text generation driven by reference images or videos", 3),
             ],
             value= index,
             label="Transformer model for Text to Video",
@@ -3108,16 +3211,17 @@ def on_tab_select(global_state, t2v_state, i2v_state, evt: gr.SelectData):
     t2v_light_sync = gr.Text()
     i2v_full_sync = gr.Text()
     t2v_full_sync = gr.Text()
-    if new_t2v or new_i2v:
-        last_tab_was_image2video =global_state.get("last_tab_was_image2video", None)
-        if last_tab_was_image2video == None or last_tab_was_image2video:
-            gen = i2v_state["gen"]
-            t2v_state["gen"] = gen
-        else:
-            gen = t2v_state["gen"]
-            i2v_state["gen"] = gen
 
+    last_tab_was_image2video =global_state.get("last_tab_was_image2video", None)
+    if last_tab_was_image2video == None or last_tab_was_image2video:
+        gen = i2v_state["gen"]
+        t2v_state["gen"] = gen
+    else:
+        gen = t2v_state["gen"]
+        i2v_state["gen"] = gen
 
+
+    if new_t2v or new_i2v:
         if last_tab_was_image2video != None and new_t2v != new_i2v:
             gen_location = gen.get("location", None)
             if  "in_progress" in gen and gen_location !=None and not (gen_location and new_i2v or not gen_location and new_t2v) :
@@ -3131,7 +3235,6 @@ def on_tab_select(global_state, t2v_state, i2v_state, evt: gr.SelectData):
                 else:
                     t2v_light_sync = gr.Text(str(time.time()))
 
-
         global_state["last_tab_was_image2video"] = new_i2v
 
     if(server_config.get("reload_model",2) == 1):
@@ -3433,7 +3536,7 @@ def create_demo():
         }
     """
     with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="sky", neutral_hue="slate", text_size="md"), title= "Wan2GP") as demo:
-        gr.Markdown("<div align=center><H1>Wan 2.1<SUP>GP</SUP> v3.4 <FONT SIZE=4>by <I>DeepBeepMeep</I></FONT> <FONT SIZE=3> (<A HREF='https://github.com/deepbeepmeep/Wan2GP'>Updates</A>)</FONT SIZE=3></H1></div>")
+        gr.Markdown("<div align=center><H1>Wan 2.1<SUP>GP</SUP> v4.0 <FONT SIZE=4>by <I>DeepBeepMeep</I></FONT> <FONT SIZE=3> (<A HREF='https://github.com/deepbeepmeep/Wan2GP'>Updates</A>)</FONT SIZE=3></H1></div>")
         gr.Markdown("<FONT SIZE=3>Welcome to Wan 2.1GP a super fast and low VRAM AI Video Generator !</FONT>")
         
         with gr.Accordion("Click here for some Info on how to use Wan2GP", open = False):
@@ -3454,7 +3557,7 @@ def create_demo():
                 i2v_loras_column, i2v_loras_choices, i2v_presets_column, i2v_lset_name, i2v_header, i2v_light_sync, i2v_full_sync, i2v_state = generate_video_tab(True)
             if not args.lock_config:
                 with gr.Tab("Downloads", id="downloads") as downloads_tab:
-                    generate_doxnload_tab(i2v_presets_column, i2v_loras_column, i2v_lset_name, i2v_loras_choices, i2v_state)
+                    generate_download_tab(i2v_presets_column, i2v_loras_column, i2v_lset_name, i2v_loras_choices, i2v_state)
                 with gr.Tab("Configuration"):
                     generate_configuration_tab()
             with gr.Tab("About"):

requirements.txt

@@ -11,11 +11,15 @@ easydict
 ftfy
 dashscope
 imageio-ffmpeg
-# flash_attn
+# flash_attn    
 gradio>=5.0.0
 numpy>=1.23.5,<2
 einops
 moviepy==1.0.3
 mmgp==3.3.4
 peft==0.14.0
-mutagen
+mutagen
+decord
+onnxruntime-gpu
+rembg[gpu]==2.0.65
+# rembg==2.0.65

wan/configs/__init__.py

@@ -40,3 +40,17 @@ SUPPORTED_SIZES = {
     'i2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
     't2i-14B': tuple(SIZE_CONFIGS.keys()),
 }
+
+VACE_SIZE_CONFIGS = {
+    '480*832': (480, 832),
+    '832*480': (832, 480),
+}
+
+VACE_MAX_AREA_CONFIGS = {
+    '480*832': 480 * 832,
+    '832*480': 832 * 480,
+}
+
+VACE_SUPPORTED_SIZES = {
+    'vace-1.3B': ('480*832', '832*480'),
+}

wan/modules/model.py

@@ -377,6 +377,7 @@ class WanI2VCrossAttention(WanSelfAttention):
         return x
 
 
+
 WAN_CROSSATTENTION_CLASSES = {
     't2v_cross_attn': WanT2VCrossAttention,
     'i2v_cross_attn': WanI2VCrossAttention,
@@ -393,7 +394,9 @@ class WanAttentionBlock(nn.Module):
                  window_size=(-1, -1),
                  qk_norm=True,
                  cross_attn_norm=False,
-                 eps=1e-6):
+                 eps=1e-6,
+                 block_id=None
+                 ):
         super().__init__()
         self.dim = dim
         self.ffn_dim = ffn_dim
@@ -422,6 +425,7 @@ class WanAttentionBlock(nn.Module):
 
         # modulation
         self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+        self.block_id = block_id
 
     def forward(
         self,
@@ -432,6 +436,8 @@ class WanAttentionBlock(nn.Module):
         freqs,
         context,
         context_lens,
+        hints= None, 
+        context_scale=1.0,
     ):
         r"""
         Args:
@@ -480,10 +486,49 @@ class WanAttentionBlock(nn.Module):
         x.addcmul_(y, e[5])
 
        
-
-        return x
-
-
+        if self.block_id is not None and hints != None:
+            if context_scale == 1:
+                x.add_(hints[self.block_id])
+            else:
+                x.add_(hints[self.block_id], alpha =context_scale)
+        return x    
+
+class VaceWanAttentionBlock(WanAttentionBlock):
+    def __init__(
+            self,
+            cross_attn_type,
+            dim,
+            ffn_dim,
+            num_heads,
+            window_size=(-1, -1),
+            qk_norm=True,
+            cross_attn_norm=False,
+            eps=1e-6,
+            block_id=0
+    ):
+        super().__init__(cross_attn_type, dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps)
+        self.block_id = block_id
+        if block_id == 0:
+            self.before_proj = nn.Linear(self.dim, self.dim)
+            nn.init.zeros_(self.before_proj.weight)
+            nn.init.zeros_(self.before_proj.bias)
+        self.after_proj = nn.Linear(self.dim, self.dim)
+        nn.init.zeros_(self.after_proj.weight)
+        nn.init.zeros_(self.after_proj.bias)
+
+    def forward(self, c, x, **kwargs):
+        # behold dbm magic !
+        if self.block_id == 0:
+            c = self.before_proj(c) + x
+            all_c = []
+        else:
+            all_c = c
+            c = all_c.pop(-1)
+        c = super().forward(c, **kwargs)
+        c_skip = self.after_proj(c)
+        all_c += [c_skip, c]
+        return all_c
+    
 class Head(nn.Module):
 
     def __init__(self, dim, out_dim, patch_size, eps=1e-6):
@@ -544,6 +589,8 @@ class WanModel(ModelMixin, ConfigMixin):
 
     @register_to_config
     def __init__(self,
+                 vace_layers=None,
+                 vace_in_dim=None,                 
                  model_type='t2v',
                  patch_size=(1, 2, 2),
                  text_len=512,
@@ -628,12 +675,13 @@ class WanModel(ModelMixin, ConfigMixin):
         self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
 
         # blocks
-        cross_attn_type = 't2v_cross_attn' if model_type == 't2v' 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)
-            for _ in range(num_layers)
-        ])
+        if vace_layers == None:
+            cross_attn_type = 't2v_cross_attn' if model_type == 't2v' 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)
+                for _ in range(num_layers)
+            ])
 
         # head
         self.head = Head(dim, out_dim, patch_size, eps)
@@ -646,6 +694,33 @@ class WanModel(ModelMixin, ConfigMixin):
         # initialize weights
         self.init_weights()
 
+        if vace_layers != None:            
+            self.vace_layers = [i for i in range(0, self.num_layers, 2)] if vace_layers is None else vace_layers
+            self.vace_in_dim = self.in_dim if vace_in_dim is None else vace_in_dim
+
+            assert 0 in self.vace_layers
+            self.vace_layers_mapping = {i: n for n, i in enumerate(self.vace_layers)}
+
+            # blocks
+            self.blocks = nn.ModuleList([
+                WanAttentionBlock('t2v_cross_attn', self.dim, self.ffn_dim, self.num_heads, self.window_size, self.qk_norm,
+                                    self.cross_attn_norm, self.eps,
+                                    block_id=self.vace_layers_mapping[i] if i in self.vace_layers else None)
+                for i in range(self.num_layers)
+            ])
+
+            # vace blocks
+            self.vace_blocks = nn.ModuleList([
+                VaceWanAttentionBlock('t2v_cross_attn', self.dim, self.ffn_dim, self.num_heads, self.window_size, self.qk_norm,
+                                        self.cross_attn_norm, self.eps, block_id=i)
+                for i in self.vace_layers
+            ])
+
+            # vace patch embeddings
+            self.vace_patch_embedding = nn.Conv3d(
+                self.vace_in_dim, self.dim, kernel_size=self.patch_size, stride=self.patch_size
+            )
+
 
     def compute_teacache_threshold(self, start_step, timesteps = None, speed_factor =0):
         rescale_func = np.poly1d(self.coefficients)         
@@ -688,6 +763,36 @@ class WanModel(ModelMixin, ConfigMixin):
         self.rel_l1_thresh = best_threshold
         print(f"Tea Cache, best threshold found:{best_threshold:0.2f} with gain x{len(timesteps)/(target_nb_steps - best_signed_diff):0.2f} for a target of x{speed_factor}")
         return best_threshold
+
+    def forward_vace(
+        self,
+        x,
+        vace_context,
+        seq_len,
+        context,
+        e,
+        kwargs
+    ):
+        # embeddings
+        c = [self.vace_patch_embedding(u.unsqueeze(0)) for u in vace_context]
+        c = [u.flatten(2).transpose(1, 2) for u in c]
+        if (len(c) == 1 and seq_len == c[0].size(1)):
+            c = c[0]
+        else:
+            c = torch.cat([
+                torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+                        dim=1) for u in c
+            ])
+
+        # arguments
+        new_kwargs = dict(x=x)
+        new_kwargs.update(kwargs)
+
+        for block in self.vace_blocks:
+            c = block(c, context= context, e= e, **new_kwargs)
+        hints = c[:-1]
+
+        return hints
     
     def forward(
         self,
@@ -695,6 +800,8 @@ class WanModel(ModelMixin, ConfigMixin):
         t,
         context,
         seq_len,
+        vace_context = None,
+        vace_context_scale=1.0,        
         clip_fea=None,
         y=None,
         freqs = None,
@@ -829,13 +936,23 @@ class WanModel(ModelMixin, ConfigMixin):
                     self.previous_residual_cond = None
                 ori_hidden_states = x_list[0].clone()
             # arguments
+
             kwargs = dict(
-                # e=e0,
                 seq_lens=seq_lens,
                 grid_sizes=grid_sizes,
                 freqs=freqs,
-                # context=context,
                 context_lens=context_lens)
+
+            if vace_context == None:
+                hints_list = [None ] *len(x_list)
+            else:
+                hints_list = []
+                for x, context in  zip(x_list,  context_list) :
+                    hints_list.append( self.forward_vace(x, vace_context, seq_len, context= context, e= e0,  kwargs= kwargs))
+                del x, context
+                kwargs['context_scale'] = vace_context_scale                
+
+            
             for block_idx, block in enumerate(self.blocks):
                 offload.shared_state["layer"] = block_idx
                 if callback != None:
@@ -852,9 +969,10 @@ class WanModel(ModelMixin, ConfigMixin):
                     x_list[0] = block(x_list[0], context = context_list[0], e= e0, **kwargs)
 
                 else:
-                    for i, (x, context) in enumerate(zip(x_list, context_list)):
-                        x_list[i] = block(x, context = context, e= e0, **kwargs)
+                    for i, (x, context, hints) in enumerate(zip(x_list, context_list, hints_list)):
+                        x_list[i] = block(x, context = context, hints= hints, e= e0, **kwargs)
                         del x
+                    del context, hints
 
             if self.enable_teacache:
                 if joint_pass:

wan/text2video.py

@@ -13,7 +13,9 @@ import torch
 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
@@ -22,6 +24,7 @@ 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):
@@ -105,8 +108,6 @@ class WanT2V:
 
         self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel, writable_tensors= False)
 
-
-
         self.model.eval().requires_grad_(False)
 
         if use_usp:
@@ -132,8 +133,148 @@ class WanT2V:
 
         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)
+
+    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, num_frames, image_size,  device, trim_video= 0):
+        image_sizes = []
+        for i, (sub_src_video, sub_src_mask) in enumerate(zip(src_video, src_mask)):
+            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)
+                src_video[i] = src_video[i].to(device)
+                src_mask[i] = src_mask[i].to(device)
+                src_video_shape = src_video[i].shape
+                if src_video_shape[1] != num_frames:
+                    src_video[i] =  torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], num_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], num_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:
+                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)
+                src_video[i] = src_video[i].to(device)
+                src_video_shape = src_video[i].shape
+                if  src_video_shape[1] != num_frames:
+                    src_video[i] =  torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], num_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1)
+                src_mask[i] = torch.ones_like(src_video[i], device=device)
+                image_sizes.append(src_video[i].shape[2:])
+
+        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(self,
                 input_prompt,
+                input_frames= None,
+                input_masks = None,
+                input_ref_images = None,        
+                context_scale=1.0,
                 size=(1280, 720),
                 frame_num=81,
                 shift=5.0,
@@ -187,14 +328,6 @@ class WanT2V:
                 - W: Frame width from size)
         """
         # preprocess
-        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] / self.sp_size) * self.sp_size
 
         if n_prompt == "":
             n_prompt = self.sample_neg_prompt
@@ -213,6 +346,29 @@ class WanT2V:
             context_null = self.text_encoder([n_prompt], torch.device('cpu'))
             context = [t.to(self.device) for t in context]
             context_null = [t.to(self.device) for t in context_null]
+        
+        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] / self.sp_size) * self.sp_size
+
 
         noise = [
             torch.randn(
@@ -261,10 +417,12 @@ class WanT2V:
         arg_c = {'context': context, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self}
         arg_null = {'context': context_null, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self}
         arg_both = {'context': context, 'context2': context_null, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self}
+        if input_frames != None:
+            vace_dict = {'vace_context' : z, 'vace_context_scale' : context_scale}
+            arg_c.update(vace_dict)
+            arg_null.update(vace_dict)
+            arg_both.update(vace_dict)
 
-        # arg_c = {'context': context, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self, "max_steps": sampling_steps}
-        # arg_null = {'context': context_null, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self, "max_steps": sampling_steps}
-        # arg_both = {'context': context, 'context2': context_null, 'seq_len': seq_len, 'freqs': freqs, 'pipeline': self, "max_steps": sampling_steps}
         if self.model.enable_teacache:
             self.model.compute_teacache_threshold(self.model.teacache_start_step, timesteps, self.model.teacache_multiplier)
         if callback != None:
@@ -281,7 +439,7 @@ class WanT2V:
             # self.model.to(self.device)
             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)
+                    latent_model_input, t=timestep,  current_step=i, slg_layers=slg_layers_local, **arg_both)
                 if self._interrupt:
                     return None
             else:
@@ -329,7 +487,11 @@ class WanT2V:
             self.model.cpu()
             torch.cuda.empty_cache()
         if self.rank == 0:
-            videos = self.vae.decode(x0, VAE_tile_size)
+
+            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 noise, latents

wan/utils/utils.py

@@ -3,21 +3,70 @@ import argparse
 import binascii
 import os
 import os.path as osp
+import torchvision.transforms.functional as TF
+import torch.nn.functional as F
 
 import imageio
 import torch
+import decord
 import torchvision
 from PIL import Image
 import numpy as np
+from rembg import remove, new_session
+
 
 __all__ = ['cache_video', 'cache_image', 'str2bool']
 
+
+
+from PIL import Image
+
+def get_video_frame(file_name, frame_no):
+    decord.bridge.set_bridge('torch')
+    reader = decord.VideoReader(file_name)
+
+    frame = reader.get_batch([frame_no]).squeeze(0)
+    img = Image.fromarray(frame.numpy().astype(np.uint8))
+    return img
+
 def resize_lanczos(img, h, w):
     img = Image.fromarray(np.clip(255. * img.movedim(0, -1).cpu().numpy(), 0, 255).astype(np.uint8))
     img = img.resize((w,h), resample=Image.Resampling.LANCZOS) 
     return torch.from_numpy(np.array(img).astype(np.float32) / 255.0).movedim(-1, 0)
 
 
+def remove_background(img, session=None):
+    if session ==None:
+        session = new_session() 
+    img = Image.fromarray(np.clip(255. * img.movedim(0, -1).cpu().numpy(), 0, 255).astype(np.uint8))
+    img = remove(img, session=session, alpha_matting = True, bgcolor=[255, 255, 255, 0]).convert('RGB')
+    return torch.from_numpy(np.array(img).astype(np.float32) / 255.0).movedim(-1, 0)
+
+
+
+
+def resize_and_remove_background(img_list, canvas_width, canvas_height, rm_background ):
+    if rm_background:
+        session = new_session() 
+
+    output_list =[]
+    for img in img_list:
+        width, height =  img.size 
+        white_canvas = np.full( (canvas_height, canvas_width, 3), 255, dtype= np.uint8 ) 
+        scale = min(canvas_height / height, canvas_width / width)
+        new_height = int(height * scale)
+        new_width = int(width * scale)
+        resized_image= img.resize((new_width,new_height), resample=Image.Resampling.LANCZOS) 
+        if rm_background:
+            resized_image = remove(resized_image, session=session, alpha_matting = True, bgcolor=[255, 255, 255, 0]).convert('RGB')
+        top = (canvas_height - new_height) // 2
+        left = (canvas_width - new_width) // 2
+        white_canvas[top:top + new_height, left:left + new_width, :] =  np.array(resized_image)
+        img = Image.fromarray(white_canvas)
+        output_list.append(img)
+    return output_list
+
+
 def rand_name(length=8, suffix=''):
     name = binascii.b2a_hex(os.urandom(length)).decode('utf-8')
     if suffix:

wan/utils/vace_preprocessor.py

@@ -0,0 +1,298 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+
+
+class VaceImageProcessor(object):
+    def __init__(self, downsample=None, seq_len=None):
+        self.downsample = downsample
+        self.seq_len = seq_len
+
+    def _pillow_convert(self, image, cvt_type='RGB'):
+        if image.mode != cvt_type:
+            if image.mode == 'P':
+                image = image.convert(f'{cvt_type}A')
+            if image.mode == f'{cvt_type}A':
+                bg = Image.new(cvt_type,
+                               size=(image.width, image.height),
+                               color=(255, 255, 255))
+                bg.paste(image, (0, 0), mask=image)
+                image = bg
+            else:
+                image = image.convert(cvt_type)
+        return image
+
+    def _load_image(self, img_path):
+        if img_path is None or img_path == '':
+            return None
+        img = Image.open(img_path)
+        img = self._pillow_convert(img)
+        return img
+
+    def _resize_crop(self, img, oh, ow, normalize=True):
+        """
+        Resize, center crop, convert to tensor, and normalize.
+        """
+        # resize and crop
+        iw, ih = img.size
+        if iw != ow or ih != oh:
+            # resize
+            scale = max(ow / iw, oh / ih)
+            img = img.resize(
+                (round(scale * iw), round(scale * ih)),
+                resample=Image.Resampling.LANCZOS
+            )
+            assert img.width >= ow and img.height >= oh
+
+            # center crop
+            x1 = (img.width - ow) // 2
+            y1 = (img.height - oh) // 2
+            img = img.crop((x1, y1, x1 + ow, y1 + oh))
+
+        # normalize
+        if normalize:
+            img = TF.to_tensor(img).sub_(0.5).div_(0.5).unsqueeze(1)
+        return img
+    
+    def _image_preprocess(self, img, oh, ow, normalize=True, **kwargs):
+        return self._resize_crop(img, oh, ow, normalize)
+
+    def load_image(self, data_key, **kwargs):
+        return self.load_image_batch(data_key, **kwargs)
+
+    def load_image_pair(self, data_key, data_key2, **kwargs):
+        return self.load_image_batch(data_key, data_key2, **kwargs)
+
+    def load_image_batch(self, *data_key_batch, normalize=True, seq_len=None, **kwargs):
+        seq_len = self.seq_len if seq_len is None else seq_len
+        imgs = []
+        for data_key in data_key_batch:
+            img = self._load_image(data_key)
+            imgs.append(img)
+        w, h = imgs[0].size
+        dh, dw = self.downsample[1:]
+
+        # compute output size
+        scale = min(1., np.sqrt(seq_len / ((h / dh) * (w / dw))))
+        oh = int(h * scale) // dh * dh
+        ow = int(w * scale) // dw * dw
+        assert (oh // dh) * (ow // dw) <= seq_len
+        imgs = [self._image_preprocess(img, oh, ow, normalize) for img in imgs]
+        return *imgs, (oh, ow)
+
+
+class VaceVideoProcessor(object):
+    def __init__(self, downsample, min_area, max_area, min_fps, max_fps, zero_start, seq_len, keep_last, **kwargs):
+        self.downsample = downsample
+        self.min_area = min_area
+        self.max_area = max_area
+        self.min_fps = min_fps
+        self.max_fps = max_fps
+        self.zero_start = zero_start
+        self.keep_last = keep_last
+        self.seq_len = seq_len
+        assert seq_len >= min_area / (self.downsample[1] * self.downsample[2])
+
+    @staticmethod
+    def resize_crop(video: torch.Tensor, oh: int, ow: int):
+        """
+        Resize, center crop and normalize for decord loaded video (torch.Tensor type)
+
+        Parameters:
+          video - video to process (torch.Tensor): Tensor from `reader.get_batch(frame_ids)`, in shape of (T, H, W, C)
+          oh - target height (int)
+          ow - target width (int)
+
+        Returns:
+            The processed video (torch.Tensor): Normalized tensor range [-1, 1], in shape of (C, T, H, W)
+
+        Raises:
+        """
+        # permute ([t, h, w, c] -> [t, c, h, w])
+        video = video.permute(0, 3, 1, 2)
+
+        # resize and crop
+        ih, iw = video.shape[2:]
+        if ih != oh or iw != ow:
+            # resize
+            scale = max(ow / iw, oh / ih)
+            video = F.interpolate(
+                video,
+                size=(round(scale * ih), round(scale * iw)),
+                mode='bicubic',
+                antialias=True
+            )
+            assert video.size(3) >= ow and video.size(2) >= oh
+
+            # center crop
+            x1 = (video.size(3) - ow) // 2
+            y1 = (video.size(2) - oh) // 2
+            video = video[:, :, y1:y1 + oh, x1:x1 + ow]
+
+        # permute ([t, c, h, w] -> [c, t, h, w]) and normalize
+        video = video.transpose(0, 1).float().div_(127.5).sub_(1.)
+        return video
+
+    def _video_preprocess(self, video, oh, ow):
+        return self.resize_crop(video, oh, ow)
+
+    def _get_frameid_bbox_default(self, fps, frame_timestamps, h, w, crop_box, rng):
+        target_fps = min(fps, self.max_fps)
+        duration = frame_timestamps[-1].mean()
+        x1, x2, y1, y2 = [0, w, 0, h] if crop_box is None else crop_box
+        h, w = y2 - y1, x2 - x1
+        ratio = h / w
+        df, dh, dw = self.downsample
+
+        # min/max area of the [latent video]
+        min_area_z = self.min_area / (dh * dw)
+        max_area_z = min(self.seq_len, self.max_area / (dh * dw), (h // dh) * (w // dw))
+
+        # sample a frame number of the [latent video]
+        rand_area_z = np.square(np.power(2, rng.uniform(
+            np.log2(np.sqrt(min_area_z)),
+            np.log2(np.sqrt(max_area_z))
+        )))
+        of = min(
+            (int(duration * target_fps) - 1) // df + 1,
+            int(self.seq_len / rand_area_z)
+        )
+
+        # deduce target shape of the [latent video]
+        target_area_z = min(max_area_z, int(self.seq_len / of))
+        oh = round(np.sqrt(target_area_z * ratio))
+        ow = int(target_area_z / oh)
+        of = (of - 1) * df + 1
+        oh *= dh
+        ow *= dw
+
+        # sample frame ids
+        target_duration = of / target_fps
+        begin = 0. if self.zero_start else rng.uniform(0, duration - target_duration)
+        timestamps = np.linspace(begin, begin + target_duration, of)
+        frame_ids = np.argmax(np.logical_and(
+            timestamps[:, None] >= frame_timestamps[None, :, 0],
+            timestamps[:, None] < frame_timestamps[None, :, 1]
+        ), axis=1).tolist()
+        return frame_ids, (x1, x2, y1, y2), (oh, ow), target_fps
+
+    def _get_frameid_bbox_adjust_last(self, fps, frame_timestamps, h, w, crop_box, rng, max_frames= 0):
+        import math
+        target_fps = self.max_fps
+        video_duration = frame_timestamps[-1][1]
+        video_frame_duration = 1 /fps
+        target_frame_duration = 1 / target_fps 
+        
+        cur_time = 0
+        target_time = 0
+        frame_no = 0
+        frame_ids =[]
+        for i in range(max_frames):
+            add_frames_count = math.ceil( (target_time -cur_time) / video_frame_duration )
+            frame_no += add_frames_count
+            frame_ids.append(frame_no)
+            cur_time += add_frames_count * video_frame_duration
+            target_time += target_frame_duration
+            if cur_time > video_duration: 
+                break
+
+        x1, x2, y1, y2 = [0, w, 0, h] if crop_box is None else crop_box
+        h, w = y2 - y1, x2 - x1
+        ratio = h / w
+        df, dh, dw = self.downsample
+        seq_len  = self.seq_len
+        # min/max area of the [latent video]
+        min_area_z = self.min_area / (dh * dw)
+        # max_area_z = min(seq_len, self.max_area / (dh * dw), (h // dh) * (w // dw))
+        max_area_z = min_area_z # workaround bug
+        # sample a frame number of the [latent video]
+        rand_area_z = np.square(np.power(2, rng.uniform(
+            np.log2(np.sqrt(min_area_z)),
+            np.log2(np.sqrt(max_area_z))
+        )))
+
+        seq_len =  max_area_z * ((max_frames- 1) // df +1)
+
+        # of = min(
+        #     (len(frame_ids) - 1) // df + 1,
+        #     int(seq_len / rand_area_z)
+        # )
+        of = (len(frame_ids) - 1) // df + 1
+
+             
+        # deduce target shape of the [latent video]
+        # target_area_z = min(max_area_z, int(seq_len / of))
+        target_area_z = max_area_z        
+        oh = round(np.sqrt(target_area_z * ratio))
+        ow = int(target_area_z / oh)
+        of = (of - 1) * df + 1
+        oh *= dh
+        ow *= dw
+
+        return frame_ids, (x1, x2, y1, y2), (oh, ow), target_fps
+
+    def _get_frameid_bbox(self, fps, frame_timestamps, h, w, crop_box, rng, max_frames= 0):
+        if self.keep_last:
+            return self._get_frameid_bbox_adjust_last(fps, frame_timestamps, h, w, crop_box, rng, max_frames= max_frames)
+        else:
+            return self._get_frameid_bbox_default(fps, frame_timestamps, h, w, crop_box, rng, max_frames= max_frames)
+
+    def load_video(self, data_key, crop_box=None, seed=2024, **kwargs):
+        return self.load_video_batch(data_key, crop_box=crop_box, seed=seed, **kwargs)
+
+    def load_video_pair(self, data_key, data_key2, crop_box=None, seed=2024, **kwargs):
+        return self.load_video_batch(data_key, data_key2, crop_box=crop_box, seed=seed, **kwargs)
+
+    def load_video_batch(self, *data_key_batch, crop_box=None, seed=2024, max_frames= 0, trim_video =0, **kwargs):
+        rng = np.random.default_rng(seed + hash(data_key_batch[0]) % 10000)
+        # read video
+        import decord
+        decord.bridge.set_bridge('torch')
+        readers = []
+        for data_k in data_key_batch:
+            reader = decord.VideoReader(data_k)
+            readers.append(reader)
+
+        fps = readers[0].get_avg_fps()
+        length = min([len(r) for r in readers])
+        frame_timestamps = [readers[0].get_frame_timestamp(i) for i in range(length)]
+        frame_timestamps = np.array(frame_timestamps, dtype=np.float32)
+        # # frame_timestamps = frame_timestamps[ :max_frames]
+        # if trim_video > 0:
+        #     frame_timestamps = frame_timestamps[ :trim_video]
+        max_frames = min(max_frames, trim_video) if trim_video > 0 else max_frames
+        h, w = readers[0].next().shape[:2]
+        frame_ids, (x1, x2, y1, y2), (oh, ow), fps = self._get_frameid_bbox(fps, frame_timestamps, h, w, crop_box, rng, max_frames=max_frames)
+
+        # preprocess video
+        videos = [reader.get_batch(frame_ids)[:, y1:y2, x1:x2, :] for reader in readers]
+        videos = [self._video_preprocess(video, oh, ow) for video in videos]
+        return *videos, frame_ids, (oh, ow), fps
+        # return videos if len(videos) > 1 else videos[0]
+
+
+def prepare_source(src_video, src_mask, src_ref_images, num_frames, image_size, device):
+    for i, (sub_src_video, sub_src_mask) in enumerate(zip(src_video, src_mask)):
+        if sub_src_video is None and sub_src_mask is None:
+            src_video[i] = torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device)
+            src_mask[i] = torch.ones((1, num_frames, image_size[0], image_size[1]), device=device)
+    for i, ref_images in enumerate(src_ref_images):
+        if ref_images is not None:
+            for j, ref_img in enumerate(ref_images):
+                if ref_img is not None and 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
+                    src_ref_images[i][j] = white_canvas
+    return src_video, src_mask, src_ref_images