From f7c5193eaad99ec85467c49e21fa2dbf4511bcfa Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 24 Mar 2025 17:20:53 +1100 Subject: [PATCH 01/69] add queues --- gradio_server.py | 836 +++++++++++++++++++++++++++-------------------- 1 file changed, 476 insertions(+), 360 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 65b14ce..7848b39 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1,5 +1,7 @@ import os import time +import threading +from typing import Dict, Any import argparse from mmgp import offload, safetensors2, profile_type try: @@ -21,6 +23,7 @@ import traceback import math import asyncio from wan.utils import prompt_parser +from typing import List PROMPT_VARS_MAX = 10 target_mmgp_version = "3.3.0" @@ -29,6 +32,194 @@ mmgp_version = version("mmgp") if mmgp_version != target_mmgp_version: print(f"Incorrect version of mmgp ({mmgp_version}), version {target_mmgp_version} is needed. Please upgrade with the command 'pip install -r requirements.txt'") exit() +queue = [] +lock = threading.Lock() +current_task_id = None +task_id = 0 +progress_tracker = {} +tracker_lock = threading.Lock() +file_list = [] + +def runner(): + global current_task_id + while True: + with lock: + for item in queue: + task_id = item['id'] + with tracker_lock: + progress = progress_tracker.get(task_id, {}) + + if item['status'] == "Processing": + current_step = progress.get('current_step', 0) + total_steps = progress.get('total_steps', 0) + elapsed = time.time() - progress.get('start_time', time.time()) + item.update({ + 'progress': f"{((current_step/total_steps)*100 if total_steps > 0 else 0):.1f}%", + 'steps': f"{current_step}/{total_steps}", + 'time': f"{elapsed:.1f}s" + }) + if not any(item['status'] == "Processing" for item in queue): + for item in queue: + if item['status'] == "Queued": + item['status'] = "Processing" + current_task_id = item['id'] + threading.Thread(target=process_task, args=(item,)).start() + break + time.sleep(1) + +def process_prompt_and_add_tasks( + prompt, + negative_prompt, + resolution, + video_length, + seed, + num_inference_steps, + guidance_scale, + flow_shift, + embedded_guidance_scale, + repeat_generation, + multi_images_gen_type, + tea_cache, + tea_cache_start_step_perc, + loras_choices, + loras_mult_choices, + image_to_continue, + image_to_end, + video_to_continue, + max_frames, + RIFLEx_setting, + slg_switch, + slg_layers, + slg_start, + slg_end, + state_arg, + image2video +): + if len(prompt) ==0: + return + prompt, errors = prompt_parser.process_template(prompt) + if len(errors) > 0: + gr.Info("Error processing prompt template: " + errors) + return + prompts = prompt.replace("\r", "").split("\n") + prompts = [prompt.strip() for prompt in prompts if len(prompt.strip())>0 and not prompt.startswith("#")] + if len(prompts) ==0: + return + + for single_prompt in prompts: + task_params = ( + single_prompt, + negative_prompt, + resolution, + video_length, + seed, + num_inference_steps, + guidance_scale, + flow_shift, + embedded_guidance_scale, + repeat_generation, + multi_images_gen_type, + tea_cache, + tea_cache_start_step_perc, + loras_choices, + loras_mult_choices, + image_to_continue, + image_to_end, + video_to_continue, + max_frames, + RIFLEx_setting, + slg_switch, + slg_layers, + slg_start, + slg_end, + state_arg, + image2video + ) + add_video_task(*task_params) + return update_queue_data() + +def process_task(task): + try: + task_id, *params = task['params'] + generate_video(task_id, *params) + finally: + with lock: + queue[:] = [item for item in queue if item['id'] != task['id']] + with tracker_lock: + if task['id'] in progress_tracker: + del progress_tracker[task['id']] + +def add_video_task(*params): + global task_id + with lock: + task_id += 1 + current_task_id = task_id + queue.append({ + "id": current_task_id, + "params": (current_task_id,) + params, + "status": "Queued", + "progress": "0.0%", + "steps": f"0/{params[5]}", + "time": "--", + "prompt": params[0] + }) + return + +def move_up(selected_indices): + if not selected_indices or len(selected_indices) == 0: + return update_queue_data() + idx = selected_indices[0] + if isinstance(idx, list): + idx = idx[0] + idx = int(idx) + with lock: + if idx > 0: + queue[idx], queue[idx-1] = queue[idx-1], queue[idx] + return update_queue_data() + +def move_down(selected_indices): + if not selected_indices or len(selected_indices) == 0: + return update_queue_data() + idx = selected_indices[0] + if isinstance(idx, list): + idx = idx[0] + idx = int(idx) + with lock: + if idx < len(queue)-1: + queue[idx], queue[idx+1] = queue[idx+1], queue[idx] + return update_queue_data() + +def remove_task(selected_indices): + if not selected_indices or len(selected_indices) == 0: + return update_queue_data() + idx = selected_indices[0] + if isinstance(idx, list): + idx = idx[0] + idx = int(idx) + with lock: + if idx < len(queue): + if idx == 0: + wan_model._interrupt = True + del queue[idx] + return update_queue_data() + +def update_queue_data(): + with lock: + data = [] + for item in queue: + data.append([ + str(item['id']), + item['status'], + item.get('progress', "0.0%"), + item.get('steps', ''), + item.get('time', '--'), + (item['prompt'][:47] + '...') if len(item['prompt']) > 50 else item['prompt'], + "↑", + "↓", + "✖" + ]) + return data + def _parse_args(): parser = argparse.ArgumentParser( description="Generate a video from a text prompt or image using Gradio") @@ -788,53 +979,22 @@ def save_video(final_frames, output_path, fps=24): final_frames = (final_frames * 255).astype(np.uint8) ImageSequenceClip(list(final_frames), fps=fps).write_videofile(output_path, verbose= False, logger = None) -def build_callback(state, pipe, progress, status, num_inference_steps): - def callback(step_idx, latents, read_state = False): - status = state["progress_status"] - if read_state: - phase, step_idx = state["progress_phase"] - else: - step_idx += 1 - if state.get("abort", False): - # pipe._interrupt = True - phase = " - Aborting" - elif step_idx == num_inference_steps: - phase = " - VAE Decoding" - else: - phase = " - Denoising" - state["progress_phase"] = (phase, step_idx) - status_msg = status + phase - if step_idx >= 0: - progress( (step_idx , num_inference_steps) , status_msg , num_inference_steps) - else: - progress(0, status_msg) - - return callback +def build_callback(task_id, total_steps): + start_time = time.time() + def update_progress(step, _): + with tracker_lock: + elapsed = time.time() - start_time + progress_tracker[task_id] = { + 'current_step': step + 1, + 'total_steps': total_steps, + 'start_time': start_time, + 'last_update': time.time() + } + return update_progress -def abort_generation(state): - if "in_progress" in state: - state["abort"] = True - state["extra_orders"] = 0 - wan_model._interrupt= True - return gr.Button(interactive= False) - else: - return gr.Button(interactive= True) - -def refresh_gallery(state, txt): - file_list = state.get("file_list", None) - prompt = state.get("prompt", "") - if len(prompt) == 0: - return file_list, gr.Text(visible= False, value="") - else: - prompts_max = state.get("prompts_max",0) - prompt_no = state.get("prompt_no",0) - if prompts_max >1 : - label = f"Current Prompt ({prompt_no+1}/{prompts_max})" - else: - label = f"Current Prompt" - return file_list, gr.Text(visible= True, value=prompt, label=label) - - +def refresh_gallery(state): + file_list = state.get("file_list", None) + return file_list def finalize_gallery(state): choice = 0 if "in_progress" in state: @@ -845,7 +1005,7 @@ def finalize_gallery(state): time.sleep(0.2) global gen_in_progress gen_in_progress = False - return gr.Gallery(selected_index=choice), gr.Button(interactive= True), gr.Button(visible= True), gr.Checkbox(visible= False), gr.Text(visible= False, value="") + return gr.Gallery(selected_index=choice), gr.Button(interactive=True), gr.Button(visible=False), gr.Checkbox(visible=False), gr.Text(visible=False, value="") def select_video(state , event_data: gr.EventData): data= event_data._data @@ -862,36 +1022,8 @@ def expand_slist(slist, num_inference_steps ): pos += inc return new_slist - -def one_more_video(state): - extra_orders = state.get("extra_orders", 0) - extra_orders += 1 - state["extra_orders"] = extra_orders - prompts_max = state.get("prompts_max",0) - if prompts_max == 0: - return state - prompt_no = state["prompt_no"] - video_no = state["video_no"] - total_video = state["total_video"] - # total_video += (prompts_max- prompt_no) - total_video += 1 - total_generation = state["total_generation"] + extra_orders - state["total_video"] = total_video - - state["progress_status"] = f"Video {video_no}/{total_video}" - offload.shared_state["refresh"] = 1 - # if (prompts_max - prompt_no) > 1: - # gr.Info(f"An extra video generation is planned for a total of {total_generation} videos for the next {prompts_max - prompt_no} prompts") - # else: - gr.Info(f"An extra video generation is planned for a total of {total_generation} videos for this prompt") - - return state - -def prepare_generate_video(): - - return gr.Button(visible= False), gr.Checkbox(visible= True) - def generate_video( + task_id, prompt, negative_prompt, resolution, @@ -921,7 +1053,6 @@ def generate_video( progress=gr.Progress() #track_tqdm= True ): - global wan_model, offloadobj reload_needed = state.get("_reload_needed", False) file_model_needed = model_needed(image2video) @@ -930,9 +1061,9 @@ def generate_video( offloadobj.release() offloadobj = None wan_model = None - yield f"Loading model {get_model_name(file_model_needed)}..." + print(f"Loading model {get_model_name(file_model_needed)}...") wan_model, offloadobj, trans = load_models(image2video) - yield f"Model loaded" + print(f"Model loaded") state["_reload_needed"] = False from PIL import Image @@ -949,8 +1080,8 @@ def generate_video( gr.Info(f"You have selected attention mode '{attention_mode}'. However it is not installed on your system. You should either install it or switch to the default 'sdpa' attention.") return - if state.get("validate_success",0) != 1: - return + #if state.get("validate_success",0) != 1: + # return width, height = resolution.split("x") width, height = int(width), int(height) @@ -997,71 +1128,11 @@ def generate_video( global gen_in_progress gen_in_progress = True temp_filename = None - if len(prompt) ==0: - return - prompt, errors = prompt_parser.process_template(prompt) - if len(errors) > 0: - gr.Info(f"Error processing prompt template: " + errors) - prompts = prompt.replace("\r", "").split("\n") - prompts = [prompt.strip() for prompt in prompts if len(prompt.strip())>0 and not prompt.startswith("#")] - if len(prompts) ==0: - return if image2video: - if image_to_continue is not None: - if isinstance(image_to_continue, list): - image_to_continue = [ tup[0] for tup in image_to_continue ] - else: - image_to_continue = [image_to_continue] - if image_to_end != None: - if isinstance(image_to_end , list): - image_to_end = [ tup[0] for tup in image_to_end ] - else: - image_to_end = [image_to_end ] - if len(image_to_continue) != len(image_to_end): - 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_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)] ) - prompts = new_prompts - image_to_continue = new_image_to_continue - if image_to_end != None: - image_to_end = new_image_to_end - else: - if len(prompts) >= len(image_to_continue): - if len(prompts) % len(image_to_continue) !=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 = [] - 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 - else: - if len(image_to_continue) % 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) - new_prompts = [] - for i, _ in enumerate(image_to_continue): - new_prompts.append( prompts[ i//rep] ) - prompts = new_prompts - - elif video_to_continue != None and len(video_to_continue) >0 : + if 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 @@ -1148,195 +1219,167 @@ def generate_video( if seed == None or seed <0: seed = random.randint(0, 999999999) - file_list = [] - state["file_list"] = file_list + global file_list + state["file_list"] = file_list global save_path os.makedirs(save_path, exist_ok=True) - video_no = 0 - total_video = repeat_generation * len(prompts) - state["total_video"] = total_video - extra_generation = 0 abort = False - start_time = time.time() - state["prompts_max"] = len(prompts) - for no, prompt in enumerate(prompts): - state["prompt"] = prompt - repeat_no = 0 - state["prompt_no"] = no - extra_generation = 0 - yield f"Prompt No{no}" - while True: - extra_orders = state.get("extra_orders",0) - state["extra_orders"] = 0 - extra_generation += extra_orders - state["total_generation"] = repeat_generation + extra_generation - # total_video += (len(prompts)- no) * extra_orders - total_video += extra_orders - if abort or repeat_no >= (repeat_generation + extra_generation): - break + with tracker_lock: + progress_tracker[task_id] = { + 'current_step': 0, + 'total_steps': num_inference_steps, + 'start_time': time.time(), + 'last_update': time.time() + } + if trans.enable_teacache: + trans.teacache_counter = 0 + trans.num_steps = num_inference_steps + trans.teacache_skipped_steps = 0 + trans.previous_residual_uncond = None + trans.previous_residual_cond = None + callback = build_callback(task_id, num_inference_steps) + offload.shared_state["callback"] = callback + gc.collect() + torch.cuda.empty_cache() + wan_model._interrupt = False + state["progress_status"] = "Starting" + try: + if image2video: + samples = wan_model.generate( + prompt, + image_to_continue.convert('RGB'), + image_to_end.convert('RGB') if image_to_end != None else None, + frame_num=(video_length // 4)* 4 + 1, + max_area=MAX_AREA_CONFIGS[resolution], + shift=flow_shift, + sampling_steps=num_inference_steps, + guide_scale=guidance_scale, + n_prompt=negative_prompt, + seed=seed, + offload_model=False, + callback=callback, + enable_RIFLEx = enable_RIFLEx, + VAE_tile_size = VAE_tile_size, + joint_pass = joint_pass, + slg_layers = slg_layers, + slg_start = slg_start/100, + slg_end = slg_end/100, + ) - if trans.enable_teacache: - trans.teacache_counter = 0 - trans.num_steps = num_inference_steps - trans.teacache_skipped_steps = 0 - trans.previous_residual_uncond = None - trans.previous_residual_cond = None + else: + samples = wan_model.generate( + prompt, + frame_num=(video_length // 4)* 4 + 1, + size=(width, height), + shift=flow_shift, + sampling_steps=num_inference_steps, + guide_scale=guidance_scale, + n_prompt=negative_prompt, + seed=seed, + offload_model=False, + callback=callback, + enable_RIFLEx = enable_RIFLEx, + VAE_tile_size = VAE_tile_size, + joint_pass = joint_pass, + slg_layers = slg_layers, + slg_start = slg_start/100, + slg_end = slg_end/100, + ) + except Exception as e: + gen_in_progress = False + if temp_filename!= None and os.path.isfile(temp_filename): + os.remove(temp_filename) + offload.last_offload_obj.unload_all() + offload.unload_loras_from_model(trans) + # if compile: + # cache_size = torch._dynamo.config.cache_size_limit + # torch.compiler.reset() + # torch._dynamo.config.cache_size_limit = cache_size - video_no += 1 - status = f"Video {video_no}/{total_video}" - state["video_no"] = video_no - state["progress_status"] = status - state["progress_phase"] = (" - Encoding Prompt", -1 ) - progress(0, desc=status + " - Encoding Prompt" ) - callback = build_callback(state, trans, progress, status, num_inference_steps) - offload.shared_state["callback"] = callback - - - gc.collect() - torch.cuda.empty_cache() - wan_model._interrupt = False - try: - if image2video: - samples = wan_model.generate( - prompt, - image_to_continue[no].convert('RGB'), - image_to_end[no].convert('RGB') if image_to_end != None else None, - frame_num=(video_length // 4)* 4 + 1, - max_area=MAX_AREA_CONFIGS[resolution], - shift=flow_shift, - sampling_steps=num_inference_steps, - guide_scale=guidance_scale, - n_prompt=negative_prompt, - seed=seed, - offload_model=False, - callback=callback, - enable_RIFLEx = enable_RIFLEx, - VAE_tile_size = VAE_tile_size, - joint_pass = joint_pass, - slg_layers = slg_layers, - slg_start = slg_start/100, - slg_end = slg_end/100, - ) - - else: - samples = wan_model.generate( - prompt, - frame_num=(video_length // 4)* 4 + 1, - size=(width, height), - shift=flow_shift, - sampling_steps=num_inference_steps, - guide_scale=guidance_scale, - n_prompt=negative_prompt, - seed=seed, - offload_model=False, - callback=callback, - enable_RIFLEx = enable_RIFLEx, - VAE_tile_size = VAE_tile_size, - joint_pass = joint_pass, - slg_layers = slg_layers, - slg_start = slg_start/100, - slg_end = slg_end/100, - ) - except Exception as e: - gen_in_progress = False - if temp_filename!= None and os.path.isfile(temp_filename): - os.remove(temp_filename) - offload.last_offload_obj.unload_all() - offload.unload_loras_from_model(trans) - # if compile: - # cache_size = torch._dynamo.config.cache_size_limit - # torch.compiler.reset() - # torch._dynamo.config.cache_size_limit = cache_size - - gc.collect() - torch.cuda.empty_cache() - s = str(e) - keyword_list = ["vram", "VRAM", "memory","allocat"] - VRAM_crash= False - if any( keyword in s for keyword in keyword_list): + gc.collect() + torch.cuda.empty_cache() + s = str(e) + keyword_list = ["vram", "VRAM", "memory","allocat"] + VRAM_crash= False + if any( keyword in s for keyword in keyword_list): + VRAM_crash = True + else: + stack = traceback.extract_stack(f=None, limit=5) + for frame in stack: + if any( keyword in frame.name for keyword in keyword_list): VRAM_crash = True - else: - stack = traceback.extract_stack(f=None, limit=5) - for frame in stack: - if any( keyword in frame.name for keyword in keyword_list): - VRAM_crash = True - break - state["prompt"] = "" - if VRAM_crash: - raise gr.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: - raise gr.Error(f"The generation of the video has encountered an error, please check your terminal for more information. '{s}'") + break + state["prompt"] = "" + if VRAM_crash: + raise gr.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: + raise gr.Error(f"The generation of the video has encountered an error, please check your terminal for more information. '{s}'") + finally: + with tracker_lock: + if task_id in progress_tracker: + del progress_tracker[task_id] - if trans.enable_teacache: - print(f"Teacache Skipped Steps:{trans.teacache_skipped_steps}/{num_inference_steps}" ) - trans.previous_residual_uncond = None - trans.previous_residual_cond = None + if trans.enable_teacache: + print(f"Teacache Skipped Steps:{trans.teacache_skipped_steps}/{num_inference_steps}" ) + trans.previous_residual_uncond = None + trans.previous_residual_cond = None - if samples != None: - samples = samples.to("cpu") - offload.last_offload_obj.unload_all() - gc.collect() - torch.cuda.empty_cache() + if samples != None: + samples = samples.to("cpu") + offload.last_offload_obj.unload_all() + gc.collect() + torch.cuda.empty_cache() - if samples == None: - end_time = time.time() - abort = True - state["prompt"] = "" - yield f"Video generation was aborted. Total Generation Time: {end_time-start_time:.1f}s" - else: - sample = samples.cpu() - # video = rearrange(sample.cpu().numpy(), "c t h w -> t h w c") + if samples == None: + end_time = time.time() + abort = True + state["prompt"] = "" + else: + sample = samples.cpu() + # video = rearrange(sample.cpu().numpy(), "c t h w -> t h w c") - time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%Hh%Mm%Ss") - if os.name == 'nt': - file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:50]).strip()}.mp4" - else: - file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:100]).strip()}.mp4" - video_path = os.path.join(save_path, file_name) - cache_video( - tensor=sample[None], - save_file=video_path, - fps=16, - nrow=1, - normalize=True, - value_range=(-1, 1)) - - configs = { - 'prompt': prompt, - 'negative_prompt': negative_prompt, - 'resolution': resolution, - 'video_length': video_length, - 'seed': seed, - 'num_inference_steps': num_inference_steps, - } + time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%Hh%Mm%Ss") + if os.name == 'nt': + file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:50]).strip()}.mp4" + else: + file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:100]).strip()}.mp4" + video_path = os.path.join(save_path, file_name) + cache_video( + tensor=sample[None], + save_file=video_path, + fps=16, + nrow=1, + normalize=True, + value_range=(-1, 1)) + + configs = { + 'prompt': prompt, + 'negative_prompt': negative_prompt, + 'resolution': resolution, + 'video_length': video_length, + 'seed': seed, + 'num_inference_steps': num_inference_steps, + } - metadata_choice = server_config["metadata_choice"] - if metadata_choice == "json": - with open(video_path.replace('.mp4', '.json'), 'w') as f: - json.dump(configs, f, indent=4) - elif metadata_choice == "metadata": - from mutagen.mp4 import MP4 - file = MP4(video_path) - file.tags['©cmt'] = [json.dumps(configs)] - file.save() + metadata_choice = server_config["metadata_choice"] + if metadata_choice == "json": + with open(video_path.replace('.mp4', '.json'), 'w') as f: + json.dump(configs, f, indent=4) + elif metadata_choice == "metadata": + from mutagen.mp4 import MP4 + file = MP4(video_path) + file.tags['©cmt'] = [json.dumps(configs)] + file.save() - print(f"New video saved to Path: "+video_path) - file_list.append(video_path) - if video_no < total_video: - yield status - else: - end_time = time.time() - state["prompt"] = "" - yield f"Total Generation Time: {end_time-start_time:.1f}s" - seed += 1 - repeat_no += 1 + print(f"New video saved to Path: "+video_path) + file_list.append(video_path) + seed += 1 if temp_filename!= None and os.path.isfile(temp_filename): os.remove(temp_filename) gen_in_progress = False offload.unload_loras_from_model(trans) - def get_new_preset_msg(advanced = True): if advanced: return "Enter here a Name for a Lora Preset or Choose one in the List" @@ -1952,16 +1995,40 @@ 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) output = gr.Gallery( label="Generated videos", show_label=False, elem_id="gallery" , columns=[3], rows=[1], object_fit="contain", height=450, selected_index=0, interactive= False) generate_btn = gr.Button("Generate") - onemore_btn = gr.Button("One More Please !", visible= False) - abort_btn = gr.Button("Abort") - gen_info = gr.Text(label="Current prompt", visible= False , interactive= False) - - + queue_df = gr.DataFrame( + headers=["ID", "Status", "Progress", "Steps", "Time", "Prompt", "", "", ""], + datatype=["str", "str", "str", "str", "str", "str", "str", "str", "str"], + interactive=False, + col_count=(9, "fixed"), + wrap=True, + value=update_queue_data, + every=1, + elem_id="queue_df" + ) + def handle_selection(evt: gr.SelectData): + cell_value = evt.value + selected_index = evt.index + if cell_value == "↑": + return move_up([selected_index]) + elif cell_value == "↓": + return move_down([selected_index]) + elif cell_value == "✖": + return remove_task([selected_index]) + return queue_df + selected_indices = gr.State([]) + queue_df.select( + fn=handle_selection, + outputs=selected_indices + ) + queue_df.change( + fn=refresh_gallery, + inputs=[state], + outputs=[output] + ) 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, 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, RIFLEx_setting, slg_switch, slg_layers, @@ -1978,48 +2045,43 @@ def generate_video_tab(image2video=False): refresh_lora_btn.click(refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) refresh_lora_btn2.click(refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) download_loras_btn.click(fn=download_loras, inputs=[], outputs=[download_status_row, download_status, presets_column, loras_column]).then(fn=refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) - gen_status.change(refresh_gallery, inputs = [state, gen_info], outputs = [output, gen_info] ) - abort_btn.click(abort_generation,state,abort_btn ) output.select(select_video, state, None ) - onemore_btn.click(fn=one_more_video,inputs=[state], outputs= [state]) - generate_btn.click(fn=prepare_generate_video,inputs=[], outputs= [generate_btn, onemore_btn] - ).then( - fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , outputs= [prompt] - ).then( - fn=generate_video, - inputs=[ - prompt, - negative_prompt, - resolution, - video_length, - seed, - num_inference_steps, - guidance_scale, - flow_shift, - embedded_guidance_scale, - repeat_generation, - multi_images_gen_type, - tea_cache_setting, - tea_cache_start_step_perc, - loras_choices, - loras_mult_choices, - image_to_continue, - image_to_end, - video_to_continue, - max_frames, - RIFLEx_setting, - slg_switch, - slg_layers, - slg_start_perc, - slg_end_perc, - state, - gr.State(image2video) - ], - outputs= [gen_status] - ).then( - finalize_gallery, - [state], - [output , abort_btn, generate_btn, onemore_btn, gen_info] + original_inputs = [ + prompt, + negative_prompt, + resolution, + video_length, + seed, + num_inference_steps, + guidance_scale, + flow_shift, + embedded_guidance_scale, + repeat_generation, + multi_images_gen_type, + tea_cache_setting, + tea_cache_start_step_perc, + loras_choices, + loras_mult_choices, + image_to_continue, + image_to_end, + video_to_continue, + max_frames, + RIFLEx_setting, + slg_switch, + slg_layers, + slg_start_perc, + slg_end_perc, + state, + gr.State(image2video) + ] + + #generate_btn.click( + # fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , outputs= [prompt] + #).then( + generate_btn.click( + fn=process_prompt_and_add_tasks, + inputs=original_inputs, + outputs=queue_df ) return loras_choices, lset_name, header, state @@ -2226,6 +2288,59 @@ def create_demo(): margin: 0 20px; white-space: nowrap; } + .queue-item { + border: 1px solid #ccc; + padding: 10px; + margin: 5px 0; + border-radius: 5px; + } + .current { + background: #f8f9fa; + border-left: 4px solid #007bff; + } + .task-header { + display: flex; + justify-content: space-between; + margin-bottom: 5px; + } + .progress-container { + height: 10px; + background: #e9ecef; + border-radius: 5px; + overflow: hidden; + } + .progress-bar { + height: 100%; + background: #007bff; + transition: width 0.3s ease; + } + .task-details { + display: flex; + justify-content: space-between; + font-size: 0.9em; + color: #6c757d; + margin-top: 5px; + } + .task-prompt { + font-size: 0.8em; + color: #868e96; + margin-top: 5px; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; + } + #queue_df td:nth-child(7), + #queue_df td:nth-child(8), + #queue_df td:nth-child(9) { + cursor: pointer; + text-align: center; + font-weight: bold; + } + #queue_df td:nth-child(7):hover, + #queue_df td:nth-child(8):hover, + #queue_df td:nth-child(9):hover { + background-color: #e0e0e0; + } """ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="sky", neutral_hue="slate", text_size="md")) as demo: gr.Markdown("

Wan 2.1GP v3.0 by DeepBeepMeep (Updates)

") @@ -2260,6 +2375,7 @@ def create_demo(): return demo if __name__ == "__main__": + threading.Thread(target=runner, daemon=True).start() os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" server_port = int(args.server_port) if os.name == "nt": From 646d9ace13db54def4e43a2dbf08b81083388c52 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 24 Mar 2025 21:00:30 +1100 Subject: [PATCH 02/69] restored original detailed status updates, fixed repeats --- gradio_server.py | 312 +++++++++++++++++++++++++---------------------- 1 file changed, 168 insertions(+), 144 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index d7037e4..2d551cb 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1,7 +1,6 @@ import os import time import threading -from typing import Dict, Any import argparse from mmgp import offload, safetensors2, profile_type try: @@ -23,7 +22,6 @@ import traceback import math import asyncio from wan.utils import prompt_parser -from typing import List PROMPT_VARS_MAX = 10 target_mmgp_version = "3.3.0" @@ -49,19 +47,23 @@ def runner(): with tracker_lock: progress = progress_tracker.get(task_id, {}) - if item['status'] == "Processing": + if item['state'] != "Queued" and item['state'] != "Finished": current_step = progress.get('current_step', 0) total_steps = progress.get('total_steps', 0) elapsed = time.time() - progress.get('start_time', time.time()) + status = progress.get('status', "") + state = progress.get("state") item.update({ 'progress': f"{((current_step/total_steps)*100 if total_steps > 0 else 0):.1f}%", 'steps': f"{current_step}/{total_steps}", - 'time': f"{elapsed:.1f}s" + 'time': f"{elapsed:.1f}s", + 'state': f"{state}", + 'status': f"{status}" }) - if not any(item['status'] == "Processing" for item in queue): + if not any(item['state'] == "Processing" for item in queue): for item in queue: - if item['status'] == "Queued": - item['status'] = "Processing" + if item['state'] == "Queued": + item['state'] = "Processing" current_task_id = item['id'] threading.Thread(target=process_task, args=(item,)).start() break @@ -157,7 +159,8 @@ def add_video_task(*params): queue.append({ "id": current_task_id, "params": (current_task_id,) + params, - "status": "Queued", + "state": "Queued", + "status": "0/0", "progress": "0.0%", "steps": f"0/{params[5]}", "time": "--", @@ -209,6 +212,7 @@ def update_queue_data(): for item in queue: data.append([ str(item['id']), + item['state'], item['status'], item.get('progress', "0.0%"), item.get('steps', ''), @@ -985,29 +989,38 @@ def save_video(final_frames, output_path, fps=24): final_frames = (final_frames * 255).astype(np.uint8) ImageSequenceClip(list(final_frames), fps=fps).write_videofile(output_path, verbose= False, logger = None) -def build_callback(task_id, total_steps): +def build_callback(state, pipe, num_inference_steps, status): start_time = time.time() - def update_progress(step, _): + def update_progress(step_idx, latents, read_state = False): with tracker_lock: + step_idx += 1 + if state.get("abort", False): + # pipe._interrupt = True + phase = "Aborting" + elif step_idx == num_inference_steps: + phase = "VAE Decoding" + else: + phase = "Denoising" elapsed = time.time() - start_time progress_tracker[task_id] = { - 'current_step': step + 1, - 'total_steps': total_steps, + 'current_step': step_idx, + 'total_steps': num_inference_steps, 'start_time': start_time, - 'last_update': time.time() + 'last_update': time.time(), + 'status': status, + 'state': phase } return update_progress def refresh_gallery(state): file_list = state.get("file_list", None) return file_list + def finalize_gallery(state): choice = 0 if "in_progress" in state: del state["in_progress"] choice = state.get("selected",0) - - state["extra_orders"] = 0 time.sleep(0.2) global gen_in_progress gen_in_progress = False @@ -1067,6 +1080,7 @@ def generate_video( offloadobj.release() offloadobj = None wan_model = None + gc.collect() print(f"Loading model {get_model_name(file_model_needed)}...") wan_model, offloadobj, trans = load_models(image2video) print(f"Model loaded") @@ -1243,143 +1257,153 @@ def generate_video( trans.teacache_skipped_steps = 0 trans.previous_residual_uncond = None trans.previous_residual_cond = None - callback = build_callback(task_id, num_inference_steps) - offload.shared_state["callback"] = callback + video_no = 0 + status = f"{video_no}/{repeat_generation}" + with tracker_lock: + if task_id in progress_tracker: + progress_tracker[task_id]['state'] = "Encoding Prompt" + progress_tracker[task_id]['status'] = status + callback = build_callback(state, trans, num_inference_steps, status) + offload.shared_state["callback"] = callback gc.collect() torch.cuda.empty_cache() wan_model._interrupt = False - state["progress_status"] = "Starting" - try: - if image2video: - samples = wan_model.generate( - prompt, - image_to_continue.convert('RGB'), - image_to_end.convert('RGB') if image_to_end != None else None, - frame_num=(video_length // 4)* 4 + 1, - max_area=MAX_AREA_CONFIGS[resolution], - shift=flow_shift, - sampling_steps=num_inference_steps, - guide_scale=guidance_scale, - n_prompt=negative_prompt, - seed=seed, - offload_model=False, - callback=callback, - enable_RIFLEx = enable_RIFLEx, - VAE_tile_size = VAE_tile_size, - joint_pass = joint_pass, - slg_layers = slg_layers, - slg_start = slg_start/100, - slg_end = slg_end/100, - ) + for i in range(repeat_generation): + try: + with tracker_lock: + if task_id in progress_tracker: + progress_tracker[task_id]['status'] = video_no + video_no += 1 + if image2video: + samples = wan_model.generate( + prompt, + image_to_continue.convert('RGB'), + image_to_end.convert('RGB') if image_to_end != None else None, + frame_num=(video_length // 4)* 4 + 1, + max_area=MAX_AREA_CONFIGS[resolution], + shift=flow_shift, + sampling_steps=num_inference_steps, + guide_scale=guidance_scale, + n_prompt=negative_prompt, + seed=seed, + offload_model=False, + callback=callback, + enable_RIFLEx = enable_RIFLEx, + VAE_tile_size = VAE_tile_size, + joint_pass = joint_pass, + slg_layers = slg_layers, + slg_start = slg_start/100, + slg_end = slg_end/100, + ) - else: - samples = wan_model.generate( - prompt, - frame_num=(video_length // 4)* 4 + 1, - size=(width, height), - shift=flow_shift, - sampling_steps=num_inference_steps, - guide_scale=guidance_scale, - n_prompt=negative_prompt, - seed=seed, - offload_model=False, - callback=callback, - enable_RIFLEx = enable_RIFLEx, - VAE_tile_size = VAE_tile_size, - joint_pass = joint_pass, - slg_layers = slg_layers, - slg_start = slg_start/100, - slg_end = slg_end/100, - ) - except Exception as e: - gen_in_progress = False - if temp_filename!= None and os.path.isfile(temp_filename): - os.remove(temp_filename) + else: + samples = wan_model.generate( + prompt, + frame_num=(video_length // 4)* 4 + 1, + size=(width, height), + shift=flow_shift, + sampling_steps=num_inference_steps, + guide_scale=guidance_scale, + n_prompt=negative_prompt, + seed=seed, + offload_model=False, + callback=callback, + enable_RIFLEx = enable_RIFLEx, + VAE_tile_size = VAE_tile_size, + joint_pass = joint_pass, + slg_layers = slg_layers, + slg_start = slg_start/100, + slg_end = slg_end/100, + ) + except Exception as e: + gen_in_progress = False + if temp_filename!= None and os.path.isfile(temp_filename): + os.remove(temp_filename) + offload.last_offload_obj.unload_all() + offload.unload_loras_from_model(trans) + # if compile: + # cache_size = torch._dynamo.config.cache_size_limit + # torch.compiler.reset() + # torch._dynamo.config.cache_size_limit = cache_size + + gc.collect() + torch.cuda.empty_cache() + s = str(e) + keyword_list = ["vram", "VRAM", "memory","allocat"] + VRAM_crash= False + if any( keyword in s for keyword in keyword_list): + VRAM_crash = True + else: + stack = traceback.extract_stack(f=None, limit=5) + for frame in stack: + if any( keyword in frame.name for keyword in keyword_list): + VRAM_crash = True + break + state["prompt"] = "" + if VRAM_crash: + raise gr.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: + raise gr.Error(f"The generation of the video has encountered an error, please check your terminal for more information. '{s}'") + finally: + with tracker_lock: + if task_id in progress_tracker: + del progress_tracker[task_id] + + if trans.enable_teacache: + print(f"Teacache Skipped Steps:{trans.teacache_skipped_steps}/{num_inference_steps}" ) + trans.previous_residual_uncond = None + trans.previous_residual_cond = None + + if samples != None: + samples = samples.to("cpu") offload.last_offload_obj.unload_all() - offload.unload_loras_from_model(trans) - # if compile: - # cache_size = torch._dynamo.config.cache_size_limit - # torch.compiler.reset() - # torch._dynamo.config.cache_size_limit = cache_size - gc.collect() torch.cuda.empty_cache() - s = str(e) - keyword_list = ["vram", "VRAM", "memory","allocat"] - VRAM_crash= False - if any( keyword in s for keyword in keyword_list): - VRAM_crash = True + + if samples == None: + end_time = time.time() + abort = True + state["prompt"] = "" else: - stack = traceback.extract_stack(f=None, limit=5) - for frame in stack: - if any( keyword in frame.name for keyword in keyword_list): - VRAM_crash = True - break - state["prompt"] = "" - if VRAM_crash: - raise gr.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: - raise gr.Error(f"The generation of the video has encountered an error, please check your terminal for more information. '{s}'") - finally: - with tracker_lock: - if task_id in progress_tracker: - del progress_tracker[task_id] + sample = samples.cpu() + # video = rearrange(sample.cpu().numpy(), "c t h w -> t h w c") - if trans.enable_teacache: - print(f"Teacache Skipped Steps:{trans.teacache_skipped_steps}/{num_inference_steps}" ) - trans.previous_residual_uncond = None - trans.previous_residual_cond = None + time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%Hh%Mm%Ss") + if os.name == 'nt': + file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:50]).strip()}.mp4" + else: + file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:100]).strip()}.mp4" + video_path = os.path.join(save_path, file_name) + cache_video( + tensor=sample[None], + save_file=video_path, + fps=16, + nrow=1, + normalize=True, + value_range=(-1, 1)) + + configs = { + 'prompt': prompt, + 'negative_prompt': negative_prompt, + 'resolution': resolution, + 'video_length': video_length, + 'seed': seed, + 'num_inference_steps': num_inference_steps, + } - if samples != None: - samples = samples.to("cpu") - offload.last_offload_obj.unload_all() - gc.collect() - torch.cuda.empty_cache() + metadata_choice = server_config.get("metadata_choice","metadata") + if metadata_choice == "json": + with open(video_path.replace('.mp4', '.json'), 'w') as f: + json.dump(configs, f, indent=4) + elif metadata_choice == "metadata": + from mutagen.mp4 import MP4 + file = MP4(video_path) + file.tags['©cmt'] = [json.dumps(configs)] + file.save() - if samples == None: - end_time = time.time() - abort = True - state["prompt"] = "" - else: - sample = samples.cpu() - # video = rearrange(sample.cpu().numpy(), "c t h w -> t h w c") - - time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%Hh%Mm%Ss") - if os.name == 'nt': - file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:50]).strip()}.mp4" - else: - file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:100]).strip()}.mp4" - video_path = os.path.join(save_path, file_name) - cache_video( - tensor=sample[None], - save_file=video_path, - fps=16, - nrow=1, - normalize=True, - value_range=(-1, 1)) - - configs = { - 'prompt': prompt, - 'negative_prompt': negative_prompt, - 'resolution': resolution, - 'video_length': video_length, - 'seed': seed, - 'num_inference_steps': num_inference_steps, - } - - metadata_choice = server_config.get("metadata_choice","metadata") - if metadata_choice == "json": - with open(video_path.replace('.mp4', '.json'), 'w') as f: - json.dump(configs, f, indent=4) - elif metadata_choice == "metadata": - from mutagen.mp4 import MP4 - file = MP4(video_path) - file.tags['©cmt'] = [json.dumps(configs)] - file.save() - - print(f"New video saved to Path: "+video_path) - file_list.append(video_path) - seed += 1 + print(f"New video saved to Path: "+video_path) + file_list.append(video_path) + seed += 1 if temp_filename!= None and os.path.isfile(temp_filename): os.remove(temp_filename) @@ -2006,10 +2030,10 @@ def generate_video_tab(image2video=False): , columns=[3], rows=[1], object_fit="contain", height=450, selected_index=0, interactive= False) generate_btn = gr.Button("Generate") queue_df = gr.DataFrame( - headers=["ID", "Status", "Progress", "Steps", "Time", "Prompt", "", "", ""], - datatype=["str", "str", "str", "str", "str", "str", "str", "str", "str"], + headers=["ID", "Status", "Repeats", "Progress", "Steps", "Time", "Prompt", "", "", ""], + datatype=["str", "str", "str", "str", "str", "str", "str", "str", "str", "str"], interactive=False, - col_count=(9, "fixed"), + col_count=(10, "fixed"), wrap=True, value=update_queue_data, every=1, From 347ab55d4aea645997056211634b3c8663dfa120 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 24 Mar 2025 22:23:51 +1100 Subject: [PATCH 03/69] add config option to set whether to reload model upon changing tabs or pressing generate --- gradio_server.py | 56 ++++++++++++++++++++++++++++++++++-------------- 1 file changed, 40 insertions(+), 16 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 2d551cb..2c71042 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -491,7 +491,8 @@ if not Path(server_config_filename).is_file(): "default_ui": "t2v", "boost" : 1, "vae_config": 0, - "profile" : profile_type.LowRAM_LowVRAM } + "profile" : profile_type.LowRAM_LowVRAM, + "reload_model": 1 } with open(server_config_filename, "w", encoding="utf-8") as writer: writer.write(json.dumps(server_config)) @@ -915,7 +916,8 @@ def apply_changes( state, vae_config_choice, metadata_choice, default_ui_choice ="t2v", - boost_choice = 1 + boost_choice = 1, + reload_choice = 1 ): if args.lock_config: return @@ -934,6 +936,7 @@ def apply_changes( state, "metadata_choice": metadata_choice, "default_ui" : default_ui_choice, "boost" : boost_choice, + "reload_model" : reload_choice, } if Path(server_config_filename).is_file(): @@ -1075,16 +1078,17 @@ def generate_video( global wan_model, offloadobj reload_needed = state.get("_reload_needed", False) file_model_needed = model_needed(image2video) - if file_model_needed != model_filename or reload_needed: - if offloadobj is not None: - offloadobj.release() - offloadobj = None - wan_model = None - gc.collect() - print(f"Loading model {get_model_name(file_model_needed)}...") - wan_model, offloadobj, trans = load_models(image2video) - print(f"Model loaded") - state["_reload_needed"] = False + if(server_config.get("reload_model",1) == 2): + if file_model_needed != model_filename or reload_needed: + if offloadobj is not None: + offloadobj.release() + offloadobj = None + wan_model = None + gc.collect() + print(f"Loading model {get_model_name(file_model_needed)}...") + wan_model, offloadobj, trans = load_models(image2video) + print(f"Model loaded") + state["_reload_needed"] = False from PIL import Image import numpy as np @@ -1190,7 +1194,7 @@ def generate_video( if len(list_mult_choices_nums ) < len(loras_choices): list_mult_choices_nums += [1.0] * ( len(loras_choices) - len(list_mult_choices_nums ) ) loras_selected = [ lora for i, lora in enumerate(loras) if str(i) in loras_choices] - pinnedLora = profile !=5 #False # # # + pinnedLora = False !=5 #False # # # offload.load_loras_into_model(trans, loras_selected, list_mult_choices_nums, activate_all_loras=True, preprocess_sd=preprocess_loras, pinnedLora=pinnedLora, split_linear_modules_map = None) errors = trans._loras_errors if len(errors) > 0: @@ -2236,6 +2240,14 @@ def generate_configuration_tab(): value=metadata, label="Metadata Handling" ) + reload_choice = gr.Dropdown( + choices=[ + ("When changing tabs", 1), + ("When pressing generate", 2), + ], + value=server_config.get("reload_model",1), + label="Reload model" + ) msg = gr.Markdown() apply_btn = gr.Button("Apply Changes") apply_btn.click( @@ -2253,6 +2265,7 @@ def generate_configuration_tab(): metadata_choice, default_ui_choice, boost_choice, + reload_choice, ], outputs= msg ) @@ -2270,13 +2283,24 @@ def generate_about_tab(): def on_tab_select(t2v_state, i2v_state, evt: gr.SelectData): global lora_model_filename, use_image2video - t2v_header = generate_header(transformer_filename_t2v, compile, attention_mode) - i2v_header = generate_header(transformer_filename_i2v, compile, attention_mode) - new_t2v = evt.index == 0 new_i2v = evt.index == 1 use_image2video = new_i2v + if(server_config.get("reload_model",1) == 1): + global wan_model, offloadobj + if wan_model is not None: + if offloadobj is not None: + offloadobj.release() + offloadobj = None + wan_model = None + gc.collect() + torch.cuda.empty_cache() + wan_model, offloadobj, trans = load_models(use_image2video) + + t2v_header = generate_header(transformer_filename_t2v, compile, attention_mode) + i2v_header = generate_header(transformer_filename_i2v, compile, attention_mode) + if new_t2v: lora_model_filename = t2v_state["loras_model"] if ("1.3B" in transformer_filename_t2v and not "1.3B" in lora_model_filename or "14B" in transformer_filename_t2v and not "14B" in lora_model_filename): From 1b125fe2331e884297aa724ffc971e82e137afc2 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 24 Mar 2025 23:18:02 +1100 Subject: [PATCH 04/69] merge latest vram fixes --- gradio_server.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 2c71042..9671334 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -854,6 +854,8 @@ def load_models(i2v): kwargs = { "extraModelsToQuantize": None} if profile == 2 or profile == 4: kwargs["budgets"] = { "transformer" : 100 if preload == 0 else preload, "text_encoder" : 100, "*" : 1000 } + if profile == 4: + kwargs["partialPinning"] = True elif profile == 3: kwargs["budgets"] = { "*" : "70%" } offloadobj = offload.profile(pipe, profile_no= profile, compile = compile, quantizeTransformer = quantizeTransformer, loras = "transformer", **kwargs) @@ -864,7 +866,7 @@ wan_model, offloadobj, transformer = load_models(use_image2video) if check_loras: setup_loras(use_image2video, transformer, get_lora_dir(use_image2video), "", None) exit() - +del transformer gen_in_progress = False def get_auto_attention(): @@ -1080,10 +1082,10 @@ def generate_video( file_model_needed = model_needed(image2video) if(server_config.get("reload_model",1) == 2): if file_model_needed != model_filename or reload_needed: + del wan_model if offloadobj is not None: offloadobj.release() - offloadobj = None - wan_model = None + del offloadobj gc.collect() print(f"Loading model {get_model_name(file_model_needed)}...") wan_model, offloadobj, trans = load_models(image2video) @@ -1194,7 +1196,7 @@ def generate_video( if len(list_mult_choices_nums ) < len(loras_choices): list_mult_choices_nums += [1.0] * ( len(loras_choices) - len(list_mult_choices_nums ) ) loras_selected = [ lora for i, lora in enumerate(loras) if str(i) in loras_choices] - pinnedLora = False !=5 #False # # # + pinnedLora = False #profile !=5 #False # # # offload.load_loras_into_model(trans, loras_selected, list_mult_choices_nums, activate_all_loras=True, preprocess_sd=preprocess_loras, pinnedLora=pinnedLora, split_linear_modules_map = None) errors = trans._loras_errors if len(errors) > 0: From e686952a84303e3e9c0c07f3fee1a0c43c33554a Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 24 Mar 2025 23:38:12 +1100 Subject: [PATCH 05/69] set default model loading to JIT --- gradio_server.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 9671334..9c4b9b9 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -492,7 +492,7 @@ if not Path(server_config_filename).is_file(): "boost" : 1, "vae_config": 0, "profile" : profile_type.LowRAM_LowVRAM, - "reload_model": 1 } + "reload_model": 2 } with open(server_config_filename, "w", encoding="utf-8") as writer: writer.write(json.dumps(server_config)) @@ -854,8 +854,8 @@ def load_models(i2v): kwargs = { "extraModelsToQuantize": None} if profile == 2 or profile == 4: kwargs["budgets"] = { "transformer" : 100 if preload == 0 else preload, "text_encoder" : 100, "*" : 1000 } - if profile == 4: - kwargs["partialPinning"] = True + # if profile == 4: + # kwargs["partialPinning"] = True elif profile == 3: kwargs["budgets"] = { "*" : "70%" } offloadobj = offload.profile(pipe, profile_no= profile, compile = compile, quantizeTransformer = quantizeTransformer, loras = "transformer", **kwargs) @@ -867,6 +867,7 @@ if check_loras: setup_loras(use_image2video, transformer, get_lora_dir(use_image2video), "", None) exit() del transformer + gen_in_progress = False def get_auto_attention(): @@ -2247,7 +2248,7 @@ def generate_configuration_tab(): ("When changing tabs", 1), ("When pressing generate", 2), ], - value=server_config.get("reload_model",1), + value=server_config.get("reload_model",2), label="Reload model" ) msg = gr.Markdown() @@ -2289,7 +2290,7 @@ def on_tab_select(t2v_state, i2v_state, evt: gr.SelectData): new_i2v = evt.index == 1 use_image2video = new_i2v - if(server_config.get("reload_model",1) == 1): + if(server_config.get("reload_model",2) == 1): global wan_model, offloadobj if wan_model is not None: if offloadobj is not None: From b5eca59a71a700458edc095832ca89dc644de2f1 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 24 Mar 2025 23:41:37 +1100 Subject: [PATCH 06/69] merge conflicts --- gradio_server.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 9c4b9b9..b3ccd94 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1086,7 +1086,7 @@ def generate_video( del wan_model if offloadobj is not None: offloadobj.release() - del offloadobj + del offloadobj gc.collect() print(f"Loading model {get_model_name(file_model_needed)}...") wan_model, offloadobj, trans = load_models(image2video) @@ -1197,8 +1197,8 @@ def generate_video( if len(list_mult_choices_nums ) < len(loras_choices): list_mult_choices_nums += [1.0] * ( len(loras_choices) - len(list_mult_choices_nums ) ) loras_selected = [ lora for i, lora in enumerate(loras) if str(i) in loras_choices] - pinnedLora = False #profile !=5 #False # # # - offload.load_loras_into_model(trans, loras_selected, list_mult_choices_nums, activate_all_loras=True, preprocess_sd=preprocess_loras, pinnedLora=pinnedLora, split_linear_modules_map = None) + pinnedLora = False #profile !=5 #False # # # + offload.load_loras_into_model(trans, loras_selected, list_mult_choices_nums, activate_all_loras=True, preprocess_sd=preprocess_loras, pinnedLora=pinnedLora, split_linear_modules_map = None) errors = trans._loras_errors if len(errors) > 0: error_files = [msg for _ , msg in errors] From 3cd0dbf4ddada5e4e2baedfc31f3facdd8a58520 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 25 Mar 2025 03:07:35 +1100 Subject: [PATCH 07/69] fix broken queue states, avoid unnecessary reloading --- gradio_server.py | 86 ++++++++++++++++++++++++++---------------------- 1 file changed, 46 insertions(+), 40 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index b3ccd94..0a0fafb 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -37,6 +37,7 @@ task_id = 0 progress_tracker = {} tracker_lock = threading.Lock() file_list = [] +last_model_type = None def runner(): global current_task_id @@ -47,22 +48,23 @@ def runner(): with tracker_lock: progress = progress_tracker.get(task_id, {}) - if item['state'] != "Queued" and item['state'] != "Finished": + if item['state'] == "Processing": current_step = progress.get('current_step', 0) total_steps = progress.get('total_steps', 0) elapsed = time.time() - progress.get('start_time', time.time()) status = progress.get('status', "") - state = progress.get("state") + repeats = progress.get("repeats") item.update({ 'progress': f"{((current_step/total_steps)*100 if total_steps > 0 else 0):.1f}%", 'steps': f"{current_step}/{total_steps}", 'time': f"{elapsed:.1f}s", - 'state': f"{state}", + 'repeats': f"{repeats}", 'status': f"{status}" }) if not any(item['state'] == "Processing" for item in queue): for item in queue: if item['state'] == "Queued": + item['status'] = "Processing" item['state'] = "Processing" current_task_id = item['id'] threading.Thread(target=process_task, args=(item,)).start() @@ -160,7 +162,8 @@ def add_video_task(*params): "id": current_task_id, "params": (current_task_id,) + params, "state": "Queued", - "status": "0/0", + "status": "Queued", + "repeats": "0/0", "progress": "0.0%", "steps": f"0/{params[5]}", "time": "--", @@ -212,8 +215,8 @@ def update_queue_data(): for item in queue: data.append([ str(item['id']), - item['state'], item['status'], + item['repeats'], item.get('progress', "0.0%"), item.get('steps', ''), item.get('time', '--'), @@ -1013,8 +1016,8 @@ def build_callback(state, pipe, num_inference_steps, status): 'total_steps': num_inference_steps, 'start_time': start_time, 'last_update': time.time(), - 'status': status, - 'state': phase + 'repeats': status, + 'status': phase } return update_progress @@ -1078,20 +1081,21 @@ def generate_video( progress=gr.Progress() #track_tqdm= True ): - global wan_model, offloadobj + global wan_model, offloadobj, last_model_type reload_needed = state.get("_reload_needed", False) file_model_needed = model_needed(image2video) - if(server_config.get("reload_model",1) == 2): - if file_model_needed != model_filename or reload_needed: - del wan_model - if offloadobj is not None: - offloadobj.release() - del offloadobj - gc.collect() - print(f"Loading model {get_model_name(file_model_needed)}...") - wan_model, offloadobj, trans = load_models(image2video) - print(f"Model loaded") - state["_reload_needed"] = False + with lock: + queue_not_empty = len(queue) > 0 + if(last_model_type != image2video and (queue_not_empty or server_config.get("reload_model",1) == 2) and (file_model_needed != model_filename or reload_needed)): + del wan_model + if offloadobj is not None: + offloadobj.release() + del offloadobj + gc.collect() + print(f"Loading model {get_model_name(file_model_needed)}...") + wan_model, offloadobj, trans = load_models(image2video) + print(f"Model loaded") + state["_reload_needed"] = False from PIL import Image import numpy as np @@ -1251,13 +1255,6 @@ def generate_video( global save_path os.makedirs(save_path, exist_ok=True) abort = False - with tracker_lock: - progress_tracker[task_id] = { - 'current_step': 0, - 'total_steps': num_inference_steps, - 'start_time': time.time(), - 'last_update': time.time() - } if trans.enable_teacache: trans.teacache_counter = 0 trans.num_steps = num_inference_steps @@ -1268,8 +1265,12 @@ def generate_video( status = f"{video_no}/{repeat_generation}" with tracker_lock: if task_id in progress_tracker: - progress_tracker[task_id]['state'] = "Encoding Prompt" - progress_tracker[task_id]['status'] = status + progress_tracker[task_id]['status'] = "Encoding Prompt" + progress_tracker[task_id]['repeats'] = status + progress_tracker[task_id]['current_step'] = 0 + progress_tracker[task_id]['total_steps'] = num_inference_steps + progress_tracker[task_id]['start_time'] = time.time() + progress_tracker[task_id]['last_update'] = time.time() callback = build_callback(state, trans, num_inference_steps, status) offload.shared_state["callback"] = callback gc.collect() @@ -1279,7 +1280,7 @@ def generate_video( try: with tracker_lock: if task_id in progress_tracker: - progress_tracker[task_id]['status'] = video_no + progress_tracker[task_id]['repeats'] = video_no video_no += 1 if image2video: samples = wan_model.generate( @@ -1326,8 +1327,8 @@ def generate_video( gen_in_progress = False if temp_filename!= None and os.path.isfile(temp_filename): os.remove(temp_filename) - offload.last_offload_obj.unload_all() - offload.unload_loras_from_model(trans) + if(offload.last_offload_obj): offload.last_offload_obj.unload_all() + if(trans): offload.unload_loras_from_model(trans) # if compile: # cache_size = torch._dynamo.config.cache_size_limit # torch.compiler.reset() @@ -1411,6 +1412,7 @@ def generate_video( print(f"New video saved to Path: "+video_path) file_list.append(video_path) seed += 1 + last_model_type = image2video if temp_filename!= None and os.path.isfile(temp_filename): os.remove(temp_filename) @@ -2291,15 +2293,19 @@ def on_tab_select(t2v_state, i2v_state, evt: gr.SelectData): use_image2video = new_i2v if(server_config.get("reload_model",2) == 1): - global wan_model, offloadobj - if wan_model is not None: - if offloadobj is not None: - offloadobj.release() - offloadobj = None - wan_model = None - gc.collect() - torch.cuda.empty_cache() - wan_model, offloadobj, trans = load_models(use_image2video) + with lock: + queue_empty = len(queue) == 0 + if queue_empty: + global wan_model, offloadobj + if wan_model is not None: + if offloadobj is not None: + offloadobj.release() + offloadobj = None + wan_model = None + gc.collect() + torch.cuda.empty_cache() + wan_model, offloadobj, trans = load_models(use_image2video) + del trans t2v_header = generate_header(transformer_filename_t2v, compile, attention_mode) i2v_header = generate_header(transformer_filename_i2v, compile, attention_mode) From 3d9f4c4326982f88b160ce9478de3d38638dd264 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 25 Mar 2025 15:59:42 +1100 Subject: [PATCH 08/69] fix broken queue states (again) --- gradio_server.py | 32 +++++++++++++++----------------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 0a0fafb..59d3096 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -215,8 +215,8 @@ def update_queue_data(): for item in queue: data.append([ str(item['id']), - item['status'], - item['repeats'], + item.get('status', "Starting"), + item.get('repeats', "0/0"), item.get('progress', "0.0%"), item.get('steps', ''), item.get('time', '--'), @@ -998,7 +998,7 @@ def save_video(final_frames, output_path, fps=24): final_frames = (final_frames * 255).astype(np.uint8) ImageSequenceClip(list(final_frames), fps=fps).write_videofile(output_path, verbose= False, logger = None) -def build_callback(state, pipe, num_inference_steps, status): +def build_callback(taskid, state, pipe, num_inference_steps, repeats): start_time = time.time() def update_progress(step_idx, latents, read_state = False): with tracker_lock: @@ -1011,12 +1011,12 @@ def build_callback(state, pipe, num_inference_steps, status): else: phase = "Denoising" elapsed = time.time() - start_time - progress_tracker[task_id] = { + progress_tracker[taskid] = { 'current_step': step_idx, 'total_steps': num_inference_steps, 'start_time': start_time, 'last_update': time.time(), - 'repeats': status, + 'repeats': repeats, 'status': phase } return update_progress @@ -1262,16 +1262,8 @@ def generate_video( trans.previous_residual_uncond = None trans.previous_residual_cond = None video_no = 0 - status = f"{video_no}/{repeat_generation}" - with tracker_lock: - if task_id in progress_tracker: - progress_tracker[task_id]['status'] = "Encoding Prompt" - progress_tracker[task_id]['repeats'] = status - progress_tracker[task_id]['current_step'] = 0 - progress_tracker[task_id]['total_steps'] = num_inference_steps - progress_tracker[task_id]['start_time'] = time.time() - progress_tracker[task_id]['last_update'] = time.time() - callback = build_callback(state, trans, num_inference_steps, status) + repeats = f"{video_no}/{repeat_generation}" + callback = build_callback(task_id, state, trans, num_inference_steps, repeats) offload.shared_state["callback"] = callback gc.collect() torch.cuda.empty_cache() @@ -1279,8 +1271,14 @@ def generate_video( for i in range(repeat_generation): try: with tracker_lock: - if task_id in progress_tracker: - progress_tracker[task_id]['repeats'] = video_no + progress_tracker[task_id] = { + 'current_step': 0, + 'total_steps': num_inference_steps, + 'start_time': time.time(), + 'last_update': time.time(), + 'repeats': f"0/{repeat_generation}", + 'status': "Encoding Prompt" + } video_no += 1 if image2video: samples = wan_model.generate( From 306229c17f5a73e82e795a22109f444179606a34 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 25 Mar 2025 18:00:32 +1100 Subject: [PATCH 09/69] improve styling --- gradio_server.py | 33 ++++++++++++++++++++++++++------- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 59d3096..b695451 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -24,7 +24,7 @@ import asyncio from wan.utils import prompt_parser PROMPT_VARS_MAX = 10 -target_mmgp_version = "3.3.0" +target_mmgp_version = "3.3.1" from importlib.metadata import version mmgp_version = version("mmgp") if mmgp_version != target_mmgp_version: @@ -39,6 +39,17 @@ tracker_lock = threading.Lock() file_list = [] last_model_type = None +def format_time(seconds): + if seconds < 60: + return f"{seconds:.1f}s" + elif seconds < 3600: + minutes = seconds / 60 + return f"{minutes:.1f}m" + else: + hours = int(seconds // 3600) + minutes = int((seconds % 3600) // 60) + return f"{hours}h {minutes}m" + def runner(): global current_task_id while True: @@ -57,7 +68,7 @@ def runner(): item.update({ 'progress': f"{((current_step/total_steps)*100 if total_steps > 0 else 0):.1f}%", 'steps': f"{current_step}/{total_steps}", - 'time': f"{elapsed:.1f}s", + 'time': format_time(elapsed), 'repeats': f"{repeats}", 'status': f"{status}" }) @@ -213,14 +224,16 @@ def update_queue_data(): with lock: data = [] for item in queue: + truncated_prompt = (item['prompt'][:97] + '...') if len(item['prompt']) > 100 else item['prompt'] + full_prompt = item['prompt'].replace('"', '"') + prompt_cell = f'{truncated_prompt}' data.append([ - str(item['id']), item.get('status', "Starting"), item.get('repeats', "0/0"), item.get('progress', "0.0%"), item.get('steps', ''), item.get('time', '--'), - (item['prompt'][:47] + '...') if len(item['prompt']) > 50 else item['prompt'], + prompt_cell, "↑", "↓", "✖" @@ -2037,10 +2050,10 @@ def generate_video_tab(image2video=False): , columns=[3], rows=[1], object_fit="contain", height=450, selected_index=0, interactive= False) generate_btn = gr.Button("Generate") queue_df = gr.DataFrame( - headers=["ID", "Status", "Repeats", "Progress", "Steps", "Time", "Prompt", "", "", ""], - datatype=["str", "str", "str", "str", "str", "str", "str", "str", "str", "str"], + headers=["Status", "Completed", "Progress", "Steps", "Time", "Prompt", "", "", ""], + datatype=["str", "str", "str", "str", "str", "markdown", "str", "str", "str"], interactive=False, - col_count=(10, "fixed"), + col_count=(9, "fixed"), wrap=True, value=update_queue_data, every=1, @@ -2390,6 +2403,12 @@ def create_demo(): overflow: hidden; text-overflow: ellipsis; } + #queue_df td:nth-child(1) { + width: 100px; + } + #queue_df td:nth-child(6) { + width: 300px; + } #queue_df td:nth-child(7), #queue_df td:nth-child(8), #queue_df td:nth-child(9) { From c7a8c82398d16862e1d709924eee169a51eb93ae Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 25 Mar 2025 18:28:35 +1100 Subject: [PATCH 10/69] improve styling 2 --- gradio_server.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/gradio_server.py b/gradio_server.py index b695451..bcc8dda 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -2403,6 +2403,18 @@ def create_demo(): overflow: hidden; text-overflow: ellipsis; } + #queue_df table { + overflow: hidden !important; + } + + #queue_df::-webkit-scrollbar { + display: none !important; + } + + #queue_df { + scrollbar-width: none !important; + -ms-overflow-style: none !important; + } #queue_df td:nth-child(1) { width: 100px; } From 7e7c9c559376e000ee9966bbc6a74f5f04c3a5c4 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 25 Mar 2025 19:04:04 +1100 Subject: [PATCH 11/69] merge with latest branch --- README.md | 9 ++-- gradio_server.py | 110 +++++++++++++++++++++++--------------- requirements.txt | 2 +- wan/image2video.py | 2 +- wan/modules/attention.py | 14 ++++- wan/modules/clip.py | 7 ++- wan/modules/sage2_core.py | 9 ++++ wan/modules/t5.py | 2 +- wan/modules/vae.py | 9 ++-- wan/text2video.py | 2 +- 10 files changed, 108 insertions(+), 58 deletions(-) diff --git a/README.md b/README.md index c69848b..e37c10d 100644 --- a/README.md +++ b/README.md @@ -19,11 +19,13 @@ In this repository, we present **Wan2.1**, a comprehensive and open suite of vid ## 🔥 Latest News!! +* Mar 19 2022: 👋 Wan2.1GP v3.1: Faster launch and RAM optimizations (should require less RAM to run)\ + You will need one more *pip install -r requirements.txt* * Mar 18 2022: 👋 Wan2.1GP v3.0: - New Tab based interface, yon can switch from i2v to t2v conversely without restarting the app - Experimental Dual Frames mode for i2v, you can also specify an End frame. It doesn't always work, so you will need a few attempts. - You can save default settings in the files *i2v_settings.json* and *t2v_settings.json* that will be used when launching the app (you can also specify the path to different settings files) - - Slight acceleration with loras + - Slight acceleration with loras\ You will need one more *pip install -r requirements.txt* Many thanks to *Tophness* who created the framework (and did a big part of the work) of the multitabs and saved settings features * Mar 18 2022: 👋 Wan2.1GP v2.11: Added more command line parameters to prefill the generation settings + customizable output directory and choice of type of metadata for generated videos. Many thanks to *Tophness* for his contributions. You will need one more *pip install -r requirements.txt* to reflect new dependencies\ @@ -255,8 +257,9 @@ You can define multiple lines of macros. If there is only one macro line, the ap --check-loras : filter loras that are incompatible (will take a few seconds while refreshing the lora list or while starting the app)\ --advanced : turn on the advanced mode while launching the app\ --i2v-settings : path to launch settings for i2v\ ---t2v-settings : path to launch settings for t2v ---listen : make server accessible on network +--t2v-settings : path to launch settings for t2v\ +--listen : make server accessible on network\ +--gpu device : run Wan on device for instance "cuda:1" ### Profiles (for power users only) You can choose between 5 profiles, but two are really relevant here : diff --git a/gradio_server.py b/gradio_server.py index bcc8dda..27d88c5 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -15,7 +15,7 @@ import json import wan from wan.configs import MAX_AREA_CONFIGS, WAN_CONFIGS, SUPPORTED_SIZES from wan.utils.utils import cache_video -from wan.modules.attention import get_attention_modes +from wan.modules.attention import get_attention_modes, get_supported_attention_modes import torch import gc import traceback @@ -24,7 +24,7 @@ import asyncio from wan.utils import prompt_parser PROMPT_VARS_MAX = 10 -target_mmgp_version = "3.3.1" +target_mmgp_version = "3.3.3" from importlib.metadata import version mmgp_version = version("mmgp") if mmgp_version != target_mmgp_version: @@ -55,16 +55,16 @@ def runner(): while True: with lock: for item in queue: - task_id = item['id'] + task_id_runner = item['id'] with tracker_lock: - progress = progress_tracker.get(task_id, {}) + progress = progress_tracker.get(task_id_runner, {}) if item['state'] == "Processing": current_step = progress.get('current_step', 0) total_steps = progress.get('total_steps', 0) elapsed = time.time() - progress.get('start_time', time.time()) status = progress.get('status', "") - repeats = progress.get("repeats") + repeats = progress.get("repeats", "0/0") item.update({ 'progress': f"{((current_step/total_steps)*100 if total_steps > 0 else 0):.1f}%", 'steps': f"{current_step}/{total_steps}", @@ -381,6 +381,13 @@ def _parse_args(): help="Server name" ) + parser.add_argument( + "--gpu", + type=str, + default="", + help="Default GPU Device" + ) + parser.add_argument( "--open-browser", action="store_true", @@ -473,7 +480,8 @@ def get_lora_dir(i2v): return lora_dir_14B return root_lora_dir -attention_modes_supported = get_attention_modes() +attention_modes_installed = get_attention_modes() +attention_modes_supported = get_supported_attention_modes() args = _parse_args() args.flow_reverse = True @@ -587,6 +595,7 @@ vae_config = server_config.get("vae_config", 0) if len(args.vae_config) > 0: vae_config = int(args.vae_config) +reload_needed = False default_ui = server_config.get("default_ui", "t2v") metadata = server_config.get("metadata_type", "metadata") save_path = server_config.get("save_path", os.path.join(os.getcwd(), "gradio_outputs")) @@ -686,7 +695,7 @@ def download_models(transformer_filename, text_encoder_filename): from huggingface_hub import hf_hub_download, snapshot_download repoId = "DeepBeepMeep/Wan2.1" sourceFolderList = ["xlm-roberta-large", "", ] - fileList = [ [], ["Wan2.1_VAE.pth", "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ] + computeList(text_encoder_filename) + computeList(transformer_filename) ] + fileList = [ [], ["Wan2.1_VAE_bf16.safetensors", "models_clip_open-clip-xlm-roberta-large-vit-huge-14-bf16.safetensors" ] + computeList(text_encoder_filename) + computeList(transformer_filename) ] targetRoot = "ckpts/" for sourceFolder, files in zip(sourceFolderList,fileList ): if len(files)==0: @@ -703,6 +712,14 @@ def download_models(transformer_filename, text_encoder_filename): offload.default_verboseLevel = verbose_level +to_remove = ["models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", "Wan2.1_VAE.pth"] +for file_name in to_remove: + file_name = os.path.join("ckpts",file_name) + if os.path.isfile(file_name): + try: + os.remove(file_name) + except: + pass download_models(transformer_filename_i2v if use_image2video else transformer_filename_t2v, text_encoder_filename) @@ -875,6 +892,8 @@ def load_models(i2v): elif profile == 3: kwargs["budgets"] = { "*" : "70%" } offloadobj = offload.profile(pipe, profile_no= profile, compile = compile, quantizeTransformer = quantizeTransformer, loras = "transformer", **kwargs) + if len(args.gpu) > 0: + torch.set_default_device(args.gpu) return wan_model, offloadobj, pipe["transformer"] @@ -914,8 +933,10 @@ def generate_header(model_filename, compile, attention_mode): header += model_name header += " (attention mode: " + (attention_mode if attention_mode!="auto" else "auto/" + get_auto_attention() ) - if attention_mode not in attention_modes_supported: + if attention_mode not in attention_modes_installed: header += " -NOT INSTALLED-" + elif attention_mode not in attention_modes_supported: + header += " -NOT SUPPORTED-" if compile: header += ", pytorch compilation ON" @@ -979,11 +1000,7 @@ def apply_changes( state, if v != v_old: changes.append(k) - state["config_changes"] = changes - state["config_new"] = server_config - state["config_old"] = old_server_config - - global attention_mode, profile, compile, transformer_filename_t2v, transformer_filename_i2v, text_encoder_filename, vae_config, boost, lora_dir + global attention_mode, profile, compile, transformer_filename_t2v, transformer_filename_i2v, text_encoder_filename, vae_config, boost, lora_dir, reload_needed attention_mode = server_config["attention_mode"] profile = server_config["profile"] compile = server_config["compile"] @@ -995,7 +1012,7 @@ def apply_changes( state, if all(change in ["attention_mode", "vae_config", "default_ui", "boost", "save_path", "metadata_choice"] for change in changes ): pass else: - state["_reload_needed"] = True + reload_needed = True yield "
The new configuration has been succesfully applied
" @@ -1013,7 +1030,7 @@ def save_video(final_frames, output_path, fps=24): def build_callback(taskid, state, pipe, num_inference_steps, repeats): start_time = time.time() - def update_progress(step_idx, latents, read_state = False): + def update_progress(step_idx, _): with tracker_lock: step_idx += 1 if state.get("abort", False): @@ -1094,8 +1111,7 @@ def generate_video( progress=gr.Progress() #track_tqdm= True ): - global wan_model, offloadobj, last_model_type - reload_needed = state.get("_reload_needed", False) + global wan_model, offloadobj, reload_needed, last_model_type file_model_needed = model_needed(image2video) with lock: queue_not_empty = len(queue) > 0 @@ -1108,7 +1124,7 @@ def generate_video( print(f"Loading model {get_model_name(file_model_needed)}...") wan_model, offloadobj, trans = load_models(image2video) print(f"Model loaded") - state["_reload_needed"] = False + reload_needed= False from PIL import Image import numpy as np @@ -1121,11 +1137,12 @@ def generate_video( elif attention_mode in attention_modes_supported: attn = attention_mode else: - gr.Info(f"You have selected attention mode '{attention_mode}'. However it is not installed on your system. You should either install it or switch to the default 'sdpa' attention.") + 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 state.get("validate_success",0) != 1: # return + raw_resolution = resolution width, height = resolution.split("x") width, height = int(width), int(height) @@ -1289,7 +1306,7 @@ def generate_video( 'total_steps': num_inference_steps, 'start_time': time.time(), 'last_update': time.time(), - 'repeats': f"0/{repeat_generation}", + 'repeats': f"{video_no}/{repeat_generation}", 'status': "Encoding Prompt" } video_no += 1 @@ -1401,14 +1418,8 @@ def generate_video( normalize=True, value_range=(-1, 1)) - configs = { - 'prompt': prompt, - 'negative_prompt': negative_prompt, - 'resolution': resolution, - 'video_length': video_length, - 'seed': seed, - 'num_inference_steps': num_inference_steps, - } + configs = get_settings_dict(state, use_image2video, prompt, 0 if image_to_end == None else 1 , video_length, raw_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, RIFLEx_setting, slg_switch, slg_layers, slg_start, slg_end) metadata_choice = server_config.get("metadata_choice","metadata") if metadata_choice == "json": @@ -1715,19 +1726,14 @@ def switch_advanced(state, new_advanced, lset_name): else: 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 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 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, loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc): - if state.get("validate_success",0) != 1: - return - - loras_choices loras = state["loras"] activated_loras = [Path( loras[int(no)]).parts[-1] for no in loras_choices ] - ui_defaults = { + ui_settings = { "prompts": prompt, - "image_prompt_type": image_prompt_type, "resolution": resolution, "video_length": video_length, "num_inference_steps": num_inference_steps, @@ -1747,10 +1753,25 @@ def save_settings(state, prompt, image_prompt_type, video_length, resolution, nu "slg_start_perc": slg_start_perc, "slg_end_perc": slg_end_perc } + + if i2v: + ui_settings["type"] = "Wan2.1GP by DeepBeepMeep - image2video" + ui_settings["image_prompt_type"] = image_prompt_type + else: + 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, + loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc): + if state.get("validate_success",0) != 1: + return + ui_defaults = get_settings_dict(state, use_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, + loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc) + defaults_filename = get_settings_file_name(use_image2video) with open(defaults_filename, "w", encoding="utf-8") as f: - json.dump(ui_defaults, f, indent=4) + json.dump(ui_settings , f, indent=4) gr.Info("New Default Settings saved") @@ -1864,7 +1885,8 @@ def generate_video_tab(image2video=False): 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 = gr.Radio( [("Use only a Start Image", 0),("Use both a Start and an End Image", 1)], value =ui_defaults["image_prompt_type"], label="Location", show_label= False, scale= 3, visible=image2video) + 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( @@ -1876,9 +1898,9 @@ def generate_video_tab(image2video=False): 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=False) + 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= False) + image_to_end = gr.Image(label= "Last Image for a new video", type ="pil", visible=image_prompt_type==1) def switch_image_prompt_type_radio(image_prompt_type_radio): if args.multiple_images: @@ -1886,7 +1908,7 @@ def generate_video_tab(image2video=False): else: return gr.Image(visible = (image_prompt_type_radio == 1) ) - image_prompt_type.change(fn=switch_image_prompt_type_radio, inputs=[image_prompt_type], outputs=[image_to_end]) + image_prompt_type_radio.change(fn=switch_image_prompt_type_radio, inputs=[image_prompt_type_radio], outputs=[image_to_end]) advanced_prompt = advanced @@ -2080,7 +2102,7 @@ def generate_video_tab(image2video=False): outputs=[output] ) 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, 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, 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, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc ], 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]) @@ -2182,8 +2204,10 @@ def generate_configuration_tab(): value=server_config.get("save_path", save_path) ) def check(mode): - if not mode in attention_modes_supported: + if not mode in attention_modes_installed: return " (NOT INSTALLED)" + elif not mode in attention_modes_supported: + return " (NOT SUPPORTED)" else: return "" attention_choice = gr.Dropdown( @@ -2435,7 +2459,7 @@ def create_demo(): } """ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="sky", neutral_hue="slate", text_size="md")) as demo: - gr.Markdown("

Wan 2.1GP v3.0 by DeepBeepMeep (Updates)

") + gr.Markdown("

Wan 2.1GP v3.1 by DeepBeepMeep (Updates)

") gr.Markdown("Welcome to Wan 2.1GP a super fast and low VRAM AI Video Generator !") with gr.Accordion("Click here for some Info on how to use Wan2GP", open = False): diff --git a/requirements.txt b/requirements.txt index b97e545..bd928de 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,6 +16,6 @@ gradio>=5.0.0 numpy>=1.23.5,<2 einops moviepy==1.0.3 -mmgp==3.3.0 +mmgp==3.3.3 peft==0.14.0 mutagen \ No newline at end of file diff --git a/wan/image2video.py b/wan/image2video.py index 506fe90..a853665 100644 --- a/wan/image2video.py +++ b/wan/image2video.py @@ -177,7 +177,7 @@ class WanI2V: logging.info(f"Creating WanModel from {model_filename}") from mmgp import offload - self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel) + self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel, writable_tensors= False) self.model.eval().requires_grad_(False) if t5_fsdp or dit_fsdp or use_usp: diff --git a/wan/modules/attention.py b/wan/modules/attention.py index e40ac63..6861283 100644 --- a/wan/modules/attention.py +++ b/wan/modules/attention.py @@ -30,6 +30,7 @@ try: max_seqlen_kv, ): return sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv) + except ImportError: sageattn_varlen_wrapper = None @@ -38,11 +39,12 @@ import warnings try: from sageattention import sageattn - from .sage2_core import sageattn as alt_sageattn + from .sage2_core import sageattn as alt_sageattn, is_sage_supported + sage_supported = is_sage_supported() except ImportError: sageattn = None alt_sageattn = None - + sage_supported = False # @torch.compiler.disable() def sageattn_wrapper( qkv_list, @@ -129,6 +131,14 @@ def get_attention_modes(): return ret +def get_supported_attention_modes(): + ret = get_attention_modes() + if not sage_supported: + if "sage" in ret: + ret.remove("sage") + if "sage2" in ret: + ret.remove("sage2") + return ret __all__ = [ 'pay_attention', diff --git a/wan/modules/clip.py b/wan/modules/clip.py index 53b76c1..fc41d85 100644 --- a/wan/modules/clip.py +++ b/wan/modules/clip.py @@ -519,8 +519,11 @@ class CLIPModel: device=device) self.model = self.model.eval().requires_grad_(False) logging.info(f'loading {checkpoint_path}') - self.model.load_state_dict( - torch.load(checkpoint_path, map_location='cpu'), assign= True) + from mmgp import offload + # self.model.load_state_dict( + # torch.load(checkpoint_path, map_location='cpu'), assign= True) + + offload.load_model_data(self.model, checkpoint_path.replace(".pth", "-bf16.safetensors"), writable_tensors= False) # init tokenizer self.tokenizer = HuggingfaceTokenizer( diff --git a/wan/modules/sage2_core.py b/wan/modules/sage2_core.py index d83d255..de94a60 100644 --- a/wan/modules/sage2_core.py +++ b/wan/modules/sage2_core.py @@ -51,6 +51,15 @@ from sageattention.quant import per_channel_fp8 from typing import Any, List, Literal, Optional, Tuple, Union import warnings +import os + +def is_sage_supported(): + device_count = torch.cuda.device_count() + for i in range(device_count): + major, minor = torch.cuda.get_device_capability(i) + if major < 8: + return False + return True def get_cuda_arch_versions(): cuda_archs = [] diff --git a/wan/modules/t5.py b/wan/modules/t5.py index 5cccbbf..110e358 100644 --- a/wan/modules/t5.py +++ b/wan/modules/t5.py @@ -496,7 +496,7 @@ class T5EncoderModel: device=device).eval().requires_grad_(False) logging.info(f'loading {checkpoint_path}') from mmgp import offload - offload.load_model_data(model,checkpoint_path ) + offload.load_model_data(model,checkpoint_path, writable_tensors= False ) self.model = model if shard_fn is not None: diff --git a/wan/modules/vae.py b/wan/modules/vae.py index e47e74c..67dcd9a 100644 --- a/wan/modules/vae.py +++ b/wan/modules/vae.py @@ -744,11 +744,12 @@ def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs): with torch.device('meta'): model = WanVAE_(**cfg) + from mmgp import offload # load checkpoint logging.info(f'loading {pretrained_path}') - model.load_state_dict( - torch.load(pretrained_path, map_location=device), assign=True) - + # model.load_state_dict( + # torch.load(pretrained_path, map_location=device), assign=True) + offload.load_model_data(model, pretrained_path.replace(".pth", "_bf16.safetensors"), writable_tensors= False) return model @@ -778,7 +779,7 @@ class WanVAE: self.model = _video_vae( pretrained_path=vae_pth, z_dim=z_dim, - ).eval().requires_grad_(False).to(device) + ).eval() #.requires_grad_(False).to(device) def encode(self, videos, tile_size = 256, any_end_frame = False): """ diff --git a/wan/text2video.py b/wan/text2video.py index 4682a4a..88046db 100644 --- a/wan/text2video.py +++ b/wan/text2video.py @@ -89,7 +89,7 @@ class WanT2V: from mmgp import offload - self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel) + self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel, writable_tensors= False) From 9fe4b30e85e970a409ec0806235b38a05f6811f5 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 25 Mar 2025 19:52:02 +1100 Subject: [PATCH 12/69] remove duplicate gpu arg left over from merge --- gradio_server.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 6d7371a..424899e 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -380,12 +380,6 @@ def _parse_args(): default="", help="Server name" ) - parser.add_argument( - "--gpu", - type=str, - default="", - help="Default GPU Device" - ) parser.add_argument( "--gpu", From 5e29595e6d8d3afc0e2f9b6985b8907fe2e6f044 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 25 Mar 2025 19:59:49 +1100 Subject: [PATCH 13/69] fix save settings change missed from merge --- gradio_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gradio_server.py b/gradio_server.py index 424899e..c22345d 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1772,7 +1772,7 @@ def save_settings(state, prompt, image_prompt_type, video_length, resolution, nu defaults_filename = get_settings_file_name(use_image2video) with open(defaults_filename, "w", encoding="utf-8") as f: - json.dump(ui_settings , f, indent=4) + json.dump(ui_defaults , f, indent=4) gr.Info("New Default Settings saved") From 4f0dd2a998ee513931e5d02da3f55e88aff08182 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 25 Mar 2025 21:09:45 +1100 Subject: [PATCH 14/69] improve styling 3 --- gradio_server.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/gradio_server.py b/gradio_server.py index c22345d..088cf12 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -2426,6 +2426,19 @@ def create_demo(): overflow: hidden; text-overflow: ellipsis; } + #queue_df td:nth-child(-n+6) { + cursor: default !important; + pointer-events: none; + } + #queue_df th { + pointer-events: none; + } + #queue_df .tabulator-col { + pointer-events: none; + } + #queue_df .tabulator-col .tabulator-arrow { + display: none; + } #queue_df table { overflow: hidden !important; } @@ -2433,7 +2446,6 @@ def create_demo(): #queue_df::-webkit-scrollbar { display: none !important; } - #queue_df { scrollbar-width: none !important; -ms-overflow-style: none !important; From ffe656d89771de15ecd588d1d622188f18f0bade Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Wed, 26 Mar 2025 17:45:12 +1100 Subject: [PATCH 15/69] fix missing comma --- gradio_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gradio_server.py b/gradio_server.py index b429867..f38169c 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -980,7 +980,7 @@ def apply_changes( state, "metadata_choice": metadata_choice, "default_ui" : default_ui_choice, "boost" : boost_choice, - "clear_file_list" : clear_file_list + "clear_file_list" : clear_file_list, "reload_model" : reload_choice, } From 85d732083749e1f99df3a9ff82c4d2827afa4e81 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Wed, 26 Mar 2025 18:05:07 +1100 Subject: [PATCH 16/69] re-removed yield statements that prevented it running --- gradio_server.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index f38169c..b9a23c4 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1431,7 +1431,7 @@ def generate_video( end_time = time.time() abort = True state["prompt"] = "" - yield f"Video generation was aborted. Total Generation Time: {end_time-start_time:.1f}s" + print(f"Video generation was aborted. Total Generation Time: {end_time-start_time:.1f}s") else: sample = samples.cpu() # video = rearrange(sample.cpu().numpy(), "c t h w -> t h w c") @@ -1465,12 +1465,6 @@ def generate_video( print(f"New video saved to Path: "+video_path) file_list.append(video_path) - if video_no < total_video: - yield status - else: - end_time = time.time() - state["prompt"] = "" - yield f"Total Generation Time: {end_time-start_time:.1f}s" seed += 1 repeat_no += 1 From 55a9c698c009b204a06d6a2884dcdba7d3e91194 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Wed, 26 Mar 2025 18:31:11 +1100 Subject: [PATCH 17/69] removed unused variable --- gradio_server.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gradio_server.py b/gradio_server.py index b9a23c4..31c0ca2 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1466,7 +1466,6 @@ def generate_video( print(f"New video saved to Path: "+video_path) file_list.append(video_path) seed += 1 - repeat_no += 1 last_model_type = image2video From 4fac11dc68ecdc4bbca7ab9c8f504e670d75f836 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Thu, 27 Mar 2025 00:24:09 +1100 Subject: [PATCH 18/69] fixed output gallery updating every second, now updates only when generation finished --- gradio_server.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 31c0ca2..da17384 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1057,8 +1057,7 @@ def build_callback(taskid, state, pipe, num_inference_steps, repeats): return update_progress def refresh_gallery(state): - file_list = state.get("file_list", None) - return file_list + return state def finalize_gallery(state): choice = 0 @@ -1465,6 +1464,7 @@ def generate_video( print(f"New video saved to Path: "+video_path) file_list.append(video_path) + state['update_gallery'] = True seed += 1 last_model_type = image2video @@ -1899,6 +1899,7 @@ def generate_video_tab(image2video=False): download_status = gr.Markdown() with gr.Row(): with gr.Column(): + gallery_update_trigger = gr.Textbox(value="0", visible=False, label="_gallery_trigger") with gr.Row(visible= len(loras)>0) as presets_column: lset_choices = [ (preset, preset) for preset in loras_presets ] + [(get_new_preset_msg(advanced), "")] with gr.Column(scale=6): @@ -2146,15 +2147,24 @@ def generate_video_tab(image2video=False): elif cell_value == "✖": return remove_task([selected_index]) return queue_df + def refresh_gallery_on_trigger(state): + if(state.get("update_gallery", False)): + state['update_gallery'] = False + return gr.update(value=state.get("file_list", [])) selected_indices = gr.State([]) queue_df.select( fn=handle_selection, outputs=selected_indices ) + gallery_update_trigger.change( + fn=refresh_gallery_on_trigger, + inputs=[state], + outputs=[output] + ) queue_df.change( fn=refresh_gallery, inputs=[state], - outputs=[output] + outputs=[gallery_update_trigger] ) 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, From 4ef20d2299cb95801ac869002d3b3a9afb400abc Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Fri, 28 Mar 2025 02:08:38 +1100 Subject: [PATCH 19/69] add image thumbnails and previews to queue items --- gradio_server.py | 169 +++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 141 insertions(+), 28 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 0ebb534..237b9dd 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -22,6 +22,9 @@ import traceback import math import asyncio from wan.utils import prompt_parser +import base64 +import io +from PIL import Image PROMPT_VARS_MAX = 10 target_mmgp_version = "3.3.4" @@ -50,6 +53,29 @@ def format_time(seconds): minutes = int((seconds % 3600) // 60) return f"{hours}h {minutes}m" +def pil_to_base64_uri(pil_image, format="png", quality=75): + if pil_image is None: + return None + buffer = io.BytesIO() + try: + img_to_save = pil_image + if format.lower() == 'jpeg' and pil_image.mode == 'RGBA': + img_to_save = pil_image.convert('RGB') + elif format.lower() == 'png' and pil_image.mode not in ['RGB', 'RGBA', 'L', 'P']: + img_to_save = pil_image.convert('RGBA') + elif pil_image.mode == 'P': + img_to_save = pil_image.convert('RGBA' if 'transparency' in pil_image.info else 'RGB') + if format.lower() == 'jpeg': + img_to_save.save(buffer, format=format, quality=quality) + else: + img_to_save.save(buffer, format=format) + img_bytes = buffer.getvalue() + encoded_string = base64.b64encode(img_bytes).decode("utf-8") + return f"data:image/{format.lower()};base64,{encoded_string}" + except Exception as e: + print(f"Error converting PIL to base64: {e}") + return None + def runner(): global current_task_id while True: @@ -175,6 +201,9 @@ def add_video_task(*params): with lock: task_id += 1 current_task_id = task_id + start_image_data = params[16] if len(params) > 16 else None + end_image_data = params[17] if len(params) > 17 else None + queue.append({ "id": current_task_id, "params": (current_task_id,) + params, @@ -184,9 +213,11 @@ def add_video_task(*params): "progress": "0.0%", "steps": f"0/{params[5]}", "time": "--", - "prompt": params[0] + "prompt": params[0], + "start_image_data": start_image_data, + "end_image_data": end_image_data }) - return + return update_queue_data() def move_up(selected_indices): if not selected_indices or len(selected_indices) == 0: @@ -233,6 +264,15 @@ def update_queue_data(): truncated_prompt = (item['prompt'][:97] + '...') if len(item['prompt']) > 100 else item['prompt'] full_prompt = item['prompt'].replace('"', '"') prompt_cell = f'{truncated_prompt}' + start_img_uri = pil_to_base64_uri(item.get('start_image_data'), format="jpeg", quality=70) + end_img_uri = pil_to_base64_uri(item.get('end_image_data'), format="jpeg", quality=70) + thumbnail_size = "50px" + start_img_md = "" + end_img_md = "" + if start_img_uri: + start_img_md = f'Start' + if end_img_uri: + end_img_md = f'End' data.append([ item.get('status', "Starting"), item.get('repeats', "0/0"), @@ -240,6 +280,8 @@ def update_queue_data(): item.get('steps', ''), item.get('time', '--'), prompt_cell, + start_img_md, + end_img_md, "↑", "↓", "✖" @@ -1143,7 +1185,6 @@ def generate_video( print(f"Model loaded") reload_needed= False - from PIL import Image import numpy as np import tempfile @@ -1905,6 +1946,10 @@ def generate_video_tab(image2video=False): download_status = gr.Markdown() with gr.Row(): with gr.Column(): + with gr.Column(visible=False, elem_id="image-modal-container") as modal_container: + with gr.Row(elem_id="image-modal-close-button-row"): + close_modal_button = gr.Button("❌", size="sm") + modal_image_display = gr.Image(label="Full Resolution Image", interactive=False, show_label=False) gallery_update_trigger = gr.Textbox(value="0", visible=False, label="_gallery_trigger") with gr.Row(visible= len(loras)>0) as presets_column: lset_choices = [ (preset, preset) for preset in loras_presets ] + [(get_new_preset_msg(advanced), "")] @@ -2134,25 +2179,49 @@ def generate_video_tab(image2video=False): , columns=[3], rows=[1], object_fit="contain", height=450, selected_index=0, interactive= False) generate_btn = gr.Button("Generate") queue_df = gr.DataFrame( - headers=["Status", "Completed", "Progress", "Steps", "Time", "Prompt", "", "", ""], - datatype=["str", "str", "str", "str", "str", "markdown", "str", "str", "str"], + headers=["Status", "Completed", "Progress", "Steps", "Time", "Prompt", "Start", "End", "", "", ""], + datatype=["str", "str", "str", "str", "str", "markdown", "markdown", "markdown", "str", "str", "str"], interactive=False, - col_count=(9, "fixed"), + col_count=(11, "fixed"), wrap=True, value=update_queue_data, every=1, elem_id="queue_df" ) def handle_selection(evt: gr.SelectData): - cell_value = evt.value - selected_index = evt.index - if cell_value == "↑": - return move_up([selected_index]) - elif cell_value == "↓": - return move_down([selected_index]) - elif cell_value == "✖": - return remove_task([selected_index]) - return queue_df + if evt.index is None: + return gr.update(), gr.update(), gr.update(visible=False) + row_index, col_index = evt.index + cell_value = None + if col_index in [8, 9, 10]: + if col_index == 8: cell_value = "↑" + elif col_index == 9: cell_value = "↓" + elif col_index == 10: cell_value = "✖" + if col_index == 8: + new_df_data = move_up([row_index]) + return new_df_data, gr.update(), gr.update(visible=False) + elif col_index == 9: + new_df_data = move_down([row_index]) + return new_df_data, gr.update(), gr.update(visible=False) + elif col_index == 10: + new_df_data = remove_task([row_index]) + return new_df_data, gr.update(), gr.update(visible=False) + start_img_col_idx = 6 + end_img_col_idx = 7 + image_data_to_show = None + if col_index == start_img_col_idx: + with lock: + if row_index < len(queue): + image_data_to_show = queue[row_index].get('start_image_data') + elif col_index == end_img_col_idx: + with lock: + if row_index < len(queue): + image_data_to_show = queue[row_index].get('end_image_data') + + if image_data_to_show: + return gr.update(), gr.update(value=image_data_to_show), gr.update(visible=True) + else: + return gr.update(), gr.update(), gr.update(visible=False) def refresh_gallery_on_trigger(state): if(state.get("update_gallery", False)): state['update_gallery'] = False @@ -2160,7 +2229,8 @@ def generate_video_tab(image2video=False): selected_indices = gr.State([]) queue_df.select( fn=handle_selection, - outputs=selected_indices + inputs=None, + outputs=[queue_df, modal_image_display, modal_container], ) gallery_update_trigger.change( fn=refresh_gallery_on_trigger, @@ -2229,6 +2299,11 @@ def generate_video_tab(image2video=False): inputs=original_inputs, outputs=queue_df ) + close_modal_button.click( + lambda: gr.update(visible=False), + inputs=[], + outputs=[modal_container] + ) return loras_choices, lset_name, header, state def generate_configuration_tab(): @@ -2524,16 +2599,9 @@ def create_demo(): #queue_df th { pointer-events: none; } - #queue_df .tabulator-col { - pointer-events: none; - } - #queue_df .tabulator-col .tabulator-arrow { - display: none; - } #queue_df table { overflow: hidden !important; } - #queue_df::-webkit-scrollbar { display: none !important; } @@ -2545,7 +2613,8 @@ def create_demo(): width: 100px; } #queue_df td:nth-child(6) { - width: 300px; + width: auto; + min-width: 200px; } #queue_df td:nth-child(7), #queue_df td:nth-child(8), @@ -2553,12 +2622,56 @@ def create_demo(): cursor: pointer; text-align: center; font-weight: bold; + width: 60px; + text-align: center; + padding: 2px !important; + cursor: pointer; } - #queue_df td:nth-child(7):hover, - #queue_df td:nth-child(8):hover, - #queue_df td:nth-child(9):hover { - background-color: #e0e0e0; + #queue_df td:nth-child(10) img, + #queue_df td:nth-child(11) img { + max-width: 50px; + max-height: 50px; + object-fit: contain; + display: block; + margin: auto; } + #image-modal-container { + position: fixed; + top: 0; + left: 0; + width: 100%; + height: 100%; + background-color: rgba(0, 0, 0, 0.7); + justify-content: center; + align-items: center; + z-index: 1000; + padding: 20px; + box-sizing: border-box; + } + #image-modal-container > div { + background-color: white; + padding: 15px; + border-radius: 8px; + max-width: 90%; + max-height: 90%; + overflow: auto; + position: relative; + display: flex; + flex-direction: column; + } + #image-modal-container img { + max-width: 100%; + max-height: 80vh; + object-fit: contain; + margin-top: 10px; + } + #image-modal-close-button-row { + display: flex; + justify-content: flex-end; + } + #image-modal-close-button-row button { + cursor: pointer; + } """ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="sky", neutral_hue="slate", text_size="md")) as demo: gr.Markdown("

Wan 2.1GP v3.2 by DeepBeepMeep (Updates)

") From 7dbc2735d8094b2622c1706b1f04853a33e41414 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Fri, 28 Mar 2025 13:57:41 +1100 Subject: [PATCH 20/69] fix styling corruption from image thumbnails --- gradio_server.py | 39 +++++++++++++++++++++------------------ 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 4e547bb..de110cb 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -2662,28 +2662,31 @@ def create_demo(): #queue_df td:nth-child(1) { width: 100px; } - #queue_df td:nth-child(6) { - width: auto; - min-width: 200px; - } - #queue_df td:nth-child(7), - #queue_df td:nth-child(8), - #queue_df td:nth-child(9) { - cursor: pointer; - text-align: center; - font-weight: bold; - width: 60px; - text-align: center; - padding: 2px !important; - cursor: pointer; - } - #queue_df td:nth-child(10) img, - #queue_df td:nth-child(11) img { + #queue_df td:nth-child(7) img, + #queue_df td:nth-child(8) img, max-width: 50px; max-height: 50px; object-fit: contain; display: block; - margin: auto; + margin: auto; + cursor: pointer; + text-align: center; + } + #queue_df td:nth-child(9), + #queue_df td:nth-child(10), + #queue_df td:nth-child(11) { + width: 60px; + padding: 2px !important; + cursor: pointer; + text-align: center; + font-weight: bold; + } + #queue_df td:nth-child(7):hover, + #queue_df td:nth-child(8):hover, + #queue_df td:nth-child(9):hover, + #queue_df td:nth-child(10):hover, + #queue_df td:nth-child(11):hover { + background-color: #e0e0e0; } #image-modal-container { position: fixed; From 74d3f42c18f7366d190ba1efbd23e08679f9e0cc Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Fri, 28 Mar 2025 14:12:07 +1100 Subject: [PATCH 21/69] resimplify data structure for generate_video --- gradio_server.py | 63 ++++++++++++++++++++++++------------------------ 1 file changed, 31 insertions(+), 32 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index de110cb..7bbec43 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -2284,44 +2284,43 @@ def generate_video_tab(image2video=False): refresh_lora_btn2.click(refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) download_loras_btn.click(fn=download_loras, inputs=[], outputs=[download_status_row, download_status, presets_column, loras_column]).then(fn=refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) output.select(select_video, state, None ) - original_inputs = [ - prompt, - negative_prompt, - resolution, - video_length, - seed, - num_inference_steps, - guidance_scale, - flow_shift, - embedded_guidance_scale, - repeat_generation, - multi_images_gen_type, - tea_cache_setting, - tea_cache_start_step_perc, - loras_choices, - loras_mult_choices, - image_prompt_type_radio, - image_to_continue, - image_to_end, - video_to_continue, - max_frames, - RIFLEx_setting, - slg_switch, - slg_layers, - slg_start_perc, - slg_end_perc, - cfg_star_switch, - cfg_zero_step, - state, - gr.State(image2video) - ] #generate_btn.click( # fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , outputs= [prompt] #).then( generate_btn.click( fn=process_prompt_and_add_tasks, - inputs=original_inputs, + inputs=[ + prompt, + negative_prompt, + resolution, + video_length, + seed, + num_inference_steps, + guidance_scale, + flow_shift, + embedded_guidance_scale, + repeat_generation, + multi_images_gen_type, + tea_cache_setting, + tea_cache_start_step_perc, + loras_choices, + loras_mult_choices, + image_prompt_type_radio, + image_to_continue, + image_to_end, + video_to_continue, + max_frames, + RIFLEx_setting, + slg_switch, + slg_layers, + slg_start_perc, + slg_end_perc, + cfg_star_switch, + cfg_zero_step, + state, + gr.State(image2video) + ], outputs=queue_df ) close_modal_button.click( From b44229a529ac87137203148c10169fd364230bbb Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Sat, 29 Mar 2025 00:42:30 +1100 Subject: [PATCH 22/69] restored validation wizard --- gradio_server.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 7bbec43..c310e1a 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -139,11 +139,15 @@ def process_prompt_and_add_tasks( state_arg, image2video ): + + if state_arg.get("validate_success",0) != 1: + print("Validation failed, not adding tasks.") + return if len(prompt) ==0: return prompt, errors = prompt_parser.process_template(prompt) if len(errors) > 0: - gr.Info("Error processing prompt template: " + errors) + print("Error processing prompt template: " + errors) return prompts = prompt.replace("\r", "").split("\n") prompts = [prompt.strip() for prompt in prompts if len(prompt.strip())>0 and not prompt.startswith("#")] @@ -1223,8 +1227,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 state.get("validate_success",0) != 1: - # return raw_resolution = resolution width, height = resolution.split("x") width, height = int(width), int(height) @@ -1546,6 +1548,7 @@ def generate_video( gen_in_progress = False offload.unload_loras_from_model(trans) + def get_new_preset_msg(advanced = True): if advanced: return "Enter here a Name for a Lora Preset or Choose one in the List" @@ -2285,10 +2288,9 @@ def generate_video_tab(image2video=False): download_loras_btn.click(fn=download_loras, inputs=[], outputs=[download_status_row, download_status, presets_column, loras_column]).then(fn=refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) output.select(select_video, state, None ) - #generate_btn.click( - # fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , outputs= [prompt] - #).then( generate_btn.click( + fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , outputs= [prompt] + ).then( fn=process_prompt_and_add_tasks, inputs=[ prompt, From ec981450df187642fdbdcef987d902b8b1185c93 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Sat, 29 Mar 2025 01:03:27 +1100 Subject: [PATCH 23/69] pre-merge for exif support --- gradio_server.py | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index c310e1a..d0fe461 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1164,6 +1164,24 @@ def expand_slist(slist, num_inference_steps ): 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 + def generate_video( task_id, prompt, @@ -1214,9 +1232,6 @@ def generate_video( print(f"Model loaded") reload_needed= False - import numpy as np - import tempfile - if wan_model == None: raise gr.Error("Unable to generate a Video while a new configuration is being applied.") if attention_mode == "auto": @@ -1413,8 +1428,8 @@ def generate_video( if image2video: samples = wan_model.generate( prompt, - image_to_continue.convert('RGB'), - image_to_end.convert('RGB') if image_to_end != None else None, + convert_image(image_to_continue), + convert_image(image_to_end) if image_to_end != None else None, frame_num=(video_length // 4)* 4 + 1, max_area=MAX_AREA_CONFIGS[resolution], shift=flow_shift, From 519b38493d3f402a6583c6331861a737a4fef24b Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Sun, 30 Mar 2025 05:33:41 +1100 Subject: [PATCH 24/69] added status and progress bars back --- gradio_server.py | 110 ++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 108 insertions(+), 2 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index d0fe461..f712ef1 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -41,6 +41,7 @@ progress_tracker = {} tracker_lock = threading.Lock() file_list = [] last_model_type = None +last_status_string = "" def format_time(seconds): if seconds < 60: @@ -292,6 +293,53 @@ def update_queue_data(): ]) return data +def create_html_progress_bar(percentage=0.0, text="Idle", is_idle=True): + bar_class = "progress-bar-custom idle" if is_idle else "progress-bar-custom" + bar_text_html = f'
{text}
' + + html = f""" +
+
+ {bar_text_html} +
+
+ """ + return html + +def refresh_progress(): + global current_task_id, progress_tracker, last_status_string + task_id_to_check = current_task_id + is_idle = True + status_string = "Starting..." + progress_percent = 0.0 + html_content = "" + + with tracker_lock: + with lock: + processing_or_queued = any(item['state'] in ["Processing", "Queued"] for item in queue) + if task_id_to_check is not None: + progress_data = progress_tracker.get(task_id_to_check) + if progress_data: + is_idle = False + current_step = progress_data.get('current_step', 0) + total_steps = progress_data.get('total_steps', 0) + status = progress_data.get('status', "Starting...") + repeats = progress_data.get("repeats", "0/0") + + if total_steps > 0: + progress_float = min(1.0, max(0.0, float(current_step) / float(total_steps))) + progress_percent = progress_float * 100 + status_string = f"{status} [{repeats}] - {progress_percent:.1f}% complete ({current_step}/{total_steps} steps)" + else: + progress_percent = 0.0 + status_string = f"{status} [{repeats}] - Initializing..." + html_content = create_html_progress_bar(progress_percent, status_string, is_idle) + return gr.update(value=html_content) + +def update_generation_status(html_content): + if(html_content): + return gr.update(value=html_content) + def _parse_args(): parser = argparse.ArgumentParser( description="Generate a video from a text prompt or image using Gradio") @@ -1212,8 +1260,7 @@ def generate_video( cfg_star_switch, cfg_zero_step, state, - image2video, - progress=gr.Progress() #track_tqdm= True + image2video ): @@ -1993,6 +2040,7 @@ def generate_video_tab(image2video=False): with gr.Row(elem_id="image-modal-close-button-row"): close_modal_button = gr.Button("❌", size="sm") modal_image_display = gr.Image(label="Full Resolution Image", interactive=False, show_label=False) + progress_update_trigger = gr.Textbox(value="0", visible=False, label="_progress_trigger") gallery_update_trigger = gr.Textbox(value="0", visible=False, label="_gallery_trigger") with gr.Row(visible= len(loras)>0) as presets_column: lset_choices = [ (preset, preset) for preset in loras_presets ] + [(get_new_preset_msg(advanced), "")] @@ -2217,6 +2265,11 @@ 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_progress_html = gr.HTML( + label="Status", + value="Idle", + elem_id="generation_progress_bar_container" + ) output = gr.Gallery( label="Generated videos", show_label=False, elem_id="gallery" , columns=[3], rows=[1], object_fit="contain", height=450, selected_index=0, interactive= False) @@ -2284,6 +2337,16 @@ def generate_video_tab(image2video=False): fn=refresh_gallery, inputs=[state], outputs=[gallery_update_trigger] + ).then( + fn=refresh_progress, + inputs=None, + outputs=[progress_update_trigger] + ) + progress_update_trigger.change( + fn=update_generation_status, + inputs=[progress_update_trigger], + outputs=[gen_progress_html], + 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, @@ -2741,6 +2804,49 @@ def create_demo(): #image-modal-close-button-row button { cursor: pointer; } + .progress-container-custom { + width: 100%; + background-color: #e9ecef; + border-radius: 0.375rem; + overflow: hidden; + height: 25px; + position: relative; + margin-top: 5px; + margin-bottom: 5px; + } + .progress-bar-custom { + height: 100%; + background-color: #0d6efd; + transition: width 0.3s ease-in-out; + display: flex; + align-items: center; + justify-content: center; + color: white; + font-size: 0.9em; + font-weight: bold; + white-space: nowrap; + overflow: hidden; + } + .progress-bar-custom.idle { + background-color: #6c757d; + } + .progress-bar-text { + position: absolute; + top: 0; + left: 0; + width: 100%; + height: 100%; + display: flex; + align-items: center; + justify-content: center; + color: white; + mix-blend-mode: difference; + font-size: 0.9em; + font-weight: bold; + white-space: nowrap; + z-index: 2; + pointer-events: none; + } """ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="sky", neutral_hue="slate", text_size="md")) as demo: gr.Markdown("

Wan 2.1GP v3.3 by DeepBeepMeep (Updates)

") From 6fc7f762a19cd51ce1fca1303dd800041da7ffa2 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Sun, 30 Mar 2025 06:03:45 +1100 Subject: [PATCH 25/69] fix gallery disappearing while next queued item generating --- gradio_server.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index f712ef1..caee14b 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1176,7 +1176,12 @@ def build_callback(taskid, state, pipe, num_inference_steps, repeats): return update_progress def refresh_gallery(state): - return state + return gr.update(value=state.get("file_list", [])) + +def refresh_gallery_on_trigger(state): + if(state.get("update_gallery", False)): + state['update_gallery'] = False + return gr.update(value=state.get("file_list", [])) def finalize_gallery(state): choice = 0 @@ -2318,10 +2323,6 @@ def generate_video_tab(image2video=False): return gr.update(), gr.update(value=image_data_to_show), gr.update(visible=True) else: return gr.update(), gr.update(), gr.update(visible=False) - def refresh_gallery_on_trigger(state): - if(state.get("update_gallery", False)): - state['update_gallery'] = False - return gr.update(value=state.get("file_list", [])) selected_indices = gr.State([]) queue_df.select( fn=handle_selection, From e16cd6204327db3b83f1899de19c9eec26effd75 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Sun, 30 Mar 2025 06:05:25 +1100 Subject: [PATCH 26/69] remove redundant gallery function --- gradio_server.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index caee14b..b509505 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1183,22 +1183,6 @@ def refresh_gallery_on_trigger(state): state['update_gallery'] = False return gr.update(value=state.get("file_list", [])) -def finalize_gallery(state): - choice = 0 - if "in_progress" in state: - del state["in_progress"] - choice = state.get("selected",0) - if state.get("last_selected", True): - file_list = state.get("file_list", []) - choice = len(file_list) - 1 - - - state["extra_orders"] = 0 - time.sleep(0.2) - global gen_in_progress - gen_in_progress = False - return gr.Gallery(selected_index=choice), gr.Button(interactive=True), gr.Button(visible=False), gr.Checkbox(visible=False), gr.Text(visible=False, value="") - def select_video(state , event_data: gr.EventData): data= event_data._data if data!=None: From e5da4fdab4cbb673251cd109662a010b4e9dc9d5 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Sun, 30 Mar 2025 07:12:00 +1100 Subject: [PATCH 27/69] fixed styling on prompt tooltips --- gradio_server.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/gradio_server.py b/gradio_server.py index b509505..05558fa 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -2706,10 +2706,13 @@ def create_demo(): overflow: hidden; text-overflow: ellipsis; } - #queue_df td:nth-child(-n+6) { + #queue_df td:nth-child(-n+5) { cursor: default !important; pointer-events: none; } + #queue_df td:nth-child(6) { + cursor: default !important; + } #queue_df th { pointer-events: none; } From f1b1836a7a28b846eee9fd188d19c978ec005b30 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Sun, 30 Mar 2025 09:33:53 +1100 Subject: [PATCH 28/69] improve styling (4) --- gradio_server.py | 73 ++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 65 insertions(+), 8 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 05558fa..6f7f691 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -2715,8 +2715,11 @@ def create_demo(): } #queue_df th { pointer-events: none; + text-align: center; + vertical-align: middle; } #queue_df table { + width: 100%; overflow: hidden !important; } #queue_df::-webkit-scrollbar { @@ -2726,27 +2729,81 @@ def create_demo(): scrollbar-width: none !important; -ms-overflow-style: none !important; } + #queue_df th:nth-child(1), #queue_df td:nth-child(1) { - width: 100px; + width: 90px; + text-align: center; + vertical-align: middle; + } + #queue_df th:nth-child(1) { + font-size: 0.8em; + } + #queue_df th:nth-child(2), + #queue_df td:nth-child(2) { + width: 85px; + text-align: center; + vertical-align: middle; + } + #queue_df th:nth-child(2) { + font-size: 0.5em; + } + #queue_df th:nth-child(3), + #queue_df td:nth-child(3) { + width: 75px; + text-align: center; + vertical-align: middle; + } + #queue_df th:nth-child(3) { + font-size: 0.6em; + } + #queue_df th:nth-child(4), + #queue_df td:nth-child(4) { + width: 65px; + text-align: center; + white-space: nowrap; + } + #queue_df th:nth-child(4) { + font-size: 0.9em; + } + #queue_df th:nth-child(5), + #queue_df td:nth-child(5) { + width: 60px; + text-align: center; + white-space: nowrap; + } + #queue_df th:nth-child(6), + #queue_df td:nth-child(6) { + width: auto; + text-align: center; + white-space: normal; + } + #queue_df th:nth-child(6) { + font-size: 0.8em; + } + #queue_df th:nth-child(7), #queue_df td:nth-child(7), + #queue_df th:nth-child(8), #queue_df td:nth-child(8) { + width: 60px; + text-align: center; + vertical-align: middle; } #queue_df td:nth-child(7) img, - #queue_df td:nth-child(8) img, + #queue_df td:nth-child(8) img { max-width: 50px; max-height: 50px; object-fit: contain; display: block; - margin: auto; + margin: auto; cursor: pointer; - text-align: center; } - #queue_df td:nth-child(9), - #queue_df td:nth-child(10), - #queue_df td:nth-child(11) { - width: 60px; + #queue_df th:nth-child(9), #queue_df td:nth-child(9), + #queue_df th:nth-child(10), #queue_df td:nth-child(10), + #queue_df th:nth-child(11), #queue_df td:nth-child(11) { + width: 20px; padding: 2px !important; cursor: pointer; text-align: center; font-weight: bold; + vertical-align: middle; } #queue_df td:nth-child(7):hover, #queue_df td:nth-child(8):hover, From 3246d8cf4d06d5d037fcd75cdf7cf985bad37134 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Tue, 1 Apr 2025 20:32:30 +0200 Subject: [PATCH 29/69] Queue adaptations --- gradio_server.py | 1557 +++++++++++++++++++++++++++--------------- wan/image2video.py | 63 +- wan/modules/model.py | 7 +- wan/text2video.py | 4 +- 4 files changed, 1061 insertions(+), 570 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index 6f7f691..e21c5cb 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1,5 +1,6 @@ import os import time +import sys import threading import argparse from mmgp import offload, safetensors2, profile_type @@ -33,15 +34,12 @@ mmgp_version = version("mmgp") if mmgp_version != target_mmgp_version: print(f"Incorrect version of mmgp ({mmgp_version}), version {target_mmgp_version} is needed. Please upgrade with the command 'pip install -r requirements.txt'") exit() -queue = [] lock = threading.Lock() current_task_id = None task_id = 0 -progress_tracker = {} -tracker_lock = threading.Lock() -file_list = [] +# progress_tracker = {} +# tracker_lock = threading.Lock() last_model_type = None -last_status_string = "" def format_time(seconds): if seconds < 60: @@ -77,37 +75,6 @@ def pil_to_base64_uri(pil_image, format="png", quality=75): print(f"Error converting PIL to base64: {e}") return None -def runner(): - global current_task_id - while True: - with lock: - for item in queue: - task_id_runner = item['id'] - with tracker_lock: - progress = progress_tracker.get(task_id_runner, {}) - - if item['state'] == "Processing": - current_step = progress.get('current_step', 0) - total_steps = progress.get('total_steps', 0) - elapsed = time.time() - progress.get('start_time', time.time()) - status = progress.get('status', "") - repeats = progress.get("repeats", "0/0") - item.update({ - 'progress': f"{((current_step/total_steps)*100 if total_steps > 0 else 0):.1f}%", - 'steps': f"{current_step}/{total_steps}", - 'time': format_time(elapsed), - 'repeats': f"{repeats}", - 'status': f"{status}" - }) - if not any(item['state'] == "Processing" for item in queue): - for item in queue: - if item['state'] == "Queued": - item['status'] = "Processing" - item['state'] = "Processing" - current_task_id = item['id'] - threading.Thread(target=process_task, args=(item,)).start() - break - time.sleep(1) def process_prompt_and_add_tasks( prompt, @@ -137,161 +104,290 @@ def process_prompt_and_add_tasks( slg_end, cfg_star_switch, cfg_zero_step, - state_arg, + state, image2video ): - - if state_arg.get("validate_success",0) != 1: - print("Validation failed, not adding tasks.") + + if state.get("validate_success",0) != 1: + gr.Info("Validation failed, not adding tasks.") return + + state["validate_success"] = 0 if len(prompt) ==0: return prompt, errors = prompt_parser.process_template(prompt) if len(errors) > 0: - print("Error processing prompt template: " + errors) + gr.Info("Error processing prompt template: " + errors) return prompts = prompt.replace("\r", "").split("\n") prompts = [prompt.strip() for prompt in prompts if len(prompt.strip())>0 and not prompt.startswith("#")] if len(prompts) ==0: return - for single_prompt in prompts: - task_params = ( - single_prompt, - negative_prompt, - resolution, - video_length, - seed, - num_inference_steps, - guidance_scale, - flow_shift, - embedded_guidance_scale, - repeat_generation, - multi_images_gen_type, - tea_cache, - tea_cache_start_step_perc, - loras_choices, - loras_mult_choices, - image_prompt_type, - image_to_continue, - image_to_end, - video_to_continue, - max_frames, - RIFLEx_setting, - slg_switch, - slg_layers, - slg_start, - slg_end, - cfg_star_switch, - cfg_zero_step, - state_arg, - image2video - ) - add_video_task(*task_params) - return update_queue_data() + file_model_needed = model_needed(image2video) + if image2video: + width, height = resolution.split("x") + width, height = int(width), int(height) -def process_task(task): - try: - task_id, *params = task['params'] - generate_video(task_id, *params) - finally: - with lock: - queue[:] = [item for item in queue if item['id'] != task['id']] - with tracker_lock: - if task['id'] in progress_tracker: - del progress_tracker[task['id']] + 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") + return + resolution = str(width) + "*" + str(height) + if resolution not in ['720*1280', '1280*720', '480*832', '832*480']: + gr.Info(f"Resolution {resolution} not supported by image 2 video") + return -def add_video_task(*params): + if "1.3B" in file_model_needed and width * height > 848*480: + 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: + 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 ] + 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 ] + else: + image_to_end = [convert_image(image_to_end) ] + if len(image_to_continue) != len(image_to_end): + 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_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)] ) + prompts = new_prompts + image_to_continue = new_image_to_continue + if image_to_end != None: + image_to_end = new_image_to_end + else: + if len(prompts) >= len(image_to_continue): + if len(prompts) % len(image_to_continue) !=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 = [] + 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 + else: + if len(image_to_continue) % 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) + new_prompts = [] + for i, _ in enumerate(image_to_continue): + 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) : + kwargs = { + "prompt" : single_prompt, + "negative_prompt" : negative_prompt, + "resolution" : resolution, + "video_length" : video_length, + "seed" : seed, + "num_inference_steps" : num_inference_steps, + "guidance_scale" : guidance_scale, + "flow_shift" : flow_shift, + "embedded_guidance_scale" : embedded_guidance_scale, + "repeat_generation" : repeat_generation, + "multi_images_gen_type" : multi_images_gen_type, + "tea_cache" : tea_cache, + "tea_cache_start_step_perc" : tea_cache_start_step_perc, + "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 , + "max_frames" : max_frames, + "RIFLEx_setting" : RIFLEx_setting, + "slg_switch" : slg_switch, + "slg_layers" : slg_layers, + "slg_start" : slg_start, + "slg_end" : slg_end, + "cfg_star_switch" : cfg_star_switch, + "cfg_zero_step" : cfg_zero_step, + "state" : state, + "image2video" : image2video + } + add_video_task(**kwargs) + + gen = get_gen_info(state) + gen["prompts_max"] = len(prompts) + gen.get("prompts_max",0) + state["validate_success"] = 1 + queue= gen.get("queue", []) + return update_queue_data(queue) + + + + +def add_video_task(**kwargs): global task_id - with lock: - task_id += 1 - current_task_id = task_id - start_image_data = params[16] if len(params) > 16 else None - end_image_data = params[17] if len(params) > 17 else None + state = kwargs["state"] + gen = get_gen_info(state) + 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"] - queue.append({ - "id": current_task_id, - "params": (current_task_id,) + params, - "state": "Queued", - "status": "Queued", - "repeats": "0/0", - "progress": "0.0%", - "steps": f"0/{params[5]}", - "time": "--", - "prompt": params[0], - "start_image_data": start_image_data, - "end_image_data": end_image_data - }) - return update_queue_data() + queue.append({ + "id": current_task_id, + "image2video": kwargs["image2video"], + "params": kwargs.copy(), + "repeats": kwargs["repeat_generation"], + "length": kwargs["video_length"], + "steps": kwargs["num_inference_steps"], + "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), + "end_image_data_base64": pil_to_base64_uri(end_image_data, format="jpeg", quality=70) + }) + return update_queue_data(queue) -def move_up(selected_indices): +def move_up(queue, selected_indices): if not selected_indices or len(selected_indices) == 0: - return update_queue_data() + return update_queue_data(queue) idx = selected_indices[0] if isinstance(idx, list): idx = idx[0] idx = int(idx) with lock: if idx > 0: + idx += 1 queue[idx], queue[idx-1] = queue[idx-1], queue[idx] - return update_queue_data() + return update_queue_data(queue) -def move_down(selected_indices): +def move_down(queue, selected_indices): if not selected_indices or len(selected_indices) == 0: - return update_queue_data() + return update_queue_data(queue) idx = selected_indices[0] if isinstance(idx, list): idx = idx[0] idx = int(idx) with lock: + idx += 1 if idx < len(queue)-1: queue[idx], queue[idx+1] = queue[idx+1], queue[idx] - return update_queue_data() + return update_queue_data(queue) -def remove_task(selected_indices): +def remove_task(queue, selected_indices): if not selected_indices or len(selected_indices) == 0: - return update_queue_data() + return update_queue_data(queue) idx = selected_indices[0] if isinstance(idx, list): idx = idx[0] - idx = int(idx) + idx = int(idx) + 1 with lock: if idx < len(queue): if idx == 0: wan_model._interrupt = True del queue[idx] - return update_queue_data() + return update_queue_data(queue) -def update_queue_data(): - with lock: - data = [] - for item in queue: - truncated_prompt = (item['prompt'][:97] + '...') if len(item['prompt']) > 100 else item['prompt'] - full_prompt = item['prompt'].replace('"', '"') - prompt_cell = f'{truncated_prompt}' - start_img_uri = pil_to_base64_uri(item.get('start_image_data'), format="jpeg", quality=70) - end_img_uri = pil_to_base64_uri(item.get('end_image_data'), format="jpeg", quality=70) - thumbnail_size = "50px" - start_img_md = "" - end_img_md = "" - if start_img_uri: - start_img_md = f'Start' - if end_img_uri: - end_img_md = f'End' - data.append([ - item.get('status', "Starting"), - item.get('repeats', "0/0"), - item.get('progress', "0.0%"), - item.get('steps', ''), - item.get('time', '--'), - prompt_cell, - start_img_md, - end_img_md, - "↑", - "↓", - "✖" - ]) - return data + + +def get_queue_table(queue): + data = [] + if len(queue) == 1: + return data + + # def td(l, content, width =None): + # if width !=None: + # l.append("" + content + "") + # else: + # l.append("" + content + "") + + # data.append("") + + for i, item in enumerate(queue): + if i==0: + continue + truncated_prompt = (item['prompt'][:97] + '...') if len(item['prompt']) > 100 else item['prompt'] + full_prompt = item['prompt'].replace('"', '"') + prompt_cell = f'{truncated_prompt}' + start_img_uri =item.get('start_image_data_base64') + end_img_uri = item.get('end_image_data_base64') + thumbnail_size = "50px" + num_steps = item.get('steps') + length = item.get('length') + start_img_md = "" + end_img_md = "" + if start_img_uri: + start_img_md = f'Start' + if end_img_uri: + end_img_md = f'End' + # if i % 2 == 1: + # data.append("") + # else: + # data.append("") + + # td(data,str(item.get('repeats', "1")) ) + # td(data, prompt_cell, "100%") + # td(data, num_steps, "100%") + # td(data, start_img_md) + # td(data, end_img_md) + # td(data, "↑") + # td(data, "↓") + # td(data, "✖") + # data.append("") + # data.append("
QtyPromptSteps
") + # return ''.join(data) + + data.append([item.get('repeats', "1"), + prompt_cell, + length, + num_steps, + start_img_md, + end_img_md, + "↑", + "↓", + "✖" + ]) + return data +def update_queue_data(queue): + + data = get_queue_table(queue) + + # if len(data) == 0: + # return gr.HTML(visible=False) + # else: + # return gr.HTML(value=data, visible= True) + if len(data) == 0: + return gr.DataFrame(visible=False) + else: + return gr.DataFrame(value=data, visible= True) def create_html_progress_bar(percentage=0.0, text="Idle", is_idle=True): bar_class = "progress-bar-custom idle" if is_idle else "progress-bar-custom" @@ -306,35 +402,35 @@ def create_html_progress_bar(percentage=0.0, text="Idle", is_idle=True): """ return html -def refresh_progress(): - global current_task_id, progress_tracker, last_status_string - task_id_to_check = current_task_id - is_idle = True - status_string = "Starting..." - progress_percent = 0.0 - html_content = "" +# def refresh_progress(): +# global current_task_id, progress_tracker, last_status_string +# task_id_to_check = current_task_id +# is_idle = True +# status_string = "Starting..." +# progress_percent = 0.0 +# html_content = "" - with tracker_lock: - with lock: - processing_or_queued = any(item['state'] in ["Processing", "Queued"] for item in queue) - if task_id_to_check is not None: - progress_data = progress_tracker.get(task_id_to_check) - if progress_data: - is_idle = False - current_step = progress_data.get('current_step', 0) - total_steps = progress_data.get('total_steps', 0) - status = progress_data.get('status', "Starting...") - repeats = progress_data.get("repeats", "0/0") +# with tracker_lock: +# with lock: +# processing_or_queued = any(item['state'] in ["Processing", "Queued"] for item in queue) +# if task_id_to_check is not None: +# progress_data = progress_tracker.get(task_id_to_check) +# if progress_data: +# is_idle = False +# current_step = progress_data.get('current_step', 0) +# total_steps = progress_data.get('total_steps', 0) +# status = progress_data.get('status', "Starting...") +# repeats = progress_data.get("repeats", 1) - if total_steps > 0: - progress_float = min(1.0, max(0.0, float(current_step) / float(total_steps))) - progress_percent = progress_float * 100 - status_string = f"{status} [{repeats}] - {progress_percent:.1f}% complete ({current_step}/{total_steps} steps)" - else: - progress_percent = 0.0 - status_string = f"{status} [{repeats}] - Initializing..." - html_content = create_html_progress_bar(progress_percent, status_string, is_idle) - return gr.update(value=html_content) +# if total_steps > 0: +# progress_float = min(1.0, max(0.0, float(current_step) / float(total_steps))) +# progress_percent = progress_float * 100 +# status_string = f"{status} [{repeats}] - {progress_percent:.1f}% complete ({current_step}/{total_steps} steps)" +# else: +# progress_percent = 0.0 +# status_string = f"{status} [{repeats}] - Initializing..." +# html_content = create_html_progress_bar(progress_percent, status_string, is_idle) +# return gr.update(value=html_content) def update_generation_status(html_content): if(html_content): @@ -736,7 +832,6 @@ if args.i2v_1_3B: only_allow_edit_in_advanced = False lora_preselected_preset = args.lora_preset -lora_preselected_preset_for_i2v = use_image2video # if args.fast : #or args.fastest # transformer_filename_t2v = transformer_choices_t2v[2] # attention_mode="sage2" if "sage2" in attention_modes_supported else "sage" @@ -749,7 +844,6 @@ if args.compile: #args.fastest or lock_ui_compile = True model_filename = "" -lora_model_filename = "" #attention_mode="sage" #attention_mode="sage2" #attention_mode="flash" @@ -758,15 +852,12 @@ lora_model_filename = "" # compile = "transformer" def preprocess_loras(sd): - if not use_image2video: - return sd - - new_sd = {} first = next(iter(sd), None) if first == None: return sd - if not first.startswith("lora_unet_"): + if not first.startswith("lora_unet_"): return sd + new_sd = {} print("Converting Lora Safetensors format to Lora Diffusers format") alphas = {} repl_list = ["cross_attn", "self_attn", "ffn"] @@ -845,14 +936,14 @@ download_models(transformer_filename_i2v if use_image2video else transformer_fil def sanitize_file_name(file_name, rep =""): return file_name.replace("/",rep).replace("\\",rep).replace(":",rep).replace("|",rep).replace("?",rep).replace("<",rep).replace(">",rep).replace("\"",rep) -def extract_preset(lset_name, loras): +def extract_preset(image2video, lset_name, loras): loras_choices = [] loras_choices_files = [] loras_mult_choices = "" prompt ="" full_prompt ="" lset_name = sanitize_file_name(lset_name) - lora_dir = get_lora_dir(use_image2video) + lora_dir = get_lora_dir(image2video) if not lset_name.endswith(".lset"): lset_name_filename = os.path.join(lora_dir, lset_name + ".lset" ) else: @@ -923,7 +1014,7 @@ def setup_loras(i2v, transformer, lora_dir, lora_preselected_preset, split_line if not os.path.isfile(os.path.join(lora_dir, lora_preselected_preset + ".lset")): raise Exception(f"Unknown preset '{lora_preselected_preset}'") default_lora_preset = lora_preselected_preset - default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, _ , error = extract_preset(default_lora_preset, loras) + default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, _ , error = extract_preset(i2v, default_lora_preset, loras) if len(error) > 0: print(error[:200]) return loras, loras_names, loras_presets, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset @@ -1010,7 +1101,7 @@ def load_models(i2v): # kwargs["partialPinning"] = True elif profile == 3: kwargs["budgets"] = { "*" : "70%" } - offloadobj = offload.profile(pipe, profile_no= profile, compile = compile, quantizeTransformer = quantizeTransformer, loras = "transformer", **kwargs) + offloadobj = offload.profile(pipe, profile_no= profile, compile = compile, quantizeTransformer = quantizeTransformer, loras = "transformer", coTenantsMap= {}, **kwargs) if len(args.gpu) > 0: torch.set_default_device(args.gpu) @@ -1087,7 +1178,7 @@ def apply_changes( state, if gen_in_progress: yield "
Unable to change config when a generation is in progress
" return - global offloadobj, wan_model, loras, loras_names, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset, loras_presets + global offloadobj, wan_model, server_config, loras, loras_names, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset, loras_presets server_config = {"attention_mode" : attention_choice, "transformer_filename": transformer_choices_t2v[transformer_t2v_choice], "transformer_filename_i2v": transformer_choices_i2v[transformer_i2v_choice], @@ -1152,44 +1243,152 @@ def save_video(final_frames, output_path, fps=24): final_frames = (final_frames * 255).astype(np.uint8) ImageSequenceClip(list(final_frames), fps=fps).write_videofile(output_path, verbose= False, logger = None) -def build_callback(taskid, state, pipe, num_inference_steps, repeats): - start_time = time.time() - def update_progress(step_idx, _): - with tracker_lock: - step_idx += 1 - if state.get("abort", False): - # pipe._interrupt = True - phase = "Aborting" - elif step_idx == num_inference_steps: - phase = "VAE Decoding" - else: - phase = "Denoising" - elapsed = time.time() - start_time - progress_tracker[taskid] = { - 'current_step': step_idx, - 'total_steps': num_inference_steps, - 'start_time': start_time, - 'last_update': time.time(), - 'repeats': repeats, - 'status': phase - } - return update_progress -def refresh_gallery(state): - return gr.update(value=state.get("file_list", [])) +def get_gen_info(state): + cache = state.get("gen", None) + if cache == None: + cache = dict() + state["gen"] = cache + return cache + +def build_callback(state, pipe, progress, status, num_inference_steps): + def callback(step_idx, force_refresh, read_state = False): + gen = get_gen_info(state) + refresh_id = gen.get("refresh", -1) + if force_refresh or step_idx >= 0: + pass + else: + refresh_id = gen.get("refresh", -1) + if refresh_id < 0: + return + UI_refresh = state.get("refresh", 0) + if UI_refresh >= refresh_id: + return + + status = gen["progress_status"] + state["refresh"] = refresh_id + if read_state: + phase, step_idx = gen["progress_phase"] + else: + step_idx += 1 + if gen.get("abort", False): + # pipe._interrupt = True + phase = " - Aborting" + elif step_idx == num_inference_steps: + phase = " - VAE Decoding" + else: + phase = " - Denoising" + gen["progress_phase"] = (phase, step_idx) + status_msg = status + phase + if step_idx >= 0: + progress_args = [(step_idx , num_inference_steps) , status_msg , num_inference_steps] + else: + progress_args = [0, status_msg] + + progress(*progress_args) + gen["progress_args"] = progress_args + + return callback +def abort_generation(state): + gen = get_gen_info(state) + if "in_progress" in gen: + + gen["abort"] = True + gen["extra_orders"] = 0 + wan_model._interrupt= True + msg = "Processing Request to abort Current Generation" + gr.Info(msg) + return msg, gr.Button(interactive= False) + else: + return "", gr.Button(interactive= True) + +def is_gen_location(state): + gen = get_gen_info(state) + + gen_location = gen.get("location",None) + if gen_location == None: + return None + return state["image2video"] == gen_location + + +def refresh_gallery(state, msg): + gen = get_gen_info(state) + + if is_gen_location(state): + gen["last_msg"] = msg + file_list = gen.get("file_list", None) + choice = gen.get("selected",0) + in_progress = "in_progress" in gen + if in_progress: + if gen.get("last_selected", True): + choice = max(len(file_list) - 1,0) + + queue = gen.get("queue", []) + abort_interactive = not gen.get("abort", False) + if not in_progress or len(queue) == 0: + return gr.Gallery(selected_index=choice, value = file_list), gr.HTML("", visible= False), gr.Button(visible=True), gr.Button(visible=False), gr.Row(visible=False), update_queue_data(queue), gr.Button(interactive= abort_interactive) + else: + task = queue[0] + start_img_md = "" + 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'Start' + if end_img_uri: + end_img_md = f'End' + + label = f"Prompt of Video being Generated" + + html = "" + if start_img_md != "": + html += "" + if end_img_md != "": + html += "" + + html += "
" + prompt + "" + start_img_md + "" + end_img_md + "
" + html_output = gr.HTML(html, visible= True) + return gr.Gallery(selected_index=choice, value = file_list), html_output, gr.Button(visible=False), gr.Button(visible=True), gr.Row(visible=True), update_queue_data(queue), gr.Button(interactive= abort_interactive) + + + +def finalize_generation(state): + gen = get_gen_info(state) + choice = gen.get("selected",0) + if "in_progress" in gen: + del gen["in_progress"] + if gen.get("last_selected", True): + file_list = gen.get("file_list", []) + choice = len(file_list) - 1 + + + gen["extra_orders"] = 0 + time.sleep(0.2) + global gen_in_progress + gen_in_progress = False + return gr.Gallery(selected_index=choice), gr.Button(interactive= True), gr.Button(visible= True), gr.Button(visible= False), gr.Column(visible= False), gr.HTML(visible= False, value="") + def refresh_gallery_on_trigger(state): - if(state.get("update_gallery", False)): - state['update_gallery'] = False - return gr.update(value=state.get("file_list", [])) + gen = get_gen_info(state) + + if(gen.get("update_gallery", False)): + gen['update_gallery'] = False + return gr.update(value=gen.get("file_list", [])) def select_video(state , event_data: gr.EventData): data= event_data._data + gen = get_gen_info(state) + if data!=None: choice = data.get("index",0) - file_list = state.get("file_list", []) - state["last_selected"] = (choice + 1) >= len(file_list) - state["selected"] = choice + file_list = gen.get("file_list", []) + gen["last_selected"] = (choice + 1) >= len(file_list) + gen["selected"] = choice return def expand_slist(slist, num_inference_steps ): @@ -1221,6 +1420,7 @@ def convert_image(image): def generate_video( task_id, + progress, prompt, negative_prompt, resolution, @@ -1254,22 +1454,29 @@ def generate_video( ): global wan_model, offloadobj, reload_needed, last_model_type + gen = get_gen_info(state) + + file_list = gen["file_list"] + prompt_no = gen["prompt_no"] + file_model_needed = model_needed(image2video) - with lock: - queue_not_empty = len(queue) > 0 - if(last_model_type != image2video and (queue_not_empty or server_config.get("reload_model",1) == 2) and (file_model_needed != model_filename or reload_needed)): + # queue = gen.get("queue", []) + # with lock: + # queue_not_empty = len(queue) > 0 + # if(last_model_type != image2video and (queue_not_empty or server_config.get("reload_model",1) == 2) and (file_model_needed != model_filename or reload_needed)): + if file_model_needed != model_filename or reload_needed: del wan_model if offloadobj is not None: offloadobj.release() del offloadobj gc.collect() - print(f"Loading model {get_model_name(file_model_needed)}...") + yield f"Loading model {get_model_name(file_model_needed)}..." wan_model, offloadobj, trans = load_models(image2video) - print(f"Model loaded") + yield f"Model loaded" reload_needed= False if wan_model == None: - raise gr.Error("Unable to generate a Video while a new configuration is being applied.") + gr.Info("Unable to generate a Video while a new configuration is being applied.") if attention_mode == "auto": attn = get_auto_attention() elif attention_mode in attention_modes_supported: @@ -1278,26 +1485,15 @@ 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 - raw_resolution = resolution - width, height = resolution.split("x") - width, height = int(width), int(height) + + + if not image2video: + width, height = resolution.split("x") + width, height = int(width), int(height) if slg_switch == 0: slg_layers = None - if image2video: - if "480p" in model_filename and not "Fun" in model_filename 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") - return - resolution = str(width) + "*" + str(height) - if resolution not in ['720*1280', '1280*720', '480*832', '832*480']: - gr.Info(f"Resolution {resolution} not supported by image 2 video") - return - - if "1.3B" in model_filename and width * height > 848*480: - gr.Info("You must use the 14B model to generate videos with a resolution equivalent to 720P") - return - offload.shared_state["_attention"] = attn # VAE Tiling @@ -1321,16 +1517,7 @@ def generate_video( trans = wan_model.model - global gen_in_progress - gen_in_progress = True temp_filename = None - if image2video: - if 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: - input_image_or_video_path = None loras = state["loras"] if len(loras) > 0: @@ -1374,10 +1561,6 @@ def generate_video( raise gr.Error("Error while loading Loras: " + ", ".join(error_files)) seed = None if seed == -1 else seed # negative_prompt = "" # not applicable in the inference - - if "abort" in state: - del state["abort"] - state["in_progress"] = True enable_RIFLEx = RIFLEx_setting == 0 and video_length > (6* 16) or RIFLEx_setting == 1 # VAE Tiling @@ -1414,49 +1597,53 @@ def generate_video( if seed == None or seed <0: seed = random.randint(0, 999999999) - global file_list - clear_file_list = server_config.get("clear_file_list", 0) - file_list = state.get("file_list", []) - if clear_file_list > 0: - file_list_current_size = len(file_list) - keep_file_from = max(file_list_current_size - clear_file_list, 0) - files_removed = keep_file_from - choice = state.get("selected",0) - choice = max(choice- files_removed, 0) - file_list = file_list[ keep_file_from: ] - else: - file_list = [] - choice = 0 - state["selected"] = choice - state["file_list"] = file_list - - global save_path os.makedirs(save_path, exist_ok=True) video_no = 0 abort = False - repeats = f"{video_no}/{repeat_generation}" - callback = build_callback(task_id, state, trans, num_inference_steps, repeats) - offload.shared_state["callback"] = callback gc.collect() torch.cuda.empty_cache() wan_model._interrupt = False - for i in range(repeat_generation): + gen["abort"] = False + gen["prompt"] = prompt + repeat_no = 0 + extra_generation = 0 + while True: + extra_generation += gen.get("extra_orders",0) + gen["extra_orders"] = 0 + total_generation = repeat_generation + extra_generation + gen["total_generation"] = total_generation + if abort or repeat_no >= total_generation: + break + repeat_no +=1 + gen["repeat_no"] = repeat_no + prompts_max = gen["prompts_max"] + status = get_generation_status(prompt_no, prompts_max, repeat_no, total_generation) + + yield status + + gen["progress_status"] = status + gen["progress_phase"] = (" - Encoding Prompt", -1 ) + callback = build_callback(state, trans, progress, status, num_inference_steps) + progress_args = [0, status + " - Encoding Prompt"] + progress(*progress_args ) + gen["progress_args"] = progress_args + try: - with tracker_lock: - start_time = time.time() - progress_tracker[task_id] = { - 'current_step': 0, - 'total_steps': num_inference_steps, - 'start_time': start_time, - 'last_update': start_time, - 'repeats': f"{video_no}/{repeat_generation}", - 'status': "Encoding Prompt" - } + start_time = time.time() + # with tracker_lock: + # progress_tracker[task_id] = { + # 'current_step': 0, + # 'total_steps': num_inference_steps, + # 'start_time': start_time, + # 'last_update': start_time, + # 'repeats': repeat_generation, # f"{video_no}/{repeat_generation}", + # 'status': "Encoding Prompt" + # } if trans.enable_teacache: trans.teacache_counter = 0 trans.num_steps = num_inference_steps - trans.teacache_skipped_steps = 0 + trans.teacache_skipped_steps = 0 trans.previous_residual_uncond = None trans.previous_residual_cond = None @@ -1464,8 +1651,8 @@ def generate_video( if image2video: samples = wan_model.generate( prompt, - convert_image(image_to_continue), - convert_image(image_to_end) if image_to_end != None else None, + image_to_continue, + image_to_end if image_to_end != None else None, frame_num=(video_length // 4)* 4 + 1, max_area=MAX_AREA_CONFIGS[resolution], shift=flow_shift, @@ -1483,7 +1670,7 @@ def generate_video( slg_end = slg_end/100, cfg_star_switch = cfg_star_switch, cfg_zero_step = cfg_zero_step, - add_frames_for_end_image = not "Fun" in transformer_filename_i2v + add_frames_for_end_image = not "Fun" in transformer_filename_i2v, ) else: samples = wan_model.generate( @@ -1507,7 +1694,6 @@ def generate_video( cfg_zero_step = cfg_zero_step, ) except Exception as e: - gen_in_progress = False if temp_filename!= None and os.path.isfile(temp_filename): os.remove(temp_filename) offload.last_offload_obj.unload_all() @@ -1530,15 +1716,23 @@ def generate_video( if any( keyword in frame.name for keyword in keyword_list): VRAM_crash = True break + + _ , exc_value, exc_traceback = sys.exc_info() + state["prompt"] = "" if VRAM_crash: - raise gr.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.") + 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: - raise gr.Error(f"The generation of the video has encountered an error, please check your terminal for more information. '{s}'") + 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] + print('\n'.join(tb)) + raise gr.Error(new_error, print_exception= False) + finally: - with tracker_lock: - if task_id in progress_tracker: - del progress_tracker[task_id] + pass + # with tracker_lock: + # if task_id in progress_tracker: + # del progress_tracker[task_id] if trans.enable_teacache: print(f"Teacache Skipped Steps:{trans.teacache_skipped_steps}/{num_inference_steps}" ) @@ -1555,7 +1749,7 @@ def generate_video( end_time = time.time() abort = True state["prompt"] = "" - print(f"Video generation was aborted. Total Generation Time: {end_time-start_time:.1f}s") + # yield f"Video generation was aborted. Total Generation Time: {end_time-start_time:.1f}s" else: sample = samples.cpu() # video = rearrange(sample.cpu().numpy(), "c t h w -> t h w c") @@ -1574,7 +1768,7 @@ def generate_video( normalize=True, value_range=(-1, 1)) - configs = get_settings_dict(state, use_image2video, prompt, 0 if image_to_end == None else 1 , video_length, raw_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, 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, loras_mult_choices, tea_cache , tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, slg_start, slg_end, cfg_star_switch, cfg_zero_step) metadata_choice = server_config.get("metadata_choice","metadata") @@ -1596,9 +1790,159 @@ def generate_video( if temp_filename!= None and os.path.isfile(temp_filename): os.remove(temp_filename) - gen_in_progress = False offload.unload_loras_from_model(trans) +def prepare_generate_video(state): + if state.get("validate_success",0) != 1: + return gr.Button(visible= True), gr.Button(visible= False), gr.Column(visible= False) + else: + return gr.Button(visible= False), gr.Button(visible= True), gr.Column(visible= True) + + +def wait_tasks_done(state, progress=gr.Progress()): + + gen = get_gen_info(state) + gen_location = is_gen_location(state) + + last_msg = gen.get("last_msg", "") + if len(last_msg) > 0: + yield last_msg + + if gen_location == None or gen_location: + return gr.Text() + + + while True: + + msg = gen.get("last_msg", "") + if len(msg) > 0 and last_msg != msg: + yield msg + last_msg = msg + progress_args = gen.get("progress_args", None) + if progress_args != None: + progress(*progress_args) + + in_progress= gen.get("in_progress", False) + if not in_progress: + break + time.sleep(0.5) + + + +def process_tasks(state, progress=gr.Progress()): + gen = get_gen_info(state) + queue = gen.get("queue", []) + + if len(queue) == 0: + return + gen = get_gen_info(state) + gen["location"] = state["image2video"] + clear_file_list = server_config.get("clear_file_list", 0) + file_list = gen.get("file_list", []) + if clear_file_list > 0: + file_list_current_size = len(file_list) + keep_file_from = max(file_list_current_size - clear_file_list, 0) + files_removed = keep_file_from + choice = gen.get("selected",0) + choice = max(choice- files_removed, 0) + file_list = file_list[ keep_file_from: ] + else: + file_list = [] + choice = 0 + gen["selected"] = choice + gen["file_list"] = file_list + + start_time = time.time() + + global gen_in_progress + gen_in_progress = True + gen["in_progress"] = True + + prompt_no = 0 + while len(queue) > 0: + prompt_no += 1 + gen["prompt_no"] = prompt_no + task = queue[0] + task_id = task["id"] + params = task['params'] + iterator = iter(generate_video(task_id, progress, **params)) + while True: + try: + ok = False + status = next(iterator, "#") + if status == "#": + break + ok = True + except Exception as e: + _ , exc_value, exc_traceback = sys.exc_info() + raise exc_value.with_traceback(exc_traceback) + finally: + if not ok: + queue.clear() + yield status + + queue[:] = [item for item in queue if item['id'] != task['id']] + + gen["prompts_max"] = 0 + gen["prompt"] = "" + end_time = time.time() + if gen.get("abort"): + yield f"Video generation was aborted. Total Generation Time: {end_time-start_time:.1f}s" + else: + yield f"Total Generation Time: {end_time-start_time:.1f}s" + + +def get_generation_status(prompt_no, prompts_max, repeat_no, repeat_max): + if prompts_max == 1: + if repeat_max == 1: + return "Video" + else: + return f"Sample {repeat_no}/{repeat_max}" + else: + if repeat_max == 1: + return f"Prompt {prompt_no}/{prompts_max}" + else: + return f"Prompt {prompt_no}/{prompts_max}, Sample {repeat_no}/{repeat_max}" + + +refresh_id = 0 + +def get_new_refresh_id(): + global refresh_id + refresh_id += 1 + return refresh_id + +def update_status(state): + gen = get_gen_info(state) + prompt_no = gen["prompt_no"] + prompts_max = gen.get("prompts_max",0) + total_generation = gen["total_generation"] + repeat_no = gen["repeat_no"] + status = get_generation_status(prompt_no, prompts_max, repeat_no, total_generation) + gen["progress_status"] = status + gen["refresh"] = get_new_refresh_id() + + +def one_more_sample(state): + gen = get_gen_info(state) + extra_orders = gen.get("extra_orders", 0) + extra_orders += 1 + gen["extra_orders"] = extra_orders + in_progress = gen.get("in_progress", False) + if not in_progress : + return state + prompt_no = gen["prompt_no"] + prompts_max = gen.get("prompts_max",0) + total_generation = gen["total_generation"] + extra_orders + repeat_no = gen["repeat_no"] + status = get_generation_status(prompt_no, prompts_max, repeat_no, total_generation) + + + gen["progress_status"] = status + gen["refresh"] = get_new_refresh_id() + gr.Info(f"An extra sample generation is planned for a total of {total_generation} videos for this prompt") + + return state def get_new_preset_msg(advanced = True): if advanced: @@ -1650,7 +1994,7 @@ def save_lset(state, lset_name, loras_choices, loras_mult_choices, prompt, save_ lset_name_filename = lset_name + ".lset" - full_lset_name_filename = os.path.join(get_lora_dir(use_image2video), lset_name_filename) + full_lset_name_filename = os.path.join(get_lora_dir(state["image2video"]), lset_name_filename) with open(full_lset_name_filename, "w", encoding="utf-8") as writer: writer.write(json.dumps(lset, indent=4)) @@ -1667,7 +2011,7 @@ def save_lset(state, lset_name, loras_choices, loras_mult_choices, prompt, save_ def delete_lset(state, lset_name): loras_presets = state["loras_presets"] - lset_name_filename = os.path.join( get_lora_dir(use_image2video), sanitize_file_name(lset_name) + ".lset" ) + lset_name_filename = os.path.join( get_lora_dir(state["image2video"]), sanitize_file_name(lset_name) + ".lset" ) if len(lset_name) > 0 and lset_name != get_new_preset_msg(True) and lset_name != get_new_preset_msg(False): if not os.path.isfile(lset_name_filename): raise gr.Error(f"Preset '{lset_name}' not found ") @@ -1688,8 +2032,8 @@ def delete_lset(state, lset_name): def refresh_lora_list(state, lset_name, loras_choices): loras_names = state["loras_names"] prev_lora_names_selected = [ loras_names[int(i)] for i in loras_choices] - - loras, loras_names, loras_presets, _, _, _, _ = setup_loras(use_image2video, None, get_lora_dir(use_image2video), lora_preselected_preset, None) + image2video= state["image2video"] + loras, loras_names, loras_presets, _, _, _, _ = setup_loras(image2video, None, get_lora_dir(image2video), lora_preselected_preset, None) state["loras"] = loras state["loras_names"] = loras_names state["loras_presets"] = loras_presets @@ -1729,7 +2073,7 @@ def apply_lset(state, wizard_prompt_activated, lset_name, loras_choices, loras_m gr.Info("Please choose a preset in the list or create one") else: loras = state["loras"] - loras_choices, loras_mult_choices, preset_prompt, full_prompt, error = extract_preset(lset_name, loras) + loras_choices, loras_mult_choices, preset_prompt, full_prompt, error = extract_preset(state["image2video"], lset_name, loras) if len(error) > 0: gr.Info(error) else: @@ -1930,10 +2274,11 @@ def save_settings(state, prompt, image_prompt_type, video_length, resolution, nu if state.get("validate_success",0) != 1: return - ui_defaults = get_settings_dict(state, use_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, + 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, loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc, cfg_star_switch, cfg_zero_step) - defaults_filename = get_settings_file_name(use_image2video) + defaults_filename = get_settings_file_name(image2video) with open(defaults_filename, "w", encoding="utf-8") as f: json.dump(ui_defaults, f, indent=4) @@ -1976,7 +2321,12 @@ def generate_video_tab(image2video=False): state_dict["advanced"] = advanced state_dict["loras_model"] = filename - preset_to_load = lora_preselected_preset if lora_preselected_preset_for_i2v == image2video else "" + state_dict["image2video"] = image2video + gen = dict() + gen["queue"] = [] + state_dict["gen"] = gen + + preset_to_load = lora_preselected_preset if use_image2video == image2video else "" loras, loras_names, loras_presets, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset = setup_loras(image2video, None, get_lora_dir(image2video), preset_to_load, None) @@ -1989,7 +2339,7 @@ def generate_video_tab(image2video=False): launch_loras = [] launch_multis_str = "" - if len(default_lora_preset) > 0 and image2video == lora_preselected_preset_for_i2v: + if len(default_lora_preset) > 0 and image2video == use_image2video: launch_preset = default_lora_preset launch_prompt = default_lora_preset_prompt launch_loras = default_loras_choices @@ -2014,15 +2364,6 @@ def generate_video_tab(image2video=False): header = gr.Markdown(generate_header(model_filename, compile, attention_mode)) - with gr.Row(visible= image2video): - with gr.Row(scale =2): - gr.Markdown("Wan2GP's Lora Festival ! Press the following button to download i2v Remade Loras collection (and bonuses Loras).") - with gr.Row(scale =1): - download_loras_btn = gr.Button("---> Let the Lora's Festival Start !", scale =1) - with gr.Row(scale =1): - gr.Markdown("") - with gr.Row(visible= image2video) as download_status_row: - download_status = gr.Markdown() with gr.Row(): with gr.Column(): with gr.Column(visible=False, elem_id="image-modal-container") as modal_container: @@ -2250,89 +2591,112 @@ def generate_video_tab(image2video=False): cfg_zero_step = gr.Slider(-1, 39, value=ui_defaults.get("cfg_zero_step",-1), step=1, label="CFG Zero below this Layer (Extra Process)") with gr.Row(): - save_settings_btn = gr.Button("Set Settings as Default") + save_settings_btn = gr.Button("Set Settings as Default", visible = not args.lock_config) 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_progress_html = gr.HTML( label="Status", value="Idle", - elem_id="generation_progress_bar_container" + elem_id="generation_progress_bar_container", visible= False ) output = gr.Gallery( label="Generated videos", show_label=False, elem_id="gallery" , columns=[3], rows=[1], object_fit="contain", height=450, selected_index=0, interactive= False) generate_btn = gr.Button("Generate") - queue_df = gr.DataFrame( - headers=["Status", "Completed", "Progress", "Steps", "Time", "Prompt", "Start", "End", "", "", ""], - datatype=["str", "str", "str", "str", "str", "markdown", "markdown", "markdown", "str", "str", "str"], - interactive=False, - col_count=(11, "fixed"), - wrap=True, - value=update_queue_data, - every=1, - elem_id="queue_df" - ) - def handle_selection(evt: gr.SelectData): - if evt.index is None: - return gr.update(), gr.update(), gr.update(visible=False) - row_index, col_index = evt.index - cell_value = None - if col_index in [8, 9, 10]: - if col_index == 8: cell_value = "↑" - elif col_index == 9: cell_value = "↓" - elif col_index == 10: cell_value = "✖" - if col_index == 8: - new_df_data = move_up([row_index]) - return new_df_data, gr.update(), gr.update(visible=False) - elif col_index == 9: - new_df_data = move_down([row_index]) - return new_df_data, gr.update(), gr.update(visible=False) - elif col_index == 10: - new_df_data = remove_task([row_index]) - return new_df_data, gr.update(), gr.update(visible=False) - start_img_col_idx = 6 - end_img_col_idx = 7 - image_data_to_show = None - if col_index == start_img_col_idx: - with lock: - if row_index < len(queue): - image_data_to_show = queue[row_index].get('start_image_data') - elif col_index == end_img_col_idx: - with lock: - if row_index < len(queue): - image_data_to_show = queue[row_index].get('end_image_data') + add_to_queue_btn = gr.Button("Add New Prompt To Queue", visible = False) - if image_data_to_show: - return gr.update(), gr.update(value=image_data_to_show), gr.update(visible=True) - else: - return gr.update(), gr.update(), gr.update(visible=False) - selected_indices = gr.State([]) - queue_df.select( - fn=handle_selection, - inputs=None, - outputs=[queue_df, modal_image_display, modal_container], - ) - gallery_update_trigger.change( - fn=refresh_gallery_on_trigger, - inputs=[state], - outputs=[output] - ) - queue_df.change( - fn=refresh_gallery, - inputs=[state], - outputs=[gallery_update_trigger] - ).then( - fn=refresh_progress, - inputs=None, - outputs=[progress_update_trigger] - ) - progress_update_trigger.change( - fn=update_generation_status, - inputs=[progress_update_trigger], - outputs=[gen_progress_html], - show_progress="hidden" - ) + with gr.Column(visible= False) as current_gen_column: + with gr.Row(): + gen_info = gr.HTML(visible=False, min_height=1) + with gr.Row(): + onemore_btn = gr.Button("One More Sample Please !") + 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"], + interactive=False, + col_count=(9, "fixed"), + wrap=True, + value=[], + visible= False, + # every=1, + elem_id="queue_df" + ) + # queue_df = gr.HTML("", + # visible= False, + # elem_id="queue_df" + # ) + + def handle_selection(state, evt: gr.SelectData): + gen = get_gen_info(state) + queue = gen.get("queue", []) + + if evt.index is None: + return gr.update(), gr.update(), gr.update(visible=False) + row_index, col_index = evt.index + cell_value = None + if col_index in [6, 7, 8]: + if col_index == 6: cell_value = "↑" + elif col_index == 7: cell_value = "↓" + elif col_index == 8: cell_value = "✖" + if col_index == 6: + new_df_data = move_up(queue, [row_index]) + return new_df_data, gr.update(), gr.update(visible=False) + elif col_index == 7: + new_df_data = move_down(queue, [row_index]) + return new_df_data, gr.update(), gr.update(visible=False) + elif col_index == 8: + new_df_data = remove_task(queue, [row_index]) + gen["prompts_max"] = gen.get("prompts_max",0) - 1 + update_status(state) + return new_df_data, gr.update(), gr.update(visible=False) + start_img_col_idx = 4 + end_img_col_idx = 5 + image_data_to_show = None + if col_index == start_img_col_idx: + with lock: + if row_index < len(queue): + image_data_to_show = queue[row_index].get('start_image_data') + elif col_index == end_img_col_idx: + with lock: + if row_index < len(queue): + image_data_to_show = queue[row_index].get('end_image_data') + + if image_data_to_show: + return gr.update(), gr.update(value=image_data_to_show), gr.update(visible=True) + else: + return gr.update(), gr.update(), gr.update(visible=False) + selected_indices = gr.State([]) + queue_df.select( + fn=handle_selection, + inputs=state, + outputs=[queue_df, modal_image_display, modal_container], + ) + # gallery_update_trigger.change( + # fn=refresh_gallery_on_trigger, + # inputs=[state], + # outputs=[output] + # ) + # queue_df.change( + # fn=refresh_gallery, + # inputs=[state], + # outputs=[gallery_update_trigger] + # ).then( + # fn=refresh_progress, + # inputs=None, + # outputs=[progress_update_trigger] + # ) + progress_update_trigger.change( + fn=update_generation_status, + inputs=[progress_update_trigger], + outputs=[gen_progress_html], + 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, loras_choices, loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, @@ -2348,53 +2712,114 @@ def generate_video_tab(image2video=False): ) refresh_lora_btn.click(refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) refresh_lora_btn2.click(refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) - download_loras_btn.click(fn=download_loras, inputs=[], outputs=[download_status_row, download_status, presets_column, loras_column]).then(fn=refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) output.select(select_video, state, None ) - generate_btn.click( - fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , outputs= [prompt] + 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, + 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, + 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, + inputs = [state, gen_status], + outputs = [output, gen_info, generate_btn, add_to_queue_btn, current_gen_column, queue_df, abort_btn] + ) + + abort_btn.click(abort_generation, [state], [gen_status, abort_btn] ) #.then(refresh_gallery, inputs = [state, gen_info], outputs = [output, gen_info, queue_df] ) + onemore_btn.click(fn=one_more_sample,inputs=[state], outputs= [state]) + + + gen_inputs=[ + prompt, + negative_prompt, + resolution, + video_length, + seed, + num_inference_steps, + guidance_scale, + flow_shift, + embedded_guidance_scale, + repeat_generation, + multi_images_gen_type, + tea_cache_setting, + tea_cache_start_step_perc, + loras_choices, + loras_mult_choices, + image_prompt_type_radio, + image_to_continue, + image_to_end, + video_to_continue, + max_frames, + RIFLEx_setting, + slg_switch, + slg_layers, + slg_start_perc, + slg_end_perc, + cfg_star_switch, + cfg_zero_step, + state, + gr.State(image2video) + ] + + generate_btn.click(fn=validate_wizard_prompt, + inputs= [state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , + outputs= [prompt] + ).then(fn=process_prompt_and_add_tasks, + inputs = gen_inputs, + outputs= queue_df + ).then(fn=prepare_generate_video, + inputs= [state], + outputs= [generate_btn, add_to_queue_btn, current_gen_column], + ).then(fn=process_tasks, + 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] + ) + + add_to_queue_btn.click(fn=validate_wizard_prompt, + inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , + outputs= [prompt] ).then( fn=process_prompt_and_add_tasks, - inputs=[ - prompt, - negative_prompt, - resolution, - video_length, - seed, - num_inference_steps, - guidance_scale, - flow_shift, - embedded_guidance_scale, - repeat_generation, - multi_images_gen_type, - tea_cache_setting, - tea_cache_start_step_perc, - loras_choices, - loras_mult_choices, - image_prompt_type_radio, - image_to_continue, - image_to_end, - video_to_continue, - max_frames, - RIFLEx_setting, - slg_switch, - slg_layers, - slg_start_perc, - slg_end_perc, - cfg_star_switch, - cfg_zero_step, - state, - gr.State(image2video) - ], + inputs = gen_inputs, outputs=queue_df + ).then( + fn=update_status, + inputs = [state], ) + + close_modal_button.click( lambda: gr.update(visible=False), inputs=[], outputs=[modal_container] ) - return loras_column, loras_choices, presets_column, lset_name, header, state + 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): + with gr.Row(): + with gr.Row(scale =2): + gr.Markdown("Wan2GP's Lora Festival ! Press the following button to download i2v Remade Loras collection (and bonuses Loras).") + with gr.Row(scale =1): + download_loras_btn = gr.Button("---> Let the Lora's Festival Start !", scale =1) + with gr.Row(scale =1): + gr.Markdown("") + with gr.Row() as download_status_row: + download_status = gr.Markdown() + + download_loras_btn.click(fn=download_loras, inputs=[], outputs=[download_status_row, download_status, presets_column, loras_column]).then(fn=refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) + + def generate_configuration_tab(): state_dict = {} state = gr.State(state_dict) @@ -2411,7 +2836,7 @@ def generate_configuration_tab(): value= index, label="Transformer model for Text to Video", interactive= not lock_ui_transformer, - visible=True #not use_image2video + visible=True ) index = transformer_choices_i2v.index(transformer_filename_i2v) index = 0 if index ==0 else index @@ -2428,7 +2853,7 @@ def generate_configuration_tab(): value= index, label="Transformer model for Image to Video", interactive= not lock_ui_transformer, - visible = True # use_image2video, + visible = True, ) index = text_encoder_choices.index(text_encoder_filename) index = 0 if index ==0 else index @@ -2524,7 +2949,7 @@ def generate_configuration_tab(): reload_choice = gr.Dropdown( choices=[ ("When changing tabs", 1), - ("When pressing generate", 2), + ("When pressing Generate", 2), ], value=server_config.get("reload_model",2), label="Reload model" @@ -2577,19 +3002,46 @@ def generate_about_tab(): gr.Markdown("- Remade_AI : for creating their awesome Loras collection") -def on_tab_select(t2v_state, i2v_state, evt: gr.SelectData): - global lora_model_filename, use_image2video - +def on_tab_select(global_state, t2v_state, i2v_state, evt: gr.SelectData): t2v_header = generate_header(transformer_filename_t2v, compile, attention_mode) i2v_header = generate_header(transformer_filename_i2v, compile, attention_mode) new_t2v = evt.index == 0 new_i2v = evt.index == 1 - use_image2video = new_i2v + i2v_light_sync = gr.Text() + 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 + + + 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) : + if new_i2v: + i2v_full_sync = gr.Text(str(time.time())) + else: + t2v_full_sync = gr.Text(str(time.time())) + else: + if new_i2v: + i2v_light_sync = gr.Text(str(time.time())) + 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): - with lock: - queue_empty = len(queue) == 0 + queue = gen.get("queue", []) + + queue_empty = len(queue) == 0 if queue_empty: global wan_model, offloadobj if wan_model is not None: @@ -2599,7 +3051,7 @@ def on_tab_select(t2v_state, i2v_state, evt: gr.SelectData): wan_model = None gc.collect() torch.cuda.empty_cache() - wan_model, offloadobj, trans = load_models(use_image2video) + wan_model, offloadobj, trans = load_models(new_i2v) del trans if new_t2v or new_i2v: @@ -2625,11 +3077,15 @@ def on_tab_select(t2v_state, i2v_state, evt: gr.SelectData): gr.Column(visible= visible), gr.Dropdown(choices=lset_choices, value=get_new_preset_msg(advanced), visible=visible), t2v_header, + t2v_light_sync, + t2v_full_sync, gr.Column(), gr.Dropdown(), gr.Column(), gr.Dropdown(), - i2v_header, + gr.Markdown(), + gr.Text(), + gr.Text(), ] else: return [ @@ -2637,16 +3093,21 @@ def on_tab_select(t2v_state, i2v_state, evt: gr.SelectData): gr.Dropdown(), gr.Column(), gr.Dropdown(), - t2v_header, + gr.Markdown(), + gr.Text(), + gr.Text(), + gr.Text(), gr.Column(visible= visible), gr.Dropdown(choices=new_loras_choices, visible=visible, value=[]), gr.Column(visible= visible), gr.Dropdown(choices=lset_choices, value=get_new_preset_msg(advanced), visible=visible), i2v_header, + i2v_light_sync, + i2v_full_sync, ] - return [gr.Column(), gr.Dropdown(), gr.Column(), gr.Dropdown(), t2v_header, - gr.Column(), gr.Dropdown(), gr.Column(), gr.Dropdown(), i2v_header] + return [gr.Column(), gr.Dropdown(), gr.Column(), gr.Dropdown(), t2v_header, t2v_light_sync, t2v_full_sync, + gr.Column(), gr.Dropdown(), gr.Column(), gr.Dropdown(), i2v_header, i2v_light_sync, i2v_full_sync] def create_demo(): @@ -2706,112 +3167,112 @@ def create_demo(): overflow: hidden; text-overflow: ellipsis; } - #queue_df td:nth-child(-n+5) { - cursor: default !important; - pointer-events: none; - } - #queue_df td:nth-child(6) { - cursor: default !important; - } - #queue_df th { - pointer-events: none; - text-align: center; - vertical-align: middle; - } - #queue_df table { - width: 100%; - overflow: hidden !important; - } - #queue_df::-webkit-scrollbar { - display: none !important; - } - #queue_df { - scrollbar-width: none !important; - -ms-overflow-style: none !important; - } - #queue_df th:nth-child(1), - #queue_df td:nth-child(1) { - width: 90px; - text-align: center; - vertical-align: middle; - } - #queue_df th:nth-child(1) { - font-size: 0.8em; - } - #queue_df th:nth-child(2), - #queue_df td:nth-child(2) { - width: 85px; - text-align: center; - vertical-align: middle; - } - #queue_df th:nth-child(2) { - font-size: 0.5em; - } - #queue_df th:nth-child(3), - #queue_df td:nth-child(3) { - width: 75px; - text-align: center; - vertical-align: middle; - } - #queue_df th:nth-child(3) { - font-size: 0.6em; - } - #queue_df th:nth-child(4), - #queue_df td:nth-child(4) { - width: 65px; - text-align: center; - white-space: nowrap; - } - #queue_df th:nth-child(4) { - font-size: 0.9em; - } - #queue_df th:nth-child(5), - #queue_df td:nth-child(5) { - width: 60px; - text-align: center; - white-space: nowrap; - } - #queue_df th:nth-child(6), - #queue_df td:nth-child(6) { - width: auto; - text-align: center; - white-space: normal; - } - #queue_df th:nth-child(6) { - font-size: 0.8em; - } - #queue_df th:nth-child(7), #queue_df td:nth-child(7), - #queue_df th:nth-child(8), #queue_df td:nth-child(8) { - width: 60px; - text-align: center; - vertical-align: middle; - } - #queue_df td:nth-child(7) img, - #queue_df td:nth-child(8) img { - max-width: 50px; - max-height: 50px; - object-fit: contain; - display: block; - margin: auto; - cursor: pointer; - } - #queue_df th:nth-child(9), #queue_df td:nth-child(9), - #queue_df th:nth-child(10), #queue_df td:nth-child(10), - #queue_df th:nth-child(11), #queue_df td:nth-child(11) { - width: 20px; - padding: 2px !important; - cursor: pointer; - text-align: center; - font-weight: bold; - vertical-align: middle; - } - #queue_df td:nth-child(7):hover, - #queue_df td:nth-child(8):hover, - #queue_df td:nth-child(9):hover, - #queue_df td:nth-child(10):hover, - #queue_df td:nth-child(11):hover { - background-color: #e0e0e0; - } + # #queue_df td:nth-child(-n+5) { + # cursor: default !important; + # pointer-events: none; + # } + # #queue_df td:nth-child(6) { + # cursor: default !important; + # } + # #queue_df th { + # pointer-events: none; + # text-align: center; + # vertical-align: middle; + # } + # #queue_df table { + # width: 100%; + # overflow: hidden !important; + # } + # #queue_df::-webkit-scrollbar { + # display: none !important; + # } + # #queue_df { + # scrollbar-width: none !important; + # -ms-overflow-style: none !important; + # } + # #queue_df th:nth-child(1), + # #queue_df td:nth-child(1) { + # width: 90px; + # text-align: center; + # vertical-align: middle; + # } + # #queue_df th:nth-child(1) { + # font-size: 0.8em; + # } + # #queue_df th:nth-child(2), + # #queue_df td:nth-child(2) { + # width: 85px; + # text-align: center; + # vertical-align: middle; + # } + # #queue_df th:nth-child(2) { + # font-size: 0.5em; + # } + # #queue_df th:nth-child(3), + # #queue_df td:nth-child(3) { + # width: 75px; + # text-align: center; + # vertical-align: middle; + # } + # #queue_df th:nth-child(3) { + # font-size: 0.6em; + # } + # #queue_df th:nth-child(4), + # #queue_df td:nth-child(4) { + # width: 65px; + # text-align: center; + # white-space: nowrap; + # } + # #queue_df th:nth-child(4) { + # font-size: 0.9em; + # } + # #queue_df th:nth-child(5), + # #queue_df td:nth-child(5) { + # width: 60px; + # text-align: center; + # white-space: nowrap; + # } + # #queue_df th:nth-child(6), + # #queue_df td:nth-child(6) { + # width: auto; + # text-align: center; + # white-space: normal; + # } + # #queue_df th:nth-child(6) { + # font-size: 0.8em; + # } + # #queue_df th:nth-child(7), #queue_df td:nth-child(7), + # #queue_df th:nth-child(8), #queue_df td:nth-child(8) { + # width: 60px; + # text-align: center; + # vertical-align: middle; + # } + # #queue_df td:nth-child(7) img, + # #queue_df td:nth-child(8) img { + # max-width: 50px; + # max-height: 50px; + # object-fit: contain; + # display: block; + # margin: auto; + # cursor: pointer; + # } + # #queue_df th:nth-child(9), #queue_df td:nth-child(9), + # #queue_df th:nth-child(10), #queue_df td:nth-child(10), + # #queue_df th:nth-child(11), #queue_df td:nth-child(11) { + # width: 20px; + # padding: 2px !important; + # cursor: pointer; + # text-align: center; + # font-weight: bold; + # vertical-align: middle; + # } + # #queue_df td:nth-child(7):hover, + # #queue_df td:nth-child(8):hover, + # #queue_df td:nth-child(9):hover, + # #queue_df td:nth-child(10):hover, + # #queue_df td:nth-child(11):hover { + # background-color: #e0e0e0; + # } #image-modal-container { position: fixed; top: 0; @@ -2893,8 +3354,8 @@ def create_demo(): pointer-events: none; } """ - with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="sky", neutral_hue="slate", text_size="md")) as demo: - gr.Markdown("

Wan 2.1GP v3.3 by DeepBeepMeep (Updates)

") + with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="sky", neutral_hue="slate", text_size="md"), title= "Wan2GP") as demo: + gr.Markdown("

Wan 2.1GP v3.4 by DeepBeepMeep (Updates)

") gr.Markdown("Welcome to Wan 2.1GP a super fast and low VRAM AI Video Generator !") with gr.Accordion("Click here for some Info on how to use Wan2GP", open = False): @@ -2904,30 +3365,34 @@ def create_demo(): gr.Markdown("- 1280 x 720 with a 14B model: 80 frames (5s): 11 GB of VRAM") gr.Markdown("It is not recommmended to generate a video longer than 8s (128 frames) even if there is still some VRAM left as some artifacts may appear") gr.Markdown("Please note that if your turn on compilation, the first denoising step of the first video generation will be slow due to the compilation. Therefore all your tests should be done with compilation turned off.") - + global_dict = {} + global_dict["last_tab_was_image2video"] = use_image2video + global_state = gr.State(global_dict) with gr.Tabs(selected="i2v" if use_image2video else "t2v") as main_tabs: with gr.Tab("Text To Video", id="t2v") as t2v_tab: - t2v_loras_column, t2v_loras_choices, t2v_presets_column, t2v_lset_name, t2v_header, t2v_state = generate_video_tab() + t2v_loras_column, t2v_loras_choices, t2v_presets_column, t2v_lset_name, t2v_header, t2v_light_sync, t2v_full_sync, t2v_state = generate_video_tab(False) with gr.Tab("Image To Video", id="i2v") as i2v_tab: - i2v_loras_column, i2v_loras_choices, i2v_presets_column, i2v_lset_name, i2v_header, i2v_state = generate_video_tab(True) + 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) with gr.Tab("Configuration"): generate_configuration_tab() with gr.Tab("About"): generate_about_tab() main_tabs.select( fn=on_tab_select, - inputs=[t2v_state, i2v_state], + inputs=[global_state, t2v_state, i2v_state], outputs=[ - t2v_loras_column, t2v_loras_choices, t2v_presets_column, t2v_lset_name, t2v_header, - i2v_loras_column, i2v_loras_choices, i2v_presets_column, i2v_lset_name, i2v_header + t2v_loras_column, t2v_loras_choices, t2v_presets_column, t2v_lset_name, t2v_header, t2v_light_sync, t2v_full_sync, + i2v_loras_column, i2v_loras_choices, i2v_presets_column, i2v_lset_name, i2v_header, i2v_light_sync, i2v_full_sync ] ) return demo if __name__ == "__main__": - threading.Thread(target=runner, daemon=True).start() + # threading.Thread(target=runner, daemon=True).start() os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" server_port = int(args.server_port) if os.name == "nt": diff --git a/wan/image2video.py b/wan/image2video.py index 2ea4310..a71e9d4 100644 --- a/wan/image2video.py +++ b/wan/image2video.py @@ -40,7 +40,12 @@ def optimized_scale(positive_flat, negative_flat): st_star = dot_product / squared_norm return st_star - + +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) + class WanI2V: @@ -90,7 +95,6 @@ class WanI2V: self.num_train_timesteps = config.num_train_timesteps self.param_dtype = config.param_dtype - shard_fn = partial(shard_model, device_id=device_id) self.text_encoder = T5EncoderModel( text_len=config.text_len, @@ -208,16 +212,16 @@ class WanI2V: - H: Frame height (from max_area) - W: Frame width from max_area) """ - img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device) + img = TF.to_tensor(img) lat_frames = int((frame_num - 1) // self.vae_stride[0] + 1) any_end_frame = img2 !=None if any_end_frame: any_end_frame = True - img2 = TF.to_tensor(img2).sub_(0.5).div_(0.5).to(self.device) + img2 = TF.to_tensor(img2) if add_frames_for_end_image: frame_num +=1 lat_frames = int((frame_num - 2) // self.vae_stride[0] + 2) - + h, w = img.shape[1:] aspect_ratio = h / w lat_h = round( @@ -229,6 +233,15 @@ class WanI2V: h = lat_h * self.vae_stride[1] w = lat_w * self.vae_stride[2] + clip_image_size = self.clip.model.image_size + img_interpolated = resize_lanczos(img, h, w).sub_(0.5).div_(0.5).unsqueeze(0).transpose(0,1).to(self.device) + img = resize_lanczos(img, clip_image_size, clip_image_size) + img = img.sub_(0.5).div_(0.5).to(self.device) + if img2!= None: + img_interpolated2 = resize_lanczos(img2, h, w).sub_(0.5).div_(0.5).unsqueeze(0).transpose(0,1).to(self.device) + img2 = resize_lanczos(img2, clip_image_size, clip_image_size) + img2 = img2.sub_(0.5).div_(0.5).to(self.device) + max_seq_len = lat_frames * lat_h * lat_w // ( self.patch_size[1] * self.patch_size[2]) max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size @@ -273,21 +286,32 @@ class WanI2V: from mmgp import offload + + # img_interpolated.save('aaa.png') + + # img_interpolated = torch.from_numpy(np.array(img_interpolated).astype(np.float32) / 255.0).movedim(-1, 0) + + # img_interpolated = torch.nn.functional.interpolate(img[None].cpu(), size=(h, w), mode='lanczos') + # img_interpolated = img_interpolated.squeeze(0).transpose(0,2).transpose(1,0) + # img_interpolated = img_interpolated.clamp(-1, 1) + # img_interpolated = (img_interpolated + 1)/2 + # img_interpolated = (img_interpolated*255).type(torch.uint8) + # img_interpolated = img_interpolated.cpu().numpy() + # xxx = Image.fromarray(img_interpolated, 'RGB') + # xxx.save('my.png') + offload.last_offload_obj.unload_all() if any_end_frame: - img_interpolated = torch.nn.functional.interpolate(img[None].cpu(), size=(h, w), mode='bicubic').transpose(0, 1).to(torch.bfloat16) - img2_interpolated = torch.nn.functional.interpolate(img2[None].cpu(), size=(h, w), mode='bicubic').transpose(0, 1).to(torch.bfloat16) mean2 = 0 enc= torch.concat([ img_interpolated, - torch.full( (3, frame_num-2, h, w), mean2, device="cpu", dtype= torch.bfloat16), - img2_interpolated, + torch.full( (3, frame_num-2, h, w), mean2, device=self.device, dtype= torch.bfloat16), + img_interpolated2, ], dim=1).to(self.device) else: enc= torch.concat([ - torch.nn.functional.interpolate( - img[None].cpu(), size=(h, w), mode='bicubic').transpose(0, 1).to(torch.bfloat16), - torch.zeros(3, frame_num-1, h, w, device="cpu", dtype= torch.bfloat16) + img_interpolated, + torch.zeros(3, frame_num-1, h, w, device=self.device, dtype= torch.bfloat16) ], dim=1).to(self.device) lat_y = self.vae.encode([enc], VAE_tile_size, any_end_frame= any_end_frame and add_frames_for_end_image)[0] @@ -333,7 +357,8 @@ class WanI2V: 'seq_len': max_seq_len, 'y': [y], 'freqs' : freqs, - 'pipeline' : self + 'pipeline' : self, + 'callback' : callback } arg_null = { @@ -342,7 +367,8 @@ class WanI2V: 'seq_len': max_seq_len, 'y': [y], 'freqs' : freqs, - 'pipeline' : self + 'pipeline' : self, + 'callback' : callback } arg_both= { @@ -352,7 +378,8 @@ class WanI2V: 'seq_len': max_seq_len, 'y': [y], 'freqs' : freqs, - 'pipeline' : self + 'pipeline' : self, + 'callback' : callback } if offload_model: @@ -363,7 +390,7 @@ class WanI2V: # self.model.to(self.device) if callback != None: - callback(-1, None) + callback(-1, True) for i, t in enumerate(tqdm(timesteps)): offload.set_step_no_for_lora(self.model, i) @@ -437,7 +464,7 @@ class WanI2V: del timestep if callback is not None: - callback(i, latent) + callback(i, False) x0 = [latent.to(self.device, dtype=torch.bfloat16)] @@ -451,7 +478,7 @@ class WanI2V: video = self.vae.decode(x0, VAE_tile_size, any_end_frame= any_end_frame and add_frames_for_end_image)[0] if any_end_frame and add_frames_for_end_image: - # video[:, -1:] = img2_interpolated + # video[:, -1:] = img_interpolated2 video = video[:, :-1] else: diff --git a/wan/modules/model.py b/wan/modules/model.py index 3e2ea5c..2daa00c 100644 --- a/wan/modules/model.py +++ b/wan/modules/model.py @@ -704,6 +704,7 @@ class WanModel(ModelMixin, ConfigMixin): is_uncond=False, max_steps = 0, slg_layers=None, + callback = None, ): r""" Forward pass through the diffusion model @@ -835,12 +836,10 @@ class WanModel(ModelMixin, ConfigMixin): freqs=freqs, # context=context, context_lens=context_lens) - for block_idx, block in enumerate(self.blocks): offload.shared_state["layer"] = block_idx - if "refresh" in offload.shared_state: - del offload.shared_state["refresh"] - offload.shared_state["callback"](-1, -1, True) + if callback != None: + callback(-1, False, True) if pipeline._interrupt: if joint_pass: return None, None diff --git a/wan/text2video.py b/wan/text2video.py index 088a9c7..cdcbd4f 100644 --- a/wan/text2video.py +++ b/wan/text2video.py @@ -268,7 +268,7 @@ class WanT2V: if self.model.enable_teacache: self.model.compute_teacache_threshold(self.model.teacache_start_step, timesteps, self.model.teacache_multiplier) if callback != None: - callback(-1, None) + callback(-1, True) for i, t in enumerate(tqdm(timesteps)): latent_model_input = latents slg_layers_local = None @@ -322,7 +322,7 @@ class WanT2V: del temp_x0 if callback is not None: - callback(i, latents) + callback(i, False) x0 = latents if offload_model: From 19641e423cca5522c783bf19c4372e096502113d Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Tue, 1 Apr 2025 22:25:08 +0200 Subject: [PATCH 30/69] Fixed bad error handling that broke the queue --- gradio_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gradio_server.py b/gradio_server.py index e21c5cb..d2dde1b 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1870,9 +1870,9 @@ def process_tasks(state, progress=gr.Progress()): try: ok = False status = next(iterator, "#") + ok = True if status == "#": break - ok = True except Exception as e: _ , exc_value, exc_traceback = sys.exc_info() raise exc_value.with_traceback(exc_traceback) From 897aaf21f0258fbbd39c2f5b2132467b873b4745 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Wed, 2 Apr 2025 08:14:46 +1100 Subject: [PATCH 31/69] fix styling --- gradio_server.py | 194 +++++++++++++++++++++-------------------------- 1 file changed, 88 insertions(+), 106 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index d2dde1b..f511102 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -3167,112 +3167,94 @@ def create_demo(): overflow: hidden; text-overflow: ellipsis; } - # #queue_df td:nth-child(-n+5) { - # cursor: default !important; - # pointer-events: none; - # } - # #queue_df td:nth-child(6) { - # cursor: default !important; - # } - # #queue_df th { - # pointer-events: none; - # text-align: center; - # vertical-align: middle; - # } - # #queue_df table { - # width: 100%; - # overflow: hidden !important; - # } - # #queue_df::-webkit-scrollbar { - # display: none !important; - # } - # #queue_df { - # scrollbar-width: none !important; - # -ms-overflow-style: none !important; - # } - # #queue_df th:nth-child(1), - # #queue_df td:nth-child(1) { - # width: 90px; - # text-align: center; - # vertical-align: middle; - # } - # #queue_df th:nth-child(1) { - # font-size: 0.8em; - # } - # #queue_df th:nth-child(2), - # #queue_df td:nth-child(2) { - # width: 85px; - # text-align: center; - # vertical-align: middle; - # } - # #queue_df th:nth-child(2) { - # font-size: 0.5em; - # } - # #queue_df th:nth-child(3), - # #queue_df td:nth-child(3) { - # width: 75px; - # text-align: center; - # vertical-align: middle; - # } - # #queue_df th:nth-child(3) { - # font-size: 0.6em; - # } - # #queue_df th:nth-child(4), - # #queue_df td:nth-child(4) { - # width: 65px; - # text-align: center; - # white-space: nowrap; - # } - # #queue_df th:nth-child(4) { - # font-size: 0.9em; - # } - # #queue_df th:nth-child(5), - # #queue_df td:nth-child(5) { - # width: 60px; - # text-align: center; - # white-space: nowrap; - # } - # #queue_df th:nth-child(6), - # #queue_df td:nth-child(6) { - # width: auto; - # text-align: center; - # white-space: normal; - # } - # #queue_df th:nth-child(6) { - # font-size: 0.8em; - # } - # #queue_df th:nth-child(7), #queue_df td:nth-child(7), - # #queue_df th:nth-child(8), #queue_df td:nth-child(8) { - # width: 60px; - # text-align: center; - # vertical-align: middle; - # } - # #queue_df td:nth-child(7) img, - # #queue_df td:nth-child(8) img { - # max-width: 50px; - # max-height: 50px; - # object-fit: contain; - # display: block; - # margin: auto; - # cursor: pointer; - # } - # #queue_df th:nth-child(9), #queue_df td:nth-child(9), - # #queue_df th:nth-child(10), #queue_df td:nth-child(10), - # #queue_df th:nth-child(11), #queue_df td:nth-child(11) { - # width: 20px; - # padding: 2px !important; - # cursor: pointer; - # text-align: center; - # font-weight: bold; - # vertical-align: middle; - # } - # #queue_df td:nth-child(7):hover, - # #queue_df td:nth-child(8):hover, - # #queue_df td:nth-child(9):hover, - # #queue_df td:nth-child(10):hover, - # #queue_df td:nth-child(11):hover { - # background-color: #e0e0e0; - # } + #queue_df th { + pointer-events: none; + text-align: center; + vertical-align: middle; + } + #queue_df table { + width: 100%; + overflow: hidden !important; + } + #queue_df::-webkit-scrollbar { + display: none !important; + } + #queue_df { + scrollbar-width: none !important; + -ms-overflow-style: none !important; + } + .selection-button { + display: none; + } + .cell-selected { + --ring-color: none; + } + #queue_df th:nth-child(1), + #queue_df td:nth-child(1) { + width: 60px; + text-align: center; + vertical-align: middle; + cursor: default !important; + pointer-events: none; + } + #queue_df th:nth-child(2), + #queue_df td:nth-child(2) { + width: auto; + text-align: center; + vertical-align: middle; + white-space: normal; + } + #queue_df td:nth-child(2) { + cursor: default !important; + } + #queue_df th:nth-child(3), + #queue_df td:nth-child(3) { + width: 60px; + text-align: center; + vertical-align: middle; + cursor: default !important; + pointer-events: none; + } + #queue_df th:nth-child(4), + #queue_df td:nth-child(4) { + width: 60px; + text-align: center; + white-space: nowrap; + cursor: default !important; + pointer-events: none; + } + #queue_df th:nth-child(5), #queue_df td:nth-child(7), + #queue_df th:nth-child(6), #queue_df td:nth-child(8) { + width: 60px; + text-align: center; + vertical-align: middle; + } + #queue_df td:nth-child(5) img, + #queue_df td:nth-child(6) img { + max-width: 50px; + max-height: 50px; + object-fit: contain; + display: block; + margin: auto; + cursor: pointer; + } + #queue_df th:nth-child(7), #queue_df td:nth-child(9), + #queue_df th:nth-child(8), #queue_df td:nth-child(10), + #queue_df th:nth-child(9), #queue_df td:nth-child(11) { + width: 20px; + padding: 2px !important; + cursor: pointer; + text-align: center; + font-weight: bold; + vertical-align: middle; + } + #queue_df td:nth-child(5):hover, + #queue_df td:nth-child(6):hover, + #queue_df td:nth-child(7):hover, + #queue_df td:nth-child(8):hover, + #queue_df td:nth-child(9):hover { + background-color: #e0e0e0; + } #image-modal-container { position: fixed; top: 0; From 986717e5be4100e5fffa74731dd8c9988ca4eaf4 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Wed, 2 Apr 2025 17:34:49 +0200 Subject: [PATCH 32/69] Added Rife Temporal upsampling and Lanczos spatial upsampling --- gradio_server.py | 181 ++++++++++++++++++++++++++++++---------- rife/IFNet_HDv3.py | 133 ++++++++++++++++++++++++++++++ rife/RIFE_HDv3.py | 84 +++++++++++++++++++ rife/inference.py | 119 +++++++++++++++++++++++++++ rife/ssim.py | 200 +++++++++++++++++++++++++++++++++++++++++++++ wan/image2video.py | 22 +---- wan/utils/utils.py | 7 ++ 7 files changed, 682 insertions(+), 64 deletions(-) create mode 100644 rife/IFNet_HDv3.py create mode 100644 rife/RIFE_HDv3.py create mode 100644 rife/inference.py create mode 100644 rife/ssim.py diff --git a/gradio_server.py b/gradio_server.py index d2dde1b..668bbcf 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -97,6 +97,8 @@ def process_prompt_and_add_tasks( image_to_end, video_to_continue, max_frames, + temporal_upsampling, + spatial_upsampling, RIFLEx_setting, slg_switch, slg_layers, @@ -230,6 +232,8 @@ def process_prompt_and_add_tasks( "image_to_end" : image_end, "video_to_continue" : video_to_continue , "max_frames" : max_frames, + "temporal_upsampling" : temporal_upsampling, + "spatial_upsampling" : spatial_upsampling, "RIFLEx_setting" : RIFLEx_setting, "slg_switch" : slg_switch, "slg_layers" : slg_layers, @@ -852,48 +856,63 @@ model_filename = "" # compile = "transformer" def preprocess_loras(sd): + if wan_model == None: + return sd + model_filename = wan_model._model_file_name + first = next(iter(sd), None) if first == None: return sd - if not first.startswith("lora_unet_"): - return sd - new_sd = {} - print("Converting Lora Safetensors format to Lora Diffusers format") - alphas = {} - repl_list = ["cross_attn", "self_attn", "ffn"] - src_list = ["_" + k + "_" for k in repl_list] - tgt_list = ["." + k + "." for k in repl_list] + + if first.startswith("lora_unet_"): + new_sd = {} + print("Converting Lora Safetensors format to Lora Diffusers format") + alphas = {} + repl_list = ["cross_attn", "self_attn", "ffn"] + src_list = ["_" + k + "_" for k in repl_list] + tgt_list = ["." + k + "." for k in repl_list] - for k,v in sd.items(): - k = k.replace("lora_unet_blocks_","diffusion_model.blocks.") + for k,v in sd.items(): + k = k.replace("lora_unet_blocks_","diffusion_model.blocks.") - for s,t in zip(src_list, tgt_list): - k = k.replace(s,t) + for s,t in zip(src_list, tgt_list): + k = k.replace(s,t) - k = k.replace("lora_up","lora_B") - k = k.replace("lora_down","lora_A") + k = k.replace("lora_up","lora_B") + k = k.replace("lora_down","lora_A") - if "alpha" in k: - alphas[k] = v - else: + if "alpha" in k: + alphas[k] = v + else: + new_sd[k] = v + + new_alphas = {} + for k,v in new_sd.items(): + if "lora_B" in k: + dim = v.shape[1] + elif "lora_A" in k: + dim = v.shape[0] + else: + continue + alpha_key = k[:-len("lora_X.weight")] +"alpha" + if alpha_key in alphas: + scale = alphas[alpha_key] / dim + new_alphas[alpha_key] = scale + else: + print(f"Lora alpha'{alpha_key}' is missing") + new_sd.update(new_alphas) + sd = new_sd + + if "text2video" in model_filename: + new_sd = {} + # convert loras for i2v to t2v + for k,v in sd.items(): + if any(layer in k for layer in ["cross_attn.k_img", "cross_attn.v_img"]): + continue new_sd[k] = v + sd = new_sd - new_alphas = {} - for k,v in new_sd.items(): - if "lora_B" in k: - dim = v.shape[1] - elif "lora_A" in k: - dim = v.shape[0] - else: - continue - alpha_key = k[:-len("lora_X.weight")] +"alpha" - if alpha_key in alphas: - scale = alphas[alpha_key] / dim - new_alphas[alpha_key] = scale - else: - print(f"Lora alpha'{alpha_key}' is missing") - new_sd.update(new_alphas) - return new_sd + return sd def download_models(transformer_filename, text_encoder_filename): @@ -905,7 +924,7 @@ def download_models(transformer_filename, text_encoder_filename): from huggingface_hub import hf_hub_download, snapshot_download repoId = "DeepBeepMeep/Wan2.1" sourceFolderList = ["xlm-roberta-large", "", ] - fileList = [ [], ["Wan2.1_VAE_bf16.safetensors", "models_clip_open-clip-xlm-roberta-large-vit-huge-14-bf16.safetensors" ] + computeList(text_encoder_filename) + computeList(transformer_filename) ] + fileList = [ [], ["Wan2.1_VAE_bf16.safetensors", "models_clip_open-clip-xlm-roberta-large-vit-huge-14-bf16.safetensors", "flownet.pkl" ] + computeList(text_encoder_filename) + computeList(transformer_filename) ] targetRoot = "ckpts/" for sourceFolder, files in zip(sourceFolderList,fileList ): if len(files)==0: @@ -1094,6 +1113,7 @@ def load_models(i2v): wan_model, pipe = load_i2v_model(model_filename, "720P" if res720P else "480P") else: wan_model, pipe = load_t2v_model(model_filename, "") + wan_model._model_file_name = model_filename kwargs = { "extraModelsToQuantize": None} if profile == 2 or profile == 4: kwargs["budgets"] = { "transformer" : 100 if preload == 0 else preload, "text_encoder" : 100, "*" : 1000 } @@ -1441,6 +1461,8 @@ def generate_video( image_to_end, video_to_continue, max_frames, + temporal_upsampling, + spatial_upsampling, RIFLEx_setting, slg_switch, slg_layers, @@ -1693,6 +1715,7 @@ def generate_video( cfg_star_switch = cfg_star_switch, cfg_zero_step = cfg_zero_step, ) + # samples = torch.empty( (1,2)) #for testing except Exception as e: if temp_filename!= None and os.path.isfile(temp_filename): os.remove(temp_filename) @@ -1717,8 +1740,6 @@ def generate_video( VRAM_crash = True break - _ , exc_value, exc_traceback = sys.exc_info() - state["prompt"] = "" if VRAM_crash: 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." @@ -1759,17 +1780,61 @@ def generate_video( file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:50]).strip()}.mp4" else: file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:100]).strip()}.mp4" - video_path = os.path.join(save_path, file_name) + video_path = os.path.join(save_path, file_name) + # if False: # for testing + # torch.save(sample, "ouput.pt") + # else: + # sample =torch.load("ouput.pt") + exp = 0 + fps = 16 + + if len(temporal_upsampling) > 0 or len(spatial_upsampling) > 0: + progress_args = [0, status + " - Upsampling"] + progress(*progress_args ) + gen["progress_args"] = progress_args + + if temporal_upsampling == "rife2": + exp = 1 + elif temporal_upsampling == "rife4": + exp = 2 + + if exp > 0: + from rife.inference import temporal_interpolation + sample = temporal_interpolation( os.path.join("ckpts", "flownet.pkl"), sample, exp, device="cuda") + fps = fps * 2**exp + + if len(spatial_upsampling) > 0: + from wan.utils.utils import resize_lanczos + if spatial_upsampling == "lanczos1.5": + scale = 1.5 + else: + scale = 2 + sample = (sample + 1) / 2 + h, w = sample.shape[-2:] + h *= scale + w *= scale + new_frames =[] + for i in range( sample.shape[1] ): + frame = sample[:, i] + frame = resize_lanczos(frame, h, w) + frame = frame.unsqueeze(1) + new_frames.append(frame) + sample = torch.cat(new_frames, dim=1) + new_frames = None + sample = sample * 2 - 1 + + cache_video( tensor=sample[None], save_file=video_path, - fps=16, + fps=fps, nrow=1, 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, - loras_mult_choices, tea_cache , tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, slg_start, slg_end, cfg_star_switch, cfg_zero_step) + 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") if metadata_choice == "json": @@ -2231,7 +2296,7 @@ def switch_advanced(state, new_advanced, 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, - loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc, cfg_star_switch, cfg_zero_step): + 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"] activated_loras = [Path( loras[int(no)]).parts[-1] for no in loras_choices ] @@ -2251,6 +2316,8 @@ def get_settings_dict(state, i2v, prompt, image_prompt_type, video_length, resol "loras_multipliers": loras_mult_choices, "tea_cache": tea_cache_setting, "tea_cache_start_step_perc": tea_cache_start_step_perc, + "temporal_upsampling" : temporal_upsampling, + "spatial_upsampling" : spatial_upsampling, "RIFLEx_setting": RIFLEx_setting, "slg_switch": slg_switch, "slg_layers": slg_layers, @@ -2269,14 +2336,14 @@ def get_settings_dict(state, i2v, prompt, image_prompt_type, video_length, resol 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, - loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc, cfg_star_switch, cfg_zero_step): + 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, - loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, slg_start_perc, slg_end_perc, cfg_star_switch, cfg_zero_step) + 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) @@ -2538,6 +2605,32 @@ 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("Upsampling") + with gr.Row(): + temporal_upsampling_choice = gr.Dropdown( + choices=[ + ("Disabled", ""), + ("Rife x2 (32 frames/s)", "rife2"), + ("Rife x4 (64 frames/s)", "rife4"), + ], + value=ui_defaults.get("temporal_upsampling", ""), + visible=True, + scale = 1, + label="Temporal Upsampling" + ) + spatial_upsampling_choice = gr.Dropdown( + choices=[ + ("Disabled", ""), + ("Lanczos x1.5", "lanczos1.5"), + ("Lanczos x2.0", "lanczos2"), + ], + value=ui_defaults.get("spatial_upsampling", ""), + visible=True, + scale = 1, + label="Spatial Upsampling" + ) + gr.Markdown("With Riflex you can generate videos longer than 5s which is the default duration of videos used to train the model") RIFLEx_setting = gr.Dropdown( choices=[ @@ -2699,7 +2792,7 @@ def generate_video_tab(image2video=False): ) 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, - loras_choices, loras_mult_choices, tea_cache_setting, tea_cache_start_step_perc, RIFLEx_setting, slg_switch, slg_layers, + 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]) confirm_save_lset_btn.click(fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , outputs= [prompt]).then( @@ -2758,6 +2851,8 @@ def generate_video_tab(image2video=False): image_to_end, video_to_continue, max_frames, + temporal_upsampling_choice, + spatial_upsampling_choice, RIFLEx_setting, slg_switch, slg_layers, diff --git a/rife/IFNet_HDv3.py b/rife/IFNet_HDv3.py new file mode 100644 index 0000000..53e512b --- /dev/null +++ b/rife/IFNet_HDv3.py @@ -0,0 +1,133 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +# from ..model.warplayer import warp + +# device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +backwarp_tenGrid = {} + +def warp(tenInput, tenFlow, device): + k = (str(tenFlow.device), str(tenFlow.size())) + if k not in backwarp_tenGrid: + tenHorizontal = torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device).view( + 1, 1, 1, tenFlow.shape[3]).expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1) + tenVertical = torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device).view( + 1, 1, tenFlow.shape[2], 1).expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3]) + backwarp_tenGrid[k] = torch.cat( + [tenHorizontal, tenVertical], 1).to(device) + + tenFlow = torch.cat([tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), + tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0)], 1) + + g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1) + return torch.nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True) + +def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, + padding=padding, dilation=dilation, bias=True), + nn.PReLU(out_planes) + ) + +def conv_bn(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, + padding=padding, dilation=dilation, bias=False), + nn.BatchNorm2d(out_planes), + nn.PReLU(out_planes) + ) + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=64): + super(IFBlock, self).__init__() + self.conv0 = nn.Sequential( + conv(in_planes, c//2, 3, 2, 1), + conv(c//2, c, 3, 2, 1), + ) + self.convblock0 = nn.Sequential( + conv(c, c), + conv(c, c) + ) + self.convblock1 = nn.Sequential( + conv(c, c), + conv(c, c) + ) + self.convblock2 = nn.Sequential( + conv(c, c), + conv(c, c) + ) + self.convblock3 = nn.Sequential( + conv(c, c), + conv(c, c) + ) + self.conv1 = nn.Sequential( + nn.ConvTranspose2d(c, c//2, 4, 2, 1), + nn.PReLU(c//2), + nn.ConvTranspose2d(c//2, 4, 4, 2, 1), + ) + self.conv2 = nn.Sequential( + nn.ConvTranspose2d(c, c//2, 4, 2, 1), + nn.PReLU(c//2), + nn.ConvTranspose2d(c//2, 1, 4, 2, 1), + ) + + def forward(self, x, flow, scale=1): + x = F.interpolate(x, scale_factor= 1. / scale, mode="bilinear", align_corners=False, recompute_scale_factor=False) + flow = F.interpolate(flow, scale_factor= 1. / scale, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 1. / scale + feat = self.conv0(torch.cat((x, flow), 1)) + feat = self.convblock0(feat) + feat + feat = self.convblock1(feat) + feat + feat = self.convblock2(feat) + feat + feat = self.convblock3(feat) + feat + flow = self.conv1(feat) + mask = self.conv2(feat) + flow = F.interpolate(flow, scale_factor=scale, mode="bilinear", align_corners=False, recompute_scale_factor=False) * scale + mask = F.interpolate(mask, scale_factor=scale, mode="bilinear", align_corners=False, recompute_scale_factor=False) + return flow, mask + +class IFNet(nn.Module): + def __init__(self): + super(IFNet, self).__init__() + self.block0 = IFBlock(7+4, c=90) + self.block1 = IFBlock(7+4, c=90) + self.block2 = IFBlock(7+4, c=90) + self.block_tea = IFBlock(10+4, c=90) + # self.contextnet = Contextnet() + # self.unet = Unet() + + def forward(self, x, scale_list=[4, 2, 1], training=False): + if training == False: + channel = x.shape[1] // 2 + img0 = x[:, :channel] + img1 = x[:, channel:] + flow_list = [] + merged = [] + mask_list = [] + warped_img0 = img0 + warped_img1 = img1 + flow = (x[:, :4]).detach() * 0 + mask = (x[:, :1]).detach() * 0 + loss_cons = 0 + block = [self.block0, self.block1, self.block2] + for i in range(3): + f0, m0 = block[i](torch.cat((warped_img0[:, :3], warped_img1[:, :3], mask), 1), flow, scale=scale_list[i]) + f1, m1 = block[i](torch.cat((warped_img1[:, :3], warped_img0[:, :3], -mask), 1), torch.cat((flow[:, 2:4], flow[:, :2]), 1), scale=scale_list[i]) + flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 + mask = mask + (m0 + (-m1)) / 2 + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2], device= flow.device) + warped_img1 = warp(img1, flow[:, 2:4], device= flow.device) + merged.append((warped_img0, warped_img1)) + ''' + c0 = self.contextnet(img0, flow[:, :2]) + c1 = self.contextnet(img1, flow[:, 2:4]) + tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) + res = tmp[:, 1:4] * 2 - 1 + ''' + for i in range(3): + mask_list[i] = torch.sigmoid(mask_list[i]) + merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) + # merged[i] = torch.clamp(merged[i] + res, 0, 1) + return flow_list, mask_list[2], merged diff --git a/rife/RIFE_HDv3.py b/rife/RIFE_HDv3.py new file mode 100644 index 0000000..75c672d --- /dev/null +++ b/rife/RIFE_HDv3.py @@ -0,0 +1,84 @@ +import torch +import torch.nn as nn +import numpy as np +from torch.optim import AdamW +import torch.optim as optim +import itertools +from torch.nn.parallel import DistributedDataParallel as DDP +from .IFNet_HDv3 import * +import torch.nn.functional as F +# from ..model.loss import * + +# device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +class Model: + def __init__(self, local_rank=-1): + self.flownet = IFNet() + # self.device() + # self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-4) + # self.epe = EPE() + # self.vgg = VGGPerceptualLoss().to(device) + # self.sobel = SOBEL() + if local_rank != -1: + self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank) + + def train(self): + self.flownet.train() + + def eval(self): + self.flownet.eval() + + def to(self, device): + self.flownet.to(device) + + def load_model(self, path, rank=0, device = "cuda"): + self.device = device + def convert(param): + if rank == -1: + return { + k.replace("module.", ""): v + for k, v in param.items() + if "module." in k + } + else: + return param + self.flownet.load_state_dict(convert(torch.load(path, map_location=device))) + + def save_model(self, path, rank=0): + if rank == 0: + torch.save(self.flownet.state_dict(),'{}/flownet.pkl'.format(path)) + + def inference(self, img0, img1, scale=1.0): + imgs = torch.cat((img0, img1), 1) + scale_list = [4/scale, 2/scale, 1/scale] + flow, mask, merged = self.flownet(imgs, scale_list) + return merged[2] + + def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): + for param_group in self.optimG.param_groups: + param_group['lr'] = learning_rate + img0 = imgs[:, :3] + img1 = imgs[:, 3:] + if training: + self.train() + else: + self.eval() + scale = [4, 2, 1] + flow, mask, merged = self.flownet(torch.cat((imgs, gt), 1), scale=scale, training=training) + loss_l1 = (merged[2] - gt).abs().mean() + loss_smooth = self.sobel(flow[2], flow[2]*0).mean() + # loss_vgg = self.vgg(merged[2], gt) + if training: + self.optimG.zero_grad() + loss_G = loss_cons + loss_smooth * 0.1 + loss_G.backward() + self.optimG.step() + else: + flow_teacher = flow[2] + return merged[2], { + 'mask': mask, + 'flow': flow[2][:, :2], + 'loss_l1': loss_l1, + 'loss_cons': loss_cons, + 'loss_smooth': loss_smooth, + } diff --git a/rife/inference.py b/rife/inference.py new file mode 100644 index 0000000..24a2bdd --- /dev/null +++ b/rife/inference.py @@ -0,0 +1,119 @@ +import os +import torch +from torch.nn import functional as F +# from .model.pytorch_msssim import ssim_matlab +from .ssim import ssim_matlab + +from .RIFE_HDv3 import Model + +def get_frame(frames, frame_no): + if frame_no >= frames.shape[1]: + return None + frame = (frames[:, frame_no] + 1) /2 + frame = frame.clip(0., 1.) + return frame + +def add_frame(frames, frame, h, w): + frame = (frame * 2) - 1 + frame = frame.clip(-1., 1.) + frame = frame.squeeze(0) + frame = frame[:, :h, :w] + frame = frame.unsqueeze(1) + frames.append(frame.cpu()) + +def process_frames(model, device, frames, exp): + pos = 0 + output_frames = [] + + lastframe = get_frame(frames, 0) + _, h, w = lastframe.shape + scale = 1 + fp16 = False + + def make_inference(I0, I1, n): + middle = model.inference(I0, I1, scale) + if n == 1: + return [middle] + first_half = make_inference(I0, middle, n=n//2) + second_half = make_inference(middle, I1, n=n//2) + if n%2: + return [*first_half, middle, *second_half] + else: + return [*first_half, *second_half] + + tmp = max(32, int(32 / scale)) + ph = ((h - 1) // tmp + 1) * tmp + pw = ((w - 1) // tmp + 1) * tmp + padding = (0, pw - w, 0, ph - h) + + def pad_image(img): + if(fp16): + return F.pad(img, padding).half() + else: + return F.pad(img, padding) + + I1 = lastframe.to(device, non_blocking=True).unsqueeze(0) + I1 = pad_image(I1) + temp = None # save lastframe when processing static frame + + while True: + if temp is not None: + frame = temp + temp = None + else: + pos += 1 + frame = get_frame(frames, pos) + if frame is None: + break + I0 = I1 + I1 = frame.to(device, non_blocking=True).unsqueeze(0) + I1 = pad_image(I1) + I0_small = F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False) + I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False) + ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) + + break_flag = False + if ssim > 0.996: + pos += 1 + frame = get_frame(frames, pos) + if frame is None: + break_flag = True + frame = lastframe + else: + temp = frame + I1 = frame.to(device, non_blocking=True).unsqueeze(0) + I1 = pad_image(I1) + I1 = model.inference(I0, I1, scale) + I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False) + ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) + frame = I1[0] + + if ssim < 0.2: + output = [] + for _ in range((2 ** exp) - 1): + output.append(I0) + else: + output = make_inference(I0, I1, 2**exp-1) if exp else [] + + add_frame(output_frames, lastframe, h, w) + for mid in output: + add_frame(output_frames, mid, h, w) + lastframe = frame + if break_flag: + break + + add_frame(output_frames, lastframe, h, w) + return torch.cat( output_frames, dim=1) + +def temporal_interpolation(model_path, frames, exp, device ="cuda"): + + model = Model() + model.load_model(model_path, -1, device=device) + + model.eval() + model.to(device=device) + + with torch.no_grad(): + output = process_frames(model, device, frames, exp) + + return output diff --git a/rife/ssim.py b/rife/ssim.py new file mode 100644 index 0000000..a4d3032 --- /dev/null +++ b/rife/ssim.py @@ -0,0 +1,200 @@ +import torch +import torch.nn.functional as F +from math import exp +import numpy as np + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +def gaussian(window_size, sigma): + gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)]) + return gauss/gauss.sum() + + +def create_window(window_size, channel=1): + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(device) + window = _2D_window.expand(channel, 1, window_size, window_size).contiguous() + return window + +def create_window_3d(window_size, channel=1): + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()) + _3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t()) + window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(device) + return window + + +def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None): + # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh). + if val_range is None: + if torch.max(img1) > 128: + max_val = 255 + else: + max_val = 1 + + if torch.min(img1) < -0.5: + min_val = -1 + else: + min_val = 0 + L = max_val - min_val + else: + L = val_range + + padd = 0 + (_, channel, height, width) = img1.size() + if window is None: + real_size = min(window_size, height, width) + window = create_window(real_size, channel=channel).to(img1.device) + + # mu1 = F.conv2d(img1, window, padding=padd, groups=channel) + # mu2 = F.conv2d(img2, window, padding=padd, groups=channel) + mu1 = F.conv2d(F.pad(img1, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel) + mu2 = F.conv2d(F.pad(img2, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = F.conv2d(F.pad(img1 * img1, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(F.pad(img2 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu2_sq + sigma12 = F.conv2d(F.pad(img1 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_mu2 + + C1 = (0.01 * L) ** 2 + C2 = (0.03 * L) ** 2 + + v1 = 2.0 * sigma12 + C2 + v2 = sigma1_sq + sigma2_sq + C2 + cs = torch.mean(v1 / v2) # contrast sensitivity + + ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2) + + if size_average: + ret = ssim_map.mean() + else: + ret = ssim_map.mean(1).mean(1).mean(1) + + if full: + return ret, cs + return ret + + +def ssim_matlab(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None): + # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh). + if val_range is None: + if torch.max(img1) > 128: + max_val = 255 + else: + max_val = 1 + + if torch.min(img1) < -0.5: + min_val = -1 + else: + min_val = 0 + L = max_val - min_val + else: + L = val_range + + padd = 0 + (_, _, height, width) = img1.size() + if window is None: + real_size = min(window_size, height, width) + window = create_window_3d(real_size, channel=1).to(img1.device) + # Channel is set to 1 since we consider color images as volumetric images + + img1 = img1.unsqueeze(1) + img2 = img2.unsqueeze(1) + + mu1 = F.conv3d(F.pad(img1, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1) + mu2 = F.conv3d(F.pad(img2, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = F.conv3d(F.pad(img1 * img1, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_sq + sigma2_sq = F.conv3d(F.pad(img2 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu2_sq + sigma12 = F.conv3d(F.pad(img1 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_mu2 + + C1 = (0.01 * L) ** 2 + C2 = (0.03 * L) ** 2 + + v1 = 2.0 * sigma12 + C2 + v2 = sigma1_sq + sigma2_sq + C2 + cs = torch.mean(v1 / v2) # contrast sensitivity + + ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2) + + if size_average: + ret = ssim_map.mean() + else: + ret = ssim_map.mean(1).mean(1).mean(1) + + if full: + return ret, cs + return ret + + +def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False): + device = img1.device + weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(device) + levels = weights.size()[0] + mssim = [] + mcs = [] + for _ in range(levels): + sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range) + mssim.append(sim) + mcs.append(cs) + + img1 = F.avg_pool2d(img1, (2, 2)) + img2 = F.avg_pool2d(img2, (2, 2)) + + mssim = torch.stack(mssim) + mcs = torch.stack(mcs) + + # Normalize (to avoid NaNs during training unstable models, not compliant with original definition) + if normalize: + mssim = (mssim + 1) / 2 + mcs = (mcs + 1) / 2 + + pow1 = mcs ** weights + pow2 = mssim ** weights + # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/ + output = torch.prod(pow1[:-1] * pow2[-1]) + return output + + +# Classes to re-use window +class SSIM(torch.nn.Module): + def __init__(self, window_size=11, size_average=True, val_range=None): + super(SSIM, self).__init__() + self.window_size = window_size + self.size_average = size_average + self.val_range = val_range + + # Assume 3 channel for SSIM + self.channel = 3 + self.window = create_window(window_size, channel=self.channel) + + def forward(self, img1, img2): + (_, channel, _, _) = img1.size() + + if channel == self.channel and self.window.dtype == img1.dtype: + window = self.window + else: + window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype) + self.window = window + self.channel = channel + + _ssim = ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average) + dssim = (1 - _ssim) / 2 + return dssim + +class MSSSIM(torch.nn.Module): + def __init__(self, window_size=11, size_average=True, channel=3): + super(MSSSIM, self).__init__() + self.window_size = window_size + self.size_average = size_average + self.channel = channel + + def forward(self, img1, img2): + return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average) diff --git a/wan/image2video.py b/wan/image2video.py index a71e9d4..ed08d44 100644 --- a/wan/image2video.py +++ b/wan/image2video.py @@ -25,8 +25,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 PIL import Image +from wan.utils.utils import resize_lanczos def optimized_scale(positive_flat, negative_flat): @@ -41,10 +40,6 @@ def optimized_scale(positive_flat, negative_flat): return st_star -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) class WanI2V: @@ -285,21 +280,6 @@ class WanI2V: self.clip.model.cpu() from mmgp import offload - - - # img_interpolated.save('aaa.png') - - # img_interpolated = torch.from_numpy(np.array(img_interpolated).astype(np.float32) / 255.0).movedim(-1, 0) - - # img_interpolated = torch.nn.functional.interpolate(img[None].cpu(), size=(h, w), mode='lanczos') - # img_interpolated = img_interpolated.squeeze(0).transpose(0,2).transpose(1,0) - # img_interpolated = img_interpolated.clamp(-1, 1) - # img_interpolated = (img_interpolated + 1)/2 - # img_interpolated = (img_interpolated*255).type(torch.uint8) - # img_interpolated = img_interpolated.cpu().numpy() - # xxx = Image.fromarray(img_interpolated, 'RGB') - # xxx.save('my.png') - offload.last_offload_obj.unload_all() if any_end_frame: mean2 = 0 diff --git a/wan/utils/utils.py b/wan/utils/utils.py index d725999..e19c298 100644 --- a/wan/utils/utils.py +++ b/wan/utils/utils.py @@ -7,9 +7,16 @@ import os.path as osp import imageio import torch import torchvision +from PIL import Image +import numpy as np __all__ = ['cache_video', 'cache_image', 'str2bool'] +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 rand_name(length=8, suffix=''): name = binascii.b2a_hex(os.urandom(length)).decode('utf-8') From 8f1e6149a834cff2f25d4495af52738218de5b92 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Thu, 3 Apr 2025 01:31:32 +0200 Subject: [PATCH 33/69] fixed bugs --- gradio_server.py | 21 +++++++++++---------- rife/inference.py | 4 ++-- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/gradio_server.py b/gradio_server.py index d883726..e85f67b 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -1773,7 +1773,6 @@ def generate_video( # yield f"Video generation was aborted. Total Generation Time: {end_time-start_time:.1f}s" else: sample = samples.cpu() - # video = rearrange(sample.cpu().numpy(), "c t h w -> t h w c") time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%Hh%Mm%Ss") if os.name == 'nt': @@ -1782,14 +1781,14 @@ def generate_video( file_name = f"{time_flag}_seed{seed}_{sanitize_file_name(prompt[:100]).strip()}.mp4" video_path = os.path.join(save_path, file_name) # if False: # for testing - # torch.save(sample, "ouput.pt") + # torch.save(sample, "output.pt") # else: - # sample =torch.load("ouput.pt") + # sample =torch.load("output.pt") exp = 0 fps = 16 if len(temporal_upsampling) > 0 or len(spatial_upsampling) > 0: - progress_args = [0, status + " - Upsampling"] + progress_args = [(num_inference_steps , num_inference_steps) , status + " - Upsampling" , num_inference_steps] progress(*progress_args ) gen["progress_args"] = progress_args @@ -1804,7 +1803,7 @@ def generate_video( fps = fps * 2**exp if len(spatial_upsampling) > 0: - from wan.utils.utils import resize_lanczos + from wan.utils.utils import resize_lanczos # need multithreading or to do lanczos with cuda if spatial_upsampling == "lanczos1.5": scale = 1.5 else: @@ -2712,10 +2711,12 @@ def generate_video_tab(image2video=False): queue_df = gr.DataFrame( headers=["Qty","Prompt", "Length","Steps","Start", "End", "", "", ""], 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"), wrap=True, value=[], + line_breaks= True, visible= False, # every=1, elem_id="queue_df" @@ -3266,15 +3267,16 @@ def create_demo(): pointer-events: none; text-align: center; vertical-align: middle; + font-size:11px; } - #queue_df table { + #xqueue_df table { width: 100%; overflow: hidden !important; } - #queue_df::-webkit-scrollbar { + #xqueue_df::-webkit-scrollbar { display: none !important; } - #queue_df { + #xqueue_df { scrollbar-width: none !important; -ms-overflow-style: none !important; } @@ -3292,9 +3294,8 @@ def create_demo(): cursor: default !important; pointer-events: none; } - #queue_df th:nth-child(2), + #xqueue_df th:nth-child(2), #queue_df td:nth-child(2) { - width: auto; text-align: center; vertical-align: middle; white-space: normal; diff --git a/rife/inference.py b/rife/inference.py index 24a2bdd..8f775d3 100644 --- a/rife/inference.py +++ b/rife/inference.py @@ -73,7 +73,7 @@ def process_frames(model, device, frames, exp): ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) break_flag = False - if ssim > 0.996: + if ssim > 0.996 or pos > 100: pos += 1 frame = get_frame(frames, pos) if frame is None: @@ -86,7 +86,7 @@ def process_frames(model, device, frames, exp): I1 = model.inference(I0, I1, scale) I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False) ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) - frame = I1[0] + frame = I1[0][:, :h, :w] if ssim < 0.2: output = [] From cf495a4aaf6b1ab8f3bb71c494effe704017b345 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Sat, 5 Apr 2025 02:06:17 +0200 Subject: [PATCH 34/69] Added Vac Contronet support --- gradio_server.py | 403 +++++++++++++++++++++------------ requirements.txt | 8 +- wan/configs/__init__.py | 14 ++ wan/modules/model.py | 144 ++++++++++-- wan/text2video.py | 194 ++++++++++++++-- wan/utils/utils.py | 49 ++++ wan/utils/vace_preprocessor.py | 298 ++++++++++++++++++++++++ 7 files changed, 929 insertions(+), 181 deletions(-) create mode 100644 wan/utils/vace_preprocessor.py diff --git a/gradio_server.py b/gradio_server.py index e85f67b..bcc0164 100644 --- a/gradio_server.py +++ b/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) + + 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_start, image_end in zip(prompts, image_to_continue, image_to_end) : + 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'{truncated_prompt}' 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'Start' - if end_img_uri: - end_img_md = f'End' + 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'Start' + if end_img_uri: + end_img_md = f'End' label = f"Prompt of Video being Generated" html = "" if start_img_md != "": html += "" - if end_img_md != "": - html += "" + if end_img_md != "": + html += "" html += "
" + prompt + "" + start_img_md + "" + end_img_md + "" + end_img_md + "
" 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) + 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) - 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) + 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) - 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_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") + + 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("Upsampling") + gr.Markdown("Upsampling - postprocessing that may improve fluidity and the size of the video") 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("Wan2GP's Lora Festival ! Press the following button to download i2v Remade 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() + + 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: - 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 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("

Wan 2.1GP v3.4 by DeepBeepMeep (Updates)

") + gr.Markdown("

Wan 2.1GP v4.0 by DeepBeepMeep (Updates)

") gr.Markdown("Welcome to Wan 2.1GP a super fast and low VRAM AI Video Generator !") 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"): diff --git a/requirements.txt b/requirements.txt index 7576271..a4cd3b3 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file +mutagen +decord +onnxruntime-gpu +rembg[gpu]==2.0.65 +# rembg==2.0.65 \ No newline at end of file diff --git a/wan/configs/__init__.py b/wan/configs/__init__.py index c72d2d0..e3f539b 100644 --- a/wan/configs/__init__.py +++ b/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'), +} diff --git a/wan/modules/model.py b/wan/modules/model.py index 2daa00c..5eba92b 100644 --- a/wan/modules/model.py +++ b/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]) + 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 - 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: diff --git a/wan/text2video.py b/wan/text2video.py index cdcbd4f..befe139 100644 --- a/wan/text2video.py +++ b/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 diff --git a/wan/utils/utils.py b/wan/utils/utils.py index e19c298..d4e237d 100644 --- a/wan/utils/utils.py +++ b/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: diff --git a/wan/utils/vace_preprocessor.py b/wan/utils/vace_preprocessor.py new file mode 100644 index 0000000..7c10719 --- /dev/null +++ b/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 From 83f1ec01d72a406e705915d34ecbf61703cc5b82 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Sat, 5 Apr 2025 03:06:41 +0200 Subject: [PATCH 35/69] Updated readme --- README.md | 39 ++++++++++++++++++++++++++++++++------- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 2b96746..175c0cc 100644 --- a/README.md +++ b/README.md @@ -10,15 +10,14 @@ Wan2.1 GP by DeepBeepMeep based on Wan2.1's Alibaba: Open and Advanced Large-Scale Video Generative Models for the GPU Poor

-In this repository, we present **Wan2.1**, a comprehensive and open suite of video foundation models that pushes the boundaries of video generation. **Wan2.1** offers these key features: -- 👍 **SOTA Performance**: **Wan2.1** consistently outperforms existing open-source models and state-of-the-art commercial solutions across multiple benchmarks. -- 👍 **Supports Consumer-grade GPUs**: The T2V-1.3B model requires only 8.19 GB VRAM, making it compatible with almost all consumer-grade GPUs. It can generate a 5-second 480P video on an RTX 4090 in about 4 minutes (without optimization techniques like quantization). Its performance is even comparable to some closed-source models. -- 👍 **Multiple Tasks**: **Wan2.1** excels in Text-to-Video, Image-to-Video, Video Editing, Text-to-Image, and Video-to-Audio, advancing the field of video generation. -- 👍 **Visual Text Generation**: **Wan2.1** is the first video model capable of generating both Chinese and English text, featuring robust text generation that enhances its practical applications. -- 👍 **Powerful Video VAE**: **Wan-VAE** delivers exceptional efficiency and performance, encoding and decoding 1080P videos of any length while preserving temporal information, making it an ideal foundation for video and image generation. + ## 🔥 Latest News!! +* April 4 2025: 👋 Wan 2.1GP v4.0: lots of goodies for you ! + - A new queuing system that lets you stack in a queue as many text2video and imag2video tasks as you want. Each task can rely on complete different generation parameters (different number of frames, steps, loras, ...). + - Temporal upsampling (Rife) and spatial upsampling (Lanczos) for a smoother video (32 fps or 64 fps) and to enlarge you video by x2 or x4. Check these new advanced options. + - Wan Vace Control Net support : with Vace you can inject in the scene people or objects, animate a person, perform inpainting or outpainting, continue a video, ... I have provided an introduction guide below. * Mar 27 2025: 👋 Added support for the new Wan Fun InP models (image2video). The 14B Fun InP has probably better end image support but unfortunately existing loras do not work so well with it. The great novelty is the Fun InP image2 1.3B model : Image 2 Video is now accessible to even lower hardware configuration. It is not as good as the 14B models but very impressive for its size. You can choose any of those models in the Configuration tab. Many thanks to the VideoX-Fun team (https://github.com/aigc-apps/VideoX-Fun) * Mar 26 2025: 👋 Good news ! Official support for RTX 50xx please check the installation instructions below. * Mar 24 2025: 👋 Wan2.1GP v3.2: @@ -224,7 +223,7 @@ python gradio_server.py --attention sdpa Every lora stored in the subfoler 'loras' for t2v and 'loras_i2v' will be automatically loaded. You will be then able to activate / desactive any of them when running the application by selecting them in the area below "Activated Loras" . -If you want to manage in differenta areas Loras for the 1.3B model and the 14B as they are not comptatible, just create the following subfolders: +If you want to manage in different areas Loras for the 1.3B model and the 14B as they are not compatible, just create the following subfolders: - loras/1.3B - loras/14B @@ -271,6 +270,32 @@ In the video, a woman is presented. The woman is in a city and looks at her watc You can define multiple lines of macros. If there is only one macro line, the app will generate a simple user interface to enter the macro variables when getting back to *Normal Mode* (advanced mode turned off) +### VACE ControlNet introduction + +Vace is a ControlNet 1.3B text2video model that allows you on top of a text prompt to provide visual hints to guide the generation. It can do more things than image2video although it is not as good for just starting a video with an image because it only a 1.3B model (in fact 3B) versus 14B and (it is not specialized for start frames). However, with Vace you can inject in the scene people or objects, animate a person, perform inpainting or outpainting, continue a video, ... + +First you need to switch the t2v model to Vace 1.3 in the Configuration Tab. Please note that Vace works well for the moment only with videos up to 5s (81 frames). + +Beside the usual Text Prompt, three new types of visual hints can be provided (and combined !): +- reference Images: use this to inject people or objects in the video. You can select multiple reference Images. The integration of the image is more efficient if the background is replaced by the full white color. You can do that with your preferred background remover or use the built in background remover by checking the box *Remove background* + +- a Video: this can be a video that contains a body pose (an animated wireframe that indicates the positions of limbs of a person), a greyed depth map video, a normal video combined with a masked video (see below),... The Vace model will detect automatically what to do depending on the video content. You can tell WanGP to use only the n first frames of this Video. All the frames beyond and up the number of requested frames will be generated by following the Text prompt and the other visual hints (for instance reference images). If the video contains area of grey color 127, they will be considered as masks and will be filled based on the Text prompt of the reference Images. There + +- a Video Mask +This offers a stronger mechanism to tell Vace which parts should be kept (black) or replaced (white). You can do as well inpainting / outpainting, fill the missing part of a video more efficientlty with just the video hint. + + +Examples: +- Inject people and / objects into a scene describe by a text promtp: Ref. Images + text Prompt +- Animate a character described in a text prompt: Body Pose Video + text Prompt +- Animate a character of your choice : Ref Images + Body Pose Video + text Prompt + + +There are lots of possible combinations. Some of them require to prepare some materials (masks on top of video, full masks, etc...). + +Vace provides on its github (https://github.com/ali-vilab/VACE/tree/main/vace/gradios) annotators / preprocessors Gradio tool that can help you build some of these materials depending on the task you want to achieve. + +There is also a guide that describes the various combination of hints (https://github.com/ali-vilab/VACE/blob/main/UserGuide.md).Good luck ! ### Command line parameters for Gradio Server --i2v : launch the image to video generator\ --t2v : launch the text to video generator (default defined in the configuration)\ From dac6796e87b95d25cfa7fe03bd66a20bff8c4741 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Sat, 5 Apr 2025 03:28:14 +0200 Subject: [PATCH 36/69] improved compatibility with former settings --- gradio_server.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/gradio_server.py b/gradio_server.py index bcc0164..7f4e345 100644 --- a/gradio_server.py +++ b/gradio_server.py @@ -2566,6 +2566,8 @@ def generate_video_tab(image2video=False): remove_background_image_ref = gr.Text(visible = False) else: image_prompt_type= ui_defaults.get("image_prompt_type","I") + if not isinstance(image_prompt_type, str): + 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", @@ -3195,9 +3197,10 @@ def generate_about_tab(): gr.Markdown("

Wan2.1GP - Wan 2.1 model for the GPU Poor by DeepBeepMeep (GitHub)

") gr.Markdown("Original Wan 2.1 Model by Alibaba (GitHub)") gr.Markdown("Many thanks to:") + gr.Markdown("- Alibaba Wan team for the best open source video generator") gr.Markdown("- Cocktail Peanuts : QA and simple installation via Pinokio.computer") + gr.Markdown("- Tophness : created multi tabs and queuing frameworks") gr.Markdown("- AmericanPresidentJimmyCarter : added original support for Skip Layer Guidance") - gr.Markdown("- Tophness : created multi tabs framework") gr.Markdown("- Remade_AI : for creating their awesome Loras collection") From fbf2793b2e5e3419a6bfa82690aac1b9a43c0cff Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Mon, 7 Apr 2025 02:15:13 +0200 Subject: [PATCH 37/69] New model selection logic / removed tabs --- README.md | 27 +- wan/utils/vace_preprocessor.py | 6 +- gradio_server.py => wgp.py | 1729 +++++++++++++++----------------- 3 files changed, 838 insertions(+), 924 deletions(-) rename gradio_server.py => wgp.py (69%) diff --git a/README.md b/README.md index 175c0cc..c0380c0 100644 --- a/README.md +++ b/README.md @@ -174,35 +174,35 @@ pip install -e . To run the text to video generator (in Low VRAM mode): ```bash -python gradio_server.py +python wgp.py.py #or -python gradio_server.py --t2v #launch the default text 2 video model +python wgp.py.py --t2v #launch the default text 2 video model #or -python gradio_server.py --t2v-14B #for the 14B model +python wgp.py.py --t2v-14B #for the 14B model #or -python gradio_server.py --t2v-1-3B #for the 1.3B model +python wgp.py.py --t2v-1-3B #for the 1.3B model ``` To run the image to video generator (in Low VRAM mode): ```bash -python gradio_server.py --i2v +python wgp.py.py --i2v ``` To run the 1.3B Fun InP image to video generator (in Low VRAM mode): ```bash -python gradio_server.py --i2v-1-3B +python wgp.py.py --i2v-1-3B ``` To be able to input multiple images with the image to video generator: ```bash -python gradio_server.py --i2v --multiple-images +python wgp.py.py --i2v --multiple-images ``` Within the application you can configure which video generator will be launched without specifying a command line switch. To run the application while loading entirely the diffusion model in VRAM (slightly faster but requires 24 GB of VRAM for a 8 bits quantized 14B model ) ```bash -python gradio_server.py --profile 3 +python wgp.py.py --profile 3 ``` **Trouble shooting**:\ @@ -215,7 +215,7 @@ Therefore you may have no choice but to fallback to sdpa attention, to do so: or - Launch the application this way: ```bash -python gradio_server.py --attention sdpa +python wgp.py.py --attention sdpa ``` ### Loras support @@ -249,7 +249,7 @@ Each preset, is a file with ".lset" extension stored in the loras directory and Last but not least you can pre activate Loras corresponding and prefill a prompt (comments only or full prompt) by specifying a preset when launching the gradio server: ```bash -python gradio_server.py --lora-preset mylorapreset.lset # where 'mylorapreset.lset' is a preset stored in the 'loras' folder +python wgp.py.py --lora-preset mylorapreset.lset # where 'mylorapreset.lset' is a preset stored in the 'loras' folder ``` You will find prebuilt Loras on https://civitai.com/ or you will be able to build them with tools such as kohya or onetrainer. @@ -274,7 +274,7 @@ You can define multiple lines of macros. If there is only one macro line, the ap Vace is a ControlNet 1.3B text2video model that allows you on top of a text prompt to provide visual hints to guide the generation. It can do more things than image2video although it is not as good for just starting a video with an image because it only a 1.3B model (in fact 3B) versus 14B and (it is not specialized for start frames). However, with Vace you can inject in the scene people or objects, animate a person, perform inpainting or outpainting, continue a video, ... -First you need to switch the t2v model to Vace 1.3 in the Configuration Tab. Please note that Vace works well for the moment only with videos up to 5s (81 frames). +First you need to select the Vace 1.3B model in the Drop Down box at the top. Please note that Vace works well for the moment only with videos up to 5s (81 frames). Beside the usual Text Prompt, three new types of visual hints can be provided (and combined !): - reference Images: use this to inject people or objects in the video. You can select multiple reference Images. The integration of the image is more efficient if the background is replaced by the full white color. You can do that with your preferred background remover or use the built in background remover by checking the box *Remove background* @@ -296,6 +296,8 @@ There are lots of possible combinations. Some of them require to prepare some ma Vace provides on its github (https://github.com/ali-vilab/VACE/tree/main/vace/gradios) annotators / preprocessors Gradio tool that can help you build some of these materials depending on the task you want to achieve. There is also a guide that describes the various combination of hints (https://github.com/ali-vilab/VACE/blob/main/UserGuide.md).Good luck ! + +It seems you will get better results if you turn on "Skip Layer Guidance" with its default configuration ### Command line parameters for Gradio Server --i2v : launch the image to video generator\ --t2v : launch the text to video generator (default defined in the configuration)\ @@ -303,6 +305,7 @@ There is also a guide that describes the various combination of hints (https://g --t2v-1-3B : launch the 1.3B model text to video generator\ --i2v-14B : launch the 14B model image to video generator\ --i2v-1-3B : launch the Fun InP 1.3B model image to video generator\ +--vace : launch the Vace ControlNet 1.3B model image to video generator\ --quantize-transformer bool: (default True) : enable / disable on the fly transformer quantization\ --lora-dir path : Path of directory that contains Loras in diffusers / safetensor format\ --lora-preset preset : name of preset gile (without the extension) to preload @@ -324,8 +327,6 @@ There is also a guide that describes the various combination of hints (https://g --slg : turn on skip layer guidance for improved quality\ --check-loras : filter loras that are incompatible (will take a few seconds while refreshing the lora list or while starting the app)\ --advanced : turn on the advanced mode while launching the app\ ---i2v-settings : path to launch settings for i2v\ ---t2v-settings : path to launch settings for t2v\ --listen : make server accessible on network\ --gpu device : run Wan on device for instance "cuda:1" diff --git a/wan/utils/vace_preprocessor.py b/wan/utils/vace_preprocessor.py index 7c10719..912ae39 100644 --- a/wan/utils/vace_preprocessor.py +++ b/wan/utils/vace_preprocessor.py @@ -183,7 +183,7 @@ class VaceVideoProcessor(object): 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_frames_count = len(frame_timestamps) video_frame_duration = 1 /fps target_frame_duration = 1 / target_fps @@ -197,9 +197,9 @@ class VaceVideoProcessor(object): frame_ids.append(frame_no) cur_time += add_frames_count * video_frame_duration target_time += target_frame_duration - if cur_time > video_duration: + if frame_no >= video_frames_count -1: break - + frame_ids = frame_ids[:video_frames_count] 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 diff --git a/gradio_server.py b/wgp.py similarity index 69% rename from gradio_server.py rename to wgp.py index 7f4e345..c81ab9a 100644 --- a/gradio_server.py +++ b/wgp.py @@ -21,7 +21,9 @@ import torch import gc import traceback import math +import typing import asyncio +import inspect from wan.utils import prompt_parser import base64 import io @@ -39,7 +41,6 @@ current_task_id = None task_id = 0 # progress_tracker = {} # tracker_lock = threading.Lock() -last_model_type = None def format_time(seconds): if seconds < 60: @@ -81,63 +82,41 @@ def pil_to_base64_uri(pil_image, format="png", quality=75): return None -def process_prompt_and_add_tasks( - prompt, - negative_prompt, - resolution, - video_length, - seed, - num_inference_steps, - guidance_scale, - flow_shift, - embedded_guidance_scale, - repeat_generation, - multi_images_gen_type, - tea_cache, - tea_cache_start_step_perc, - loras_choices, - loras_mult_choices, - image_prompt_type, - image_source1, - image_source2, - image_source3, - max_frames, - remove_background_image_ref, - temporal_upsampling, - spatial_upsampling, - RIFLEx_setting, - slg_switch, - slg_layers, - slg_start, - slg_end, - cfg_star_switch, - cfg_zero_step, - state, - image2video -): +def process_prompt_and_add_tasks(state, model_choice): if state.get("validate_success",0) != 1: - gr.Info("Validation failed, not adding tasks.") return state["validate_success"] = 0 + + model_filename = state["model_filename"] + + if model_choice != get_model_type(model_filename): + raise gr.Error("Webform can not be used as the App has been restarted since the form was displayed. Please refresh the page") + + inputs = state.get(get_model_type(model_filename), None) + inputs["state"] = state + if inputs == None: + return + prompt = inputs["prompt"] if len(prompt) ==0: return prompt, errors = prompt_parser.process_template(prompt) if len(errors) > 0: gr.Info("Error processing prompt template: " + errors) return + + inputs["model_filename"] = model_filename prompts = prompt.replace("\r", "").split("\n") prompts = [prompt.strip() for prompt in prompts if len(prompt.strip())>0 and not prompt.startswith("#")] if len(prompts) ==0: return - file_model_needed = model_needed(image2video) + resolution = inputs["resolution"] width, height = resolution.split("x") width, height = int(width), int(height) - if image2video: - - if "480p" in file_model_needed and not "Fun" in file_model_needed and width * height > 848*480: + if test_class_i2v(model_filename): + if "480p" in model_filename and not "Fun" in model_filename 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") return resolution = str(width) + "*" + str(height) @@ -145,133 +124,124 @@ def process_prompt_and_add_tasks( gr.Info(f"Resolution {resolution} not supported by image 2 video") return - if "1.3B" in file_model_needed and width * height > 848*480: + if "1.3B" in model_filename and width * height > 848*480: gr.Info("You must use the 14B model to generate videos with a resolution equivalent to 720P") return - - if not image2video: - if "Vace" in file_model_needed and "1.3B" in file_model_needed : + + + if "Vace" in model_filename: + video_prompt_type = inputs["video_prompt_type"] + image_refs = inputs["image_refs"] + video_guide = inputs["video_guide"] + video_mask = inputs["video_mask"] + if "Vace" in model_filename and "1.3B" in model_filename : 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 video_prompt_type: + image_refs = None + if not "V" in video_prompt_type: + video_guide = None + if not "M" in video_prompt_type: + video_mask = None - 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 ] + if isinstance(image_refs, list): + image_refs = [ convert_image(tup[0]) for tup in image_refs ] 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_refs = resize_and_remove_background(image_refs, width, height, inputs["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: + for single_prompt in prompts: + extra_inputs = { + "prompt" : single_prompt, + "image_refs": image_refs, + "video_guide" : video_guide, + "video_mask" : video_mask , + } + inputs.update(extra_inputs) + add_video_task(**inputs) + elif "image2video" in model_filename or "Fun_InP" in model_filename : + image_prompt_type = inputs["image_prompt_type"] + + image_start = inputs["image_start"] + image_end = inputs["image_end"] + if image_start == None or isinstance(image_start, list) and len(image_start) == 0: return - if image_prompt_type == 0: - image_source2 = None - if isinstance(image_source1, list): - image_source1 = [ convert_image(tup[0]) for tup in image_source1 ] + if not "E" in image_prompt_type: + image_end = None + if isinstance(image_start, list): + image_start = [ convert_image(tup[0]) for tup in image_start ] else: - 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 ] + image_start = [convert_image(image_start)] + if image_end != None: + if isinstance(image_end , list): + image_end = [ convert_image(tup[0]) for tup in image_end ] else: - image_source2 = [convert_image(image_source2) ] - if len(image_source1) != len(image_source2): + image_end = [convert_image(image_end) ] + if len(image_start) != len(image_end): gr.Info("The number of start and end images should be the same ") return - if multi_images_gen_type == 0: + if inputs["multi_images_gen_type"] == 0: new_prompts = [] - new_image_source1 = [] - new_image_source2 = [] - for i in range(len(prompts) * len(image_source1) ): + new_image_start = [] + new_image_end = [] + for i in range(len(prompts) * len(image_start) ): new_prompts.append( prompts[ 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)] ) + new_image_start.append(image_start[i // len(prompts)] ) + if image_end != None: + new_image_end.append(image_end[i // len(prompts)] ) prompts = new_prompts - image_source1 = new_image_source1 - if image_source2 != None: - image_source2 = new_image_source2 + image_start = new_image_start + if image_end != None: + image_end = new_image_end else: - if len(prompts) >= len(image_source1): - if len(prompts) % len(image_source1) !=0: + if len(prompts) >= len(image_start): + if len(prompts) % len(image_start) != 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_source1) - new_image_source1 = [] - new_image_source2 = [] + rep = len(prompts) // len(image_start) + new_image_start = [] + new_image_end = [] for i, _ in enumerate(prompts): - 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 + new_image_start.append(image_start[i//rep] ) + if image_end != None: + new_image_end.append(image_end[i//rep] ) + image_start = new_image_start + if image_end != None: + image_end = new_image_end else: - if len(image_source1) % len(prompts) !=0: + if len(image_start) % 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_source1) // len(prompts) + rep = len(image_start) // len(prompts) new_prompts = [] - for i, _ in enumerate(image_source1): + for i, _ in enumerate(image_start): new_prompts.append( prompts[ i//rep] ) prompts = new_prompts - 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) + if image_start == None: + image_start = [None] * len(prompts) + if image_end == None: + image_end = [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, - "resolution" : resolution, - "video_length" : video_length, - "seed" : seed, - "num_inference_steps" : num_inference_steps, - "guidance_scale" : guidance_scale, - "flow_shift" : flow_shift, - "embedded_guidance_scale" : embedded_guidance_scale, - "repeat_generation" : repeat_generation, - "multi_images_gen_type" : multi_images_gen_type, - "tea_cache" : tea_cache, - "tea_cache_start_step_perc" : tea_cache_start_step_perc, - "loras_choices" : loras_choices, - "loras_mult_choices" : loras_mult_choices, - "image_prompt_type" : image_prompt_type, - "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, - "slg_switch" : slg_switch, - "slg_layers" : slg_layers, - "slg_start" : slg_start, - "slg_end" : slg_end, - "cfg_star_switch" : cfg_star_switch, - "cfg_zero_step" : cfg_zero_step, - "state" : state, - "image2video" : image2video - } - add_video_task(**kwargs) + for single_prompt, start, end in zip(prompts, image_start, image_end) : + extra_inputs = { + "prompt" : single_prompt, + "image_start": start, + "image_end" : end, + } + inputs.update(extra_inputs) + add_video_task(**inputs) + else: + for single_prompt in prompts : + extra_inputs = { + "prompt" : single_prompt, + } + inputs.update(extra_inputs) + add_video_task(**inputs) gen = get_gen_info(state) gen["prompts_max"] = len(prompts) + gen.get("prompts_max",0) @@ -282,29 +252,37 @@ def process_prompt_and_add_tasks( -def add_video_task(**kwargs): +def add_video_task(**inputs): global task_id - state = kwargs["state"] + state = inputs["state"] gen = get_gen_info(state) queue = gen["queue"] task_id += 1 current_task_id = task_id - 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"] + inputs_to_query = ["image_start", "image_end", "image_refs", "video_guide", "video_mask"] + start_image_data = None + end_image_data = None + for name in inputs_to_query: + image= inputs.get(name, None) + if image != None: + image= [image] if not isinstance(image, list) else image + if start_image_data == None: + start_image_data = image + else: + end_image_data = image + break queue.append({ "id": current_task_id, - "image2video": kwargs["image2video"], - "params": kwargs.copy(), - "repeats": kwargs["repeat_generation"], - "length": kwargs["video_length"], - "steps": kwargs["num_inference_steps"], - "prompt": kwargs["prompt"], + "params": inputs.copy(), + "repeats": inputs["repeat_generation"], + "length": inputs["video_length"], + "steps": inputs["num_inference_steps"], + "prompt": inputs["prompt"], "start_image_data": start_image_data, "end_image_data": end_image_data, - "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) + "start_image_data_base64": [pil_to_base64_uri(img, format="jpeg", quality=70) for img in start_image_data] if start_image_data != None else None, + "end_image_data_base64": [pil_to_base64_uri(img, format="jpeg", quality=70) for img in end_image_data] if end_image_data != None else None }) return update_queue_data(queue) @@ -372,6 +350,7 @@ def get_queue_table(queue): 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') + end_img_uri = end_img_uri[0] if end_img_uri !=None else None thumbnail_size = "50px" num_steps = item.get('steps') length = item.get('length') @@ -435,35 +414,6 @@ def create_html_progress_bar(percentage=0.0, text="Idle", is_idle=True): """ return html -# def refresh_progress(): -# global current_task_id, progress_tracker, last_status_string -# task_id_to_check = current_task_id -# is_idle = True -# status_string = "Starting..." -# progress_percent = 0.0 -# html_content = "" - -# with tracker_lock: -# with lock: -# processing_or_queued = any(item['state'] in ["Processing", "Queued"] for item in queue) -# if task_id_to_check is not None: -# progress_data = progress_tracker.get(task_id_to_check) -# if progress_data: -# is_idle = False -# current_step = progress_data.get('current_step', 0) -# total_steps = progress_data.get('total_steps', 0) -# status = progress_data.get('status', "Starting...") -# repeats = progress_data.get("repeats", 1) - -# if total_steps > 0: -# progress_float = min(1.0, max(0.0, float(current_step) / float(total_steps))) -# progress_percent = progress_float * 100 -# status_string = f"{status} [{repeats}] - {progress_percent:.1f}% complete ({current_step}/{total_steps} steps)" -# else: -# progress_percent = 0.0 -# status_string = f"{status} [{repeats}] - Initializing..." -# html_content = create_html_progress_bar(progress_percent, status_string, is_idle) -# return gr.update(value=html_content) def update_generation_status(html_content): if(html_content): @@ -534,19 +484,19 @@ def _parse_args(): help="Lora preset to preload" ) - parser.add_argument( - "--i2v-settings", - type=str, - default="i2v_settings.json", - help="Path to settings file for i2v" - ) + # parser.add_argument( + # "--i2v-settings", + # type=str, + # default="i2v_settings.json", + # help="Path to settings file for i2v" + # ) - parser.add_argument( - "--t2v-settings", - type=str, - default="t2v_settings.json", - help="Path to settings file for t2v" - ) + # parser.add_argument( + # "--t2v-settings", + # type=str, + # default="t2v_settings.json", + # help="Path to settings file for t2v" + # ) # parser.add_argument( # "--lora-preset-i2v", @@ -645,6 +595,12 @@ def _parse_args(): action="store_true", help="text to video mode 1.3B model" ) + + parser.add_argument( + "--vace-1-3B", + action="store_true", + help="Vace ControlNet 1.3B model" + ) parser.add_argument( "--i2v-1-3B", action="store_true", @@ -700,8 +656,9 @@ def _parse_args(): return args -def get_lora_dir(i2v): +def get_lora_dir(model_filename): lora_dir =args.lora_dir + i2v = test_class_i2v(model_filename) if i2v and len(lora_dir)==0: lora_dir =args.lora_dir_i2v if len(lora_dir) > 0: @@ -709,7 +666,7 @@ def get_lora_dir(i2v): root_lora_dir = "loras_i2v" if i2v else "loras" - if "1.3B" in (transformer_filename_i2v if i2v else transformer_filename_t2v) : + if "1.3B" in model_filename : lora_dir_1_3B = os.path.join(root_lora_dir, "1.3B") if os.path.isdir(lora_dir_1_3B ): return lora_dir_1_3B @@ -719,6 +676,7 @@ def get_lora_dir(i2v): return lora_dir_14B return root_lora_dir + attention_modes_installed = get_attention_modes() attention_modes_supported = get_supported_attention_modes() args = _parse_args() @@ -740,14 +698,14 @@ 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", "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", ] +transformer_choices = transformer_choices_t2v + transformer_choices_i2v text_encoder_choices = ["ckpts/models_t5_umt5-xxl-enc-bf16.safetensors", "ckpts/models_t5_umt5-xxl-enc-quanto_int8.safetensors"] - server_config_filename = "gradio_config.json" if not Path(server_config_filename).is_file(): server_config = {"attention_mode" : "auto", - "transformer_filename": transformer_choices_t2v[0], - "transformer_filename_i2v": transformer_choices_i2v[1], + "transformer_type": "t2v", + "transformer_quantization": "int8", "text_encoder_filename" : text_encoder_choices[1], "save_path": os.path.join(os.getcwd(), "gradio_outputs"), "compile" : "", @@ -766,23 +724,68 @@ else: text = reader.read() server_config = json.loads(text) -def get_settings_file_name(i2v): - return args.i2v_settings if i2v else args.t2v_settings -def get_default_settings(filename, i2v): - def get_default_prompt(i2v): +model_types = [ "t2v_1.3B", "vace_1.3B", "fun_inp_1.3B", "t2v", "i2v", "i2v_720p", "fun_inp"] +model_signatures = {"t2v": "text2video_14B", "t2v_1.3B" : "text2video_1.3B", "fun_inp_1.3B" : "Fun_InP_1.3B", "fun_inp" : "Fun_InP_14B", + "i2v" : "image2video_480p", "i2v_720p" : "image2video_720p" , "vace_1.3B" : "Vace_1.3B" } + + +def get_model_type(model_filename): + if "text2video" in model_filename and "14B" in model_filename: + return "t2v" + elif "text2video" in model_filename and "1.3B" in model_filename: + return "t2v_1.3B" + elif "Fun_InP" in model_filename and "1.3B" in model_filename: + return "fun_inp_1.3B" + elif "Fun_InP" in model_filename and "14B" in model_filename: + return "fun_inp" + elif "image2video_480p" in model_filename : + return "i2v" + elif "image2video_720p" in model_filename : + return "i2v_720p" + elif "Vace" in model_filename and "1.3B" in model_filename: + return "vace_1.3B" + elif "Vace" in model_filename and "14B" in model_filename: + return "vace" + else: + raise Exception("Unknown model:" + model_filename) + +def test_class_i2v(model_filename): + return "image2video" in model_filename or "Fun_InP" in model_filename + + +def get_model_filename(model_type, quantization): + signature = model_signatures[model_type] + + choices = [ name for name in transformer_choices if signature in name] + if len(quantization) == 0: + quantization = "bf16" + + if len(choices) <= 1: + return choices[0] + + sub_choices = [ name for name in choices if quantization in name] + if len(sub_choices) > 0: + return sub_choices[0] + else: + return choices[0] + +def get_settings_file_name(model_filename): + return get_model_type(model_filename) + "_settings.json" + +def get_default_settings(filename): + def get_default_prompt(i2v): if i2v: return "Several giant wooly mammoths approach treading through a snowy meadow, their long wooly fur lightly blows in the wind as they walk, snow covered trees and dramatic snow capped mountains in the distance, mid afternoon light with wispy clouds and a sun high in the distance creates a warm glow, the low camera view is stunning capturing the large furry mammal with beautiful photography, depth of field." else: return "A large orange octopus is seen resting on the bottom of the ocean floor, blending in with the sandy and rocky terrain. Its tentacles are spread out around its body, and its eyes are closed. The octopus is unaware of a king crab that is crawling towards it from behind a rock, its claws raised and ready to attack. The crab is brown and spiny, with long legs and antennae. The scene is captured from a wide angle, showing the vastness and depth of the ocean. The water is clear and blue, with rays of sunlight filtering through. The shot is sharp and crisp, with a high dynamic range. The octopus and the crab are in focus, while the background is slightly blurred, creating a depth of field effect." - - defaults_filename = get_settings_file_name(i2v) + i2v = "image2video" in file_name + defaults_filename = get_settings_file_name(filename) if not Path(defaults_filename).is_file(): ui_defaults = { "prompts": get_default_prompt(i2v), "resolution": "832x480", "video_length": 81, - "image_prompt_type" : 0 if i2v else "", "num_inference_steps": 30, "seed": -1, "repeat_generation": 1, @@ -805,6 +808,12 @@ def get_default_settings(filename, i2v): else: with open(defaults_filename, "r", encoding="utf-8") as f: ui_defaults = json.load(f) + prompts = ui_defaults.get("prompts", "") + if len(prompts) > 0: + ui_defaults["prompt"] = prompts + image_prompt_type = ui_defaults.get("image_prompt_type", None) + if image_prompt_type !=None and not isinstance(image_prompt_type, str): + ui_defaults["image_prompt_type"] = "S" if image_prompt_type == 0 else "SE" default_seed = args.seed if default_seed > -1: @@ -817,9 +826,9 @@ def get_default_settings(filename, i2v): ui_defaults["num_inference_steps"] = default_number_steps return ui_defaults -transformer_filename_t2v = server_config["transformer_filename"] -transformer_filename_i2v = server_config.get("transformer_filename_i2v", transformer_choices_i2v[1]) - +transformer_type = server_config.get("transformer_type", "t2v") +transformer_quantization =server_config.get("transformer_quantization", "int8") +transformer_filename = get_model_filename(transformer_type, transformer_quantization) text_encoder_filename = server_config["text_encoder_filename"] attention_mode = server_config["attention_mode"] if len(args.attention)> 0: @@ -839,41 +848,27 @@ if len(args.vae_config) > 0: reload_needed = False default_ui = server_config.get("default_ui", "t2v") save_path = server_config.get("save_path", os.path.join(os.getcwd(), "gradio_outputs")) -use_image2video = default_ui != "t2v" -if args.t2v: - use_image2video = False -if args.i2v: - use_image2video = True -if args.t2v_14B: - use_image2video = False - if not "14B" in transformer_filename_t2v: - transformer_filename_t2v = transformer_choices_t2v[2] - lock_ui_transformer = False +reload_model = server_config.get("reload_model", 2) -if args.i2v_14B: - use_image2video = True - if not "14B" in transformer_filename_i2v: - transformer_filename_i2v = transformer_choices_t2v[3] - lock_ui_transformer = False + +if args.t2v_14B or args.t2v: + transformer_filename = get_model_filename("t2v", transformer_quantization) + +if args.i2v_14B or args.i2v: + transformer_filename = get_model_filename("i2v", transformer_quantization) if args.t2v_1_3B: - transformer_filename_t2v = transformer_choices_t2v[0] - use_image2video = False - lock_ui_transformer = False + transformer_filename = get_model_filename("t2v_1.3B", transformer_quantization) if args.i2v_1_3B: - transformer_filename_i2v = transformer_choices_i2v[4] - use_image2video = True - lock_ui_transformer = False + transformer_filename = get_model_filename("fun_inp_1.3B", transformer_quantization) + +if args.vace_1_3B: + transformer_filename = get_model_filename("vace_1.3B", transformer_quantization) only_allow_edit_in_advanced = False lora_preselected_preset = args.lora_preset -# if args.fast : #or args.fastest -# transformer_filename_t2v = transformer_choices_t2v[2] -# attention_mode="sage2" if "sage2" in attention_modes_supported else "sage" -# default_tea_cache = 0.15 -# lock_ui_attention = True -# lock_ui_transformer = True +lora_preset_model = transformer_filename if args.compile: #args.fastest or compile="transformer" @@ -982,19 +977,19 @@ for file_name in to_remove: except: pass -download_models(transformer_filename_i2v if use_image2video else transformer_filename_t2v, text_encoder_filename) +download_models(transformer_filename, text_encoder_filename) def sanitize_file_name(file_name, rep =""): return file_name.replace("/",rep).replace("\\",rep).replace(":",rep).replace("|",rep).replace("?",rep).replace("<",rep).replace(">",rep).replace("\"",rep) -def extract_preset(image2video, lset_name, loras): +def extract_preset(model_filename, lset_name, loras): loras_choices = [] loras_choices_files = [] loras_mult_choices = "" prompt ="" full_prompt ="" lset_name = sanitize_file_name(lset_name) - lora_dir = get_lora_dir(image2video) + lora_dir = get_lora_dir(model_filename) if not lset_name.endswith(".lset"): lset_name_filename = os.path.join(lora_dir, lset_name + ".lset" ) else: @@ -1028,7 +1023,7 @@ def extract_preset(image2video, lset_name, loras): -def setup_loras(i2v, transformer, lora_dir, lora_preselected_preset, split_linear_modules_map = None): +def setup_loras(model_filename, transformer, lora_dir, lora_preselected_preset, split_linear_modules_map = None): loras =[] loras_names = [] default_loras_choices = [] @@ -1039,7 +1034,7 @@ def setup_loras(i2v, transformer, lora_dir, lora_preselected_preset, split_line from pathlib import Path - lora_dir = get_lora_dir(i2v) + lora_dir = get_lora_dir(model_filename) if lora_dir != None : if not os.path.isdir(lora_dir): raise Exception("--lora-dir should be a path to a directory that contains Loras") @@ -1065,7 +1060,7 @@ def setup_loras(i2v, transformer, lora_dir, lora_preselected_preset, split_line if not os.path.isfile(os.path.join(lora_dir, lora_preselected_preset + ".lset")): raise Exception(f"Unknown preset '{lora_preselected_preset}'") default_lora_preset = lora_preselected_preset - default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, _ , error = extract_preset(i2v, default_lora_preset, loras) + default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, _ , error = extract_preset(model_filename, default_lora_preset, loras) if len(error) > 0: print(error[:200]) return loras, loras_names, loras_presets, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset @@ -1133,14 +1128,14 @@ def load_i2v_model(model_filename, value): raise Exception("Model i2v {value} not supported") return wan_model, pipe -def model_needed(i2v): - return transformer_filename_i2v if i2v else transformer_filename_t2v -def load_models(i2v): - global model_filename - model_filename = model_needed(i2v) + +def load_models(model_filename): + global transformer_filename + + transformer_filename = model_filename download_models(model_filename, text_encoder_filename) - if i2v: + if test_class_i2v(model_filename): res720P = "720p" in model_filename wan_model, pipe = load_i2v_model(model_filename, "720P" if res720P else "480P") else: @@ -1159,11 +1154,15 @@ def load_models(i2v): return wan_model, offloadobj, pipe["transformer"] -wan_model, offloadobj, transformer = load_models(use_image2video) -if check_loras: - setup_loras(use_image2video, transformer, get_lora_dir(use_image2video), "", None) - exit() -del transformer +if reload_model ==3: + wan_model, offloadobj, transformer = None, None, None + reload_needed = True +else: + wan_model, offloadobj, transformer = load_models(transformer_filename) + if check_loras: + setup_loras(model_filename, transformer, get_lora_dir(transformer_filename), "", None) + exit() + del transformer gen_in_progress = False @@ -1182,7 +1181,7 @@ def get_model_name(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 = "Vace ControlNet" model_name += " 14B" if "14B" in model_filename else " 1.3B" elif "image" in model_filename: model_name = "Wan2.1 image2video" @@ -1193,14 +1192,30 @@ def get_model_name(model_filename): return model_name -def generate_header(model_filename, compile, attention_mode): +# def generate_header(model_filename, compile, attention_mode): - header = "

" +# header = "

" - model_name = get_model_name(model_filename) +# model_name = get_model_name(model_filename) - header += model_name - header += " (attention mode: " + (attention_mode if attention_mode!="auto" else "auto/" + get_auto_attention() ) +# header += model_name +# header += " (attention mode: " + (attention_mode if attention_mode!="auto" else "auto/" + get_auto_attention() ) +# if attention_mode not in attention_modes_installed: +# header += " -NOT INSTALLED-" +# elif attention_mode not in attention_modes_supported: +# header += " -NOT SUPPORTED-" + +# if compile: +# header += ", pytorch compilation ON" +# header += ")

" + + +# return header + + +def generate_header(compile, attention_mode): + + header = "
Attention mode: " + (attention_mode if attention_mode!="auto" else "auto/" + get_auto_attention() ) if attention_mode not in attention_modes_installed: header += " -NOT INSTALLED-" elif attention_mode not in attention_modes_supported: @@ -1208,14 +1223,12 @@ def generate_header(model_filename, compile, attention_mode): if compile: header += ", pytorch compilation ON" - header += ")

" - + header += "
" return header def apply_changes( state, - transformer_t2v_choice, - transformer_i2v_choice, + transformer_type_choice, text_encoder_choice, save_path_choice, attention_choice, @@ -1223,7 +1236,7 @@ def apply_changes( state, profile_choice, vae_config_choice, metadata_choice, - default_ui_choice ="t2v", + quantization_choice, boost_choice = 1, clear_file_list = 0, reload_choice = 1 @@ -1235,15 +1248,14 @@ def apply_changes( state, return global offloadobj, wan_model, server_config, loras, loras_names, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset, loras_presets server_config = {"attention_mode" : attention_choice, - "transformer_filename": transformer_choices_t2v[transformer_t2v_choice], - "transformer_filename_i2v": transformer_choices_i2v[transformer_i2v_choice], + "transformer_type": transformer_type_choice, "text_encoder_filename" : text_encoder_choices[text_encoder_choice], "save_path" : save_path_choice, "compile" : compile_choice, "profile" : profile_choice, "vae_config" : vae_config_choice, "metadata_choice": metadata_choice, - "default_ui" : default_ui_choice, + "transformer_quantization" : quantization_choice, "boost" : boost_choice, "clear_file_list" : clear_file_list, "reload_model" : reload_choice, @@ -1255,7 +1267,6 @@ def apply_changes( state, old_server_config = json.loads(text) if lock_ui_transformer: server_config["transformer_filename"] = old_server_config["transformer_filename"] - server_config["transformer_filename_i2v"] = old_server_config["transformer_filename_i2v"] if lock_ui_attention: server_config["attention_mode"] = old_server_config["attention_mode"] if lock_ui_compile: @@ -1270,15 +1281,17 @@ def apply_changes( state, if v != v_old: changes.append(k) - global attention_mode, profile, compile, transformer_filename_t2v, transformer_filename_i2v, text_encoder_filename, vae_config, boost, lora_dir, reload_needed + global attention_mode, profile, compile, transformer_filename, text_encoder_filename, vae_config, boost, lora_dir, reload_needed, reload_model, transformer_quantization, transformer_type attention_mode = server_config["attention_mode"] profile = server_config["profile"] compile = server_config["compile"] - transformer_filename_t2v = server_config["transformer_filename"] - transformer_filename_i2v = server_config["transformer_filename_i2v"] text_encoder_filename = server_config["text_encoder_filename"] vae_config = server_config["vae_config"] boost = server_config["boost"] + reload_model = server_config["reload_model"] + transformer_quantization = server_config["transformer_quantization"] + transformer_filename = get_model_filename(transformer_type, transformer_quantization) + if all(change in ["attention_mode", "vae_config", "default_ui", "boost", "save_path", "metadata_choice", "clear_file_list"] for change in changes ): pass else: @@ -1357,20 +1370,12 @@ def abort_generation(state): else: return "", gr.Button(interactive= True) -def is_gen_location(state): - gen = get_gen_info(state) - gen_location = gen.get("location",None) - if gen_location == None: - return None - return state["image2video"] == gen_location - def refresh_gallery(state, msg): gen = get_gen_info(state) - if is_gen_location(state): - gen["last_msg"] = msg + gen["last_msg"] = msg file_list = gen.get("file_list", None) choice = gen.get("selected",0) in_progress = "in_progress" in gen @@ -1391,6 +1396,7 @@ def refresh_gallery(state, msg): 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') + end_img_uri = end_img_uri[0] if end_img_uri !=None else None thumbnail_size = "100px" if start_img_uri: start_img_md = f'Start' @@ -1489,14 +1495,17 @@ def generate_video( embedded_guidance_scale, repeat_generation, multi_images_gen_type, - tea_cache, + tea_cache_setting, tea_cache_start_step_perc, - loras_choices, - loras_mult_choices, + activated_loras, + loras_multipliers, image_prompt_type, - image_source1, - image_source2, - image_source3, + image_start, + image_end, + video_prompt_type, + image_refs, + video_guide, + video_mask, max_frames, remove_background_image_ref, temporal_upsampling, @@ -1504,39 +1513,41 @@ def generate_video( RIFLEx_setting, slg_switch, slg_layers, - slg_start, - slg_end, + slg_start_perc, + slg_end_perc, cfg_star_switch, cfg_zero_step, state, - image2video + model_filename ): - global wan_model, offloadobj, reload_needed, last_model_type + global wan_model, offloadobj, reload_needed gen = get_gen_info(state) file_list = gen["file_list"] prompt_no = gen["prompt_no"] - - file_model_needed = model_needed(image2video) - # queue = gen.get("queue", []) - # with lock: - # queue_not_empty = len(queue) > 0 - # if(last_model_type != image2video and (queue_not_empty or server_config.get("reload_model",1) == 2) and (file_model_needed != model_filename or reload_needed)): - if file_model_needed != model_filename or reload_needed: - del wan_model + + + # if wan_model == None: + # gr.Info("Unable to generate a Video while a new configuration is being applied.") + # return + + if reload_model !=3 : + while wan_model == None: + time.sleep(1) + + if model_filename != transformer_filename or reload_needed: + wan_model = None if offloadobj is not None: offloadobj.release() - del offloadobj + offloadobj = None gc.collect() - yield f"Loading model {get_model_name(file_model_needed)}..." - wan_model, offloadobj, trans = load_models(image2video) + yield f"Loading model {get_model_name(model_filename)}..." + wan_model, offloadobj, trans = load_models(model_filename) yield f"Model loaded" reload_needed= False - if wan_model == None: - gr.Info("Unable to generate a Video while a new configuration is being applied.") if attention_mode == "auto": attn = get_auto_attention() elif attention_mode in attention_modes_supported: @@ -1544,11 +1555,10 @@ def generate_video( else: 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") - width, height = int(width), int(height) + width, height = resolution.split("x") + width, height = int(width), int(height) + resolution_reformated = str(height) + "*" + str(width) if slg_switch == 0: slg_layers = None @@ -1589,11 +1599,11 @@ def generate_video( except ValueError: return False list_mult_choices_nums = [] - if len(loras_mult_choices) > 0: - loras_mult_choices_list = loras_mult_choices.replace("\r", "").split("\n") + if len(loras_multipliers) > 0: + loras_mult_choices_list = loras_multipliers.replace("\r", "").split("\n") loras_mult_choices_list = [multi for multi in loras_mult_choices_list if len(multi)>0 and not multi.startswith("#")] - loras_mult_choices = " ".join(loras_mult_choices_list) - list_mult_choices_str = loras_mult_choices.split(" ") + loras_multipliers = " ".join(loras_mult_choices_list) + list_mult_choices_str = loras_multipliers.split(" ") for i, mult in enumerate(list_mult_choices_str): mult = mult.strip() if "," in mult: @@ -1609,9 +1619,9 @@ def generate_video( if not is_float(mult): raise gr.Error(f"Lora Multiplier no {i+1} ({mult}) is invalid") list_mult_choices_nums.append(float(mult)) - if len(list_mult_choices_nums ) < len(loras_choices): - list_mult_choices_nums += [1.0] * ( len(loras_choices) - len(list_mult_choices_nums ) ) - loras_selected = [ lora for i, lora in enumerate(loras) if str(i) in loras_choices] + if len(list_mult_choices_nums ) < len(activated_loras): + list_mult_choices_nums += [1.0] * ( len(activated_loras) - len(list_mult_choices_nums ) ) + loras_selected = [ lora for lora in loras if os.path.basename(lora) in activated_loras] pinnedLora = profile !=5 #False # # # offload.load_loras_into_model(trans, loras_selected, list_mult_choices_nums, activate_all_loras=True, preprocess_sd=preprocess_loras, pinnedLora=pinnedLora, split_linear_modules_map = None) errors = trans._loras_errors @@ -1620,43 +1630,42 @@ def generate_video( raise gr.Error("Error while loading Loras: " + ", ".join(error_files)) seed = None if seed == -1 else seed # negative_prompt = "" # not applicable in the inference - + image2video = test_class_i2v(model_filename) enable_RIFLEx = RIFLEx_setting == 0 and video_length > (6* 16) or RIFLEx_setting == 1 # VAE Tiling device_mem_capacity = torch.cuda.get_device_properties(None).total_memory / 1048576 joint_pass = boost ==1 #and profile != 1 and profile != 3 # TeaCache - trans.enable_teacache = tea_cache > 0 + trans.enable_teacache = tea_cache_setting > 0 if trans.enable_teacache: - trans.teacache_multiplier = tea_cache + trans.teacache_multiplier = tea_cache_setting trans.rel_l1_thresh = 0 trans.teacache_start_step = int(tea_cache_start_step_perc*num_inference_steps/100) if image2video: - if '480p' in transformer_filename_i2v: + if '480p' in model_filename: # teacache_thresholds = [0.13, .19, 0.26] trans.coefficients = [-3.02331670e+02, 2.23948934e+02, -5.25463970e+01, 5.87348440e+00, -2.01973289e-01] - elif '720p' in transformer_filename_i2v: + elif '720p' in model_filename: teacache_thresholds = [0.18, 0.2 , 0.3] trans.coefficients = [-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683] else: raise gr.Error("Teacache not supported for this model") else: - if '1.3B' in transformer_filename_t2v: + if '1.3B' in model_filename: # teacache_thresholds= [0.05, 0.07, 0.08] trans.coefficients = [2.39676752e+03, -1.31110545e+03, 2.01331979e+02, -8.29855975e+00, 1.37887774e-01] - elif '14B' in transformer_filename_t2v: + elif '14B' in model_filename: # teacache_thresholds = [0.14, 0.15, 0.2] trans.coefficients = [-5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404] 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], + src_video, src_mask, src_ref_images = wan_model.prepare_source([video_guide], + [video_mask], + [image_refs], video_length, VACE_SIZE_CONFIGS[resolution_reformated], "cpu", trim_video=max_frames) else: @@ -1721,10 +1730,10 @@ def generate_video( if image2video: samples = wan_model.generate( prompt, - image_source1, - image_source2 if image_source2 != None else None, + image_start, + image_end if image_end != None else None, frame_num=(video_length // 4)* 4 + 1, - max_area=MAX_AREA_CONFIGS[resolution], + max_area=MAX_AREA_CONFIGS[resolution_reformated], shift=flow_shift, sampling_steps=num_inference_steps, guide_scale=guidance_scale, @@ -1736,11 +1745,11 @@ def generate_video( VAE_tile_size = VAE_tile_size, joint_pass = joint_pass, slg_layers = slg_layers, - slg_start = slg_start/100, - slg_end = slg_end/100, + slg_start = slg_start_perc/100, + slg_end = slg_end_perc/100, cfg_star_switch = cfg_star_switch, cfg_zero_step = cfg_zero_step, - add_frames_for_end_image = not "Fun" in transformer_filename_i2v, + add_frames_for_end_image = not "Fun_InP" in model_filename, ) else: samples = wan_model.generate( @@ -1761,8 +1770,8 @@ def generate_video( VAE_tile_size = VAE_tile_size, joint_pass = joint_pass, slg_layers = slg_layers, - slg_start = slg_start/100, - slg_end = slg_end/100, + slg_start = slg_start_perc/100, + slg_end = slg_end_perc/100, cfg_star_switch = cfg_star_switch, cfg_zero_step = cfg_zero_step, ) @@ -1863,6 +1872,8 @@ def generate_video( h, w = sample.shape[-2:] h *= scale w *= scale + h = int(h) + w = int(w) new_frames =[] for i in range( sample.shape[1] ): frame = sample[:, i] @@ -1881,9 +1892,10 @@ def generate_video( nrow=1, normalize=True, value_range=(-1, 1)) - - 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) + + inputs = get_function_arguments(generate_video, locals()) + inputs.pop("progress") + configs = prepare_inputs_dict("metadata", inputs) metadata_choice = server_config.get("metadata_choice","metadata") if metadata_choice == "json": @@ -1899,8 +1911,6 @@ def generate_video( file_list.append(video_path) state['update_gallery'] = True seed += 1 - - last_model_type = image2video if temp_filename!= None and os.path.isfile(temp_filename): os.remove(temp_filename) @@ -1913,36 +1923,6 @@ def prepare_generate_video(state): return gr.Button(visible= False), gr.Button(visible= True), gr.Column(visible= True) -def wait_tasks_done(state, progress=gr.Progress()): - - gen = get_gen_info(state) - gen_location = is_gen_location(state) - - last_msg = gen.get("last_msg", "") - if len(last_msg) > 0: - yield last_msg - - if gen_location == None or gen_location: - return gr.Text() - - - while True: - - msg = gen.get("last_msg", "") - if len(msg) > 0 and last_msg != msg: - yield msg - last_msg = msg - progress_args = gen.get("progress_args", None) - if progress_args != None: - progress(*progress_args) - - in_progress= gen.get("in_progress", False) - if not in_progress: - break - time.sleep(0.5) - - - def process_tasks(state, progress=gr.Progress()): gen = get_gen_info(state) queue = gen.get("queue", []) @@ -1950,7 +1930,6 @@ def process_tasks(state, progress=gr.Progress()): if len(queue) == 0: return gen = get_gen_info(state) - gen["location"] = state["image2video"] clear_file_list = server_config.get("clear_file_list", 0) file_list = gen.get("file_list", []) if clear_file_list > 0: @@ -2108,7 +2087,7 @@ def save_lset(state, lset_name, loras_choices, loras_mult_choices, prompt, save_ lset_name_filename = lset_name + ".lset" - full_lset_name_filename = os.path.join(get_lora_dir(state["image2video"]), lset_name_filename) + full_lset_name_filename = os.path.join(get_lora_dir(state["model_filename"]), lset_name_filename) with open(full_lset_name_filename, "w", encoding="utf-8") as writer: writer.write(json.dumps(lset, indent=4)) @@ -2125,7 +2104,7 @@ def save_lset(state, lset_name, loras_choices, loras_mult_choices, prompt, save_ def delete_lset(state, lset_name): loras_presets = state["loras_presets"] - lset_name_filename = os.path.join( get_lora_dir(state["image2video"]), sanitize_file_name(lset_name) + ".lset" ) + lset_name_filename = os.path.join( get_lora_dir(state["model_filename"]), sanitize_file_name(lset_name) + ".lset" ) if len(lset_name) > 0 and lset_name != get_new_preset_msg(True) and lset_name != get_new_preset_msg(False): if not os.path.isfile(lset_name_filename): raise gr.Error(f"Preset '{lset_name}' not found ") @@ -2146,8 +2125,8 @@ def delete_lset(state, lset_name): def refresh_lora_list(state, lset_name, loras_choices): loras_names = state["loras_names"] prev_lora_names_selected = [ loras_names[int(i)] for i in loras_choices] - image2video= state["image2video"] - loras, loras_names, loras_presets, _, _, _, _ = setup_loras(image2video, None, get_lora_dir(image2video), lora_preselected_preset, None) + model_filename= state["model_filename"] + loras, loras_names, loras_presets, _, _, _, _ = setup_loras(model_filename, None, get_lora_dir(model_filename), lora_preselected_preset, None) state["loras"] = loras state["loras_names"] = loras_names state["loras_presets"] = loras_presets @@ -2187,7 +2166,7 @@ def apply_lset(state, wizard_prompt_activated, lset_name, loras_choices, loras_m gr.Info("Please choose a preset in the list or create one") else: loras = state["loras"] - loras_choices, loras_mult_choices, preset_prompt, full_prompt, error = extract_preset(state["image2video"], lset_name, loras) + loras_choices, loras_mult_choices, preset_prompt, full_prompt, error = extract_preset(state["model_filename"], lset_name, loras) if len(error) > 0: gr.Info(error) else: @@ -2344,71 +2323,110 @@ 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, 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): - +def prepare_inputs_dict(target, inputs ): + + state = inputs.pop("state") loras = state["loras"] - activated_loras = [Path( loras[int(no)]).parts[-1] for no in loras_choices ] + if "loras_choices" in inputs: + loras_choices = inputs.pop("loras_choices") + inputs.pop("model_filename", None) + activated_loras = [Path( loras[int(no)]).parts[-1] for no in loras_choices ] + inputs["activated_loras"] = activated_loras - ui_settings = { - "prompts": prompt, - "resolution": resolution, - "video_length": video_length, - "num_inference_steps": num_inference_steps, - "seed": seed, - "repeat_generation": repeat_generation, - "multi_images_gen_type": multi_images_gen_type, - "guidance_scale": guidance_scale, - "flow_shift": flow_shift, - "negative_prompt": negative_prompt, - "activated_loras": activated_loras, - "loras_multipliers": loras_mult_choices, - "tea_cache": tea_cache_setting, - "tea_cache_start_step_perc": tea_cache_start_step_perc, - "temporal_upsampling" : temporal_upsampling, - "spatial_upsampling" : spatial_upsampling, - "RIFLEx_setting": RIFLEx_setting, - "slg_switch": slg_switch, - "slg_layers": slg_layers, - "slg_start_perc": slg_start_perc, - "slg_end_perc": slg_end_perc, - "cfg_star_switch": cfg_star_switch, - "cfg_zero_step": cfg_zero_step - } + if target == "state": + return inputs - if i2v: - 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" + unsaved_params = ["image_start", "image_end", "image_refs", "video_guide", "video_mask"] + for k in unsaved_params: + inputs.pop(k) - return ui_settings + model_filename = state["model_filename"] + inputs["type"] = "Wan2.1GP by DeepBeepMeep - " + get_model_name(model_filename) -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 target == "settings": + return inputs + + if not any(k in model_filename for k in ["image2video", "Fun_InP"]): + inputs.pop("image_prompt_type") - if state.get("validate_success",0) != 1: - return - image2video = state["image2video"] - 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) + if not "Vace" in model_filename: + unsaved_params = ["video_prompt_type", "max_frames", "remove_background_image_ref"] + for k in unsaved_params: + inputs.pop(k) - defaults_filename = get_settings_file_name(image2video) + if target == "metadata": + inputs = {k: v for k,v in inputs.items() if v != None } - with open(defaults_filename, "w", encoding="utf-8") as f: - json.dump(ui_defaults, f, indent=4) + return inputs - gr.Info("New Default Settings saved") +def get_function_arguments(func, locals): + args_names = list(inspect.signature(func).parameters) + kwargs = typing.OrderedDict() + for k in args_names: + kwargs[k] = locals[k] + return kwargs + + +def save_inputs( + target, + prompt, + negative_prompt, + resolution, + video_length, + seed, + num_inference_steps, + guidance_scale, + flow_shift, + embedded_guidance_scale, + repeat_generation, + multi_images_gen_type, + tea_cache_setting, + tea_cache_start_step_perc, + loras_choices, + loras_multipliers, + image_prompt_type, + image_start, + image_end, + video_prompt_type, + image_refs, + video_guide, + video_mask, + max_frames, + remove_background_image_ref, + temporal_upsampling, + spatial_upsampling, + RIFLEx_setting, + slg_switch, + slg_layers, + slg_start_perc, + slg_end_perc, + cfg_star_switch, + cfg_zero_step, + state, +): + + + # if state.get("validate_success",0) != 1: + # return + model_filename = state["model_filename"] + inputs = get_function_arguments(save_inputs, locals()) + inputs.pop("target") + cleaned_inputs = prepare_inputs_dict(target, inputs) + if target == "settings": + defaults_filename = get_settings_file_name(model_filename) + + with open(defaults_filename, "w", encoding="utf-8") as f: + json.dump(cleaned_inputs, f, indent=4) + + gr.Info("New Default Settings saved") + elif target == "state": + state[get_model_type(model_filename)] = cleaned_inputs def download_loras(): from huggingface_hub import snapshot_download yield gr.Row(visible=True), "Please wait while the Loras are being downloaded", *[gr.Column(visible=False)] * 2 - lora_dir = get_lora_dir(True) + lora_dir = get_lora_dir(get_model_filename("i2v"), quantizeTransformer) log_path = os.path.join(lora_dir, "log.txt") if not os.path.isfile(log_path): import shutil @@ -2433,42 +2451,130 @@ 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): +def refresh_image_prompt_type(state, image_prompt_type): if args.multiple_images: - return gr.Gallery(visible = (image_prompt_type_radio == 1) ) + return gr.Gallery(visible = "S" in image_prompt_type ), gr.Gallery(visible = "E" in image_prompt_type ) else: - return gr.Image(visible = (image_prompt_type_radio == 1) ) + return gr.Image(visible = "S" in image_prompt_type ), gr.Image(visible = "E" in image_prompt_type ) -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 refresh_video_prompt_type(state, video_prompt_type): + return gr.Gallery(visible = "I" in video_prompt_type), gr.Video(visible= "V" in video_prompt_type),gr.Video(visible= "M" in video_prompt_type ), gr.Text(visible= "V" in video_prompt_type) , gr.Checkbox(visible= "I" in video_prompt_type) -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())) + +def handle_celll_selection(state, evt: gr.SelectData): + gen = get_gen_info(state) + queue = gen.get("queue", []) + + if evt.index is None: + return gr.update(), gr.update(), gr.update(visible=False) + row_index, col_index = evt.index + cell_value = None + if col_index in [6, 7, 8]: + if col_index == 6: cell_value = "↑" + elif col_index == 7: cell_value = "↓" + elif col_index == 8: cell_value = "✖" + if col_index == 6: + new_df_data = move_up(queue, [row_index]) + return new_df_data, gr.update(), gr.update(visible=False) + elif col_index == 7: + new_df_data = move_down(queue, [row_index]) + return new_df_data, gr.update(), gr.update(visible=False) + elif col_index == 8: + new_df_data = remove_task(queue, [row_index]) + gen["prompts_max"] = gen.get("prompts_max",0) - 1 + update_status(state) + return new_df_data, gr.update(), gr.update(visible=False) + start_img_col_idx = 4 + end_img_col_idx = 5 + image_data_to_show = None + if col_index == start_img_col_idx: + with lock: + row_index += 1 + if row_index < len(queue): + image_data_to_show = queue[row_index].get('start_image_data') + elif col_index == end_img_col_idx: + with lock: + row_index += 1 + if row_index < len(queue): + image_data_to_show = queue[row_index].get('end_image_data') + + if image_data_to_show: + return gr.update(), gr.update(value=image_data_to_show[0]), gr.update(visible=True) + else: + return gr.update(), gr.update(), gr.update(visible=False) + + +def change_model(state, model_choice): + model_filename = "" + for filename in model_list: + if get_model_type(filename) == model_choice: + model_filename = filename + break + if len(model_filename) == 0: + return + + state["model_filename"] = model_filename + header = generate_header(compile=compile, attention_mode=attention_mode) + return header + +def fill_inputs(state): + model_filename = state["model_filename"] + prefix = get_model_type(model_filename) + ui_defaults = state.get(prefix, None) + if ui_defaults == None: + ui_defaults = get_default_settings(model_filename) + + return generate_video_tab(update_form = True, state_dict = state, ui_defaults = ui_defaults) + +def preload_model(state): + global reload_needed, wan_model, offloadobj + if reload_model == 1: + model_filename = state["model_filename"] + if state["model_filename"] != transformer_filename: + wan_model = None + if offloadobj is not None: + offloadobj.release() + offloadobj = None + gc.collect() + yield f"Loading model {get_model_name(model_filename)}..." + wan_model, offloadobj, _ = load_models(model_filename) + yield f"Model loaded" + reload_needed= False + return 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) +def unload_model_if_needed(state): + global reload_needed, wan_model, offloadobj + if reload_model == 3: + if wan_model != None: + wan_model = None + if offloadobj is not None: + offloadobj.release() + offloadobj = None + gc.collect() + reload_needed= True - state_dict = {} - state_dict["advanced"] = advanced - state_dict["loras_model"] = filename - state_dict["image_input_type_model"] = filename - state_dict["image2video"] = image2video - gen = dict() - gen["queue"] = [] - state_dict["gen"] = gen +def generate_video_tab(update_form = False, state_dict = None, ui_defaults = None, model_choice = None, header = None): + global inputs_names #, advanced - preset_to_load = lora_preselected_preset if use_image2video == image2video else "" + if update_form: + model_filename = state_dict["model_filename"] + advanced_ui = state_dict["advanced"] + else: + model_filename = transformer_filename + advanced_ui = advanced + ui_defaults= get_default_settings(model_filename) + state_dict = {} + state_dict["model_filename"] = model_filename + state_dict["advanced"] = advanced_ui + gen = dict() + gen["queue"] = [] + state_dict["gen"] = gen - loras, loras_names, loras_presets, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset = setup_loras(image2video, None, get_lora_dir(image2video), preset_to_load, None) + preset_to_load = lora_preselected_preset if lora_preset_model == model_filename else "" + + loras, loras_names, loras_presets, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset = setup_loras(model_filename, None, get_lora_dir(model_filename), preset_to_load, None) state_dict["loras"] = loras state_dict["loras_presets"] = loras_presets @@ -2479,18 +2585,19 @@ def generate_video_tab(image2video=False): launch_loras = [] launch_multis_str = "" - if len(default_lora_preset) > 0 and image2video == use_image2video: + if update_form: + pass + if len(default_lora_preset) > 0 and lora_preset_model == model_filename: launch_preset = default_lora_preset launch_prompt = default_lora_preset_prompt launch_loras = default_loras_choices launch_multis_str = default_loras_multis_str if len(launch_prompt) == 0: - launch_prompt = ui_defaults["prompts"] + launch_prompt = ui_defaults.get("prompt","") if len(launch_loras) == 0: - activated_loras = ui_defaults["activated_loras"] - launch_multis_str = ui_defaults["loras_multipliers"] - + launch_multis_str = ui_defaults.get("loras_multipliers","") + activated_loras = ui_defaults.get("activated_loras",[]) if len(activated_loras) > 0: lora_filenames = [os.path.basename(lora_path) for lora_path in loras] activated_indices = [] @@ -2502,24 +2609,20 @@ def generate_video_tab(image2video=False): print(f"Warning: Lora file {lora_file} from config not found in loras directory") launch_loras = activated_indices - - header = gr.Markdown(generate_header(model_filename, compile, attention_mode)) with gr.Row(): with gr.Column(): with gr.Column(visible=False, elem_id="image-modal-container") as modal_container: with gr.Row(elem_id="image-modal-close-button-row"): close_modal_button = gr.Button("❌", size="sm") modal_image_display = gr.Image(label="Full Resolution Image", interactive=False, show_label=False) - progress_update_trigger = gr.Textbox(value="0", visible=False, label="_progress_trigger") - gallery_update_trigger = gr.Textbox(value="0", visible=False, label="_gallery_trigger") - with gr.Row(visible= len(loras)>0) as presets_column: - lset_choices = [ (preset, preset) for preset in loras_presets ] + [(get_new_preset_msg(advanced), "")] + with gr.Row(visible= True): #len(loras)>0) as presets_column: + lset_choices = [ (preset, preset) for preset in loras_presets ] + [(get_new_preset_msg(advanced_ui), "")] with gr.Column(scale=6): lset_name = gr.Dropdown(show_label=False, allow_custom_value= True, scale=5, filterable=True, choices= lset_choices, value=launch_preset) with gr.Column(scale=1): with gr.Row(height=17): apply_lset_btn = gr.Button("Apply Lora Preset", size="sm", min_width= 1) - refresh_lora_btn = gr.Button("Refresh", size="sm", min_width= 1, visible=advanced or not only_allow_edit_in_advanced) + refresh_lora_btn = gr.Button("Refresh", size="sm", min_width= 1, visible=advanced_ui or not only_allow_edit_in_advanced) save_lset_prompt_drop= gr.Dropdown( choices=[ ("Save Prompt Comments Only", 0), @@ -2529,62 +2632,50 @@ def generate_video_tab(image2video=False): with gr.Row(height=17, visible=False) as refresh2_row: refresh_lora_btn2 = gr.Button("Refresh", size="sm", min_width= 1) - with gr.Row(height=17, visible=advanced or not only_allow_edit_in_advanced) as preset_buttons_rows: + with gr.Row(height=17, visible=advanced_ui or not only_allow_edit_in_advanced) as preset_buttons_rows: confirm_save_lset_btn = gr.Button("Go Ahead Save it !", size="sm", min_width= 1, visible=False) confirm_delete_lset_btn = gr.Button("Go Ahead Delete it !", size="sm", min_width= 1, visible=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) - state = gr.State(state_dict) - vace_model = "Vace" in filename and not image2video + if not update_form: + state = gr.State(state_dict) 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) + with gr.Column(visible= "image2video" in model_filename or "Fun_InP" in model_filename ) as image_prompt_column: + image_prompt_type_value= ui_defaults.get("image_prompt_type","S") + image_prompt_type = gr.Radio( [("Use only a Start Image", "S"),("Use both a Start and an End Image", "SE")], value =image_prompt_type_value, label="Location", show_label= False, scale= 3) - 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") - - 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) + if args.multiple_images: + image_start = 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, value= ui_defaults.get("image_start", None), visible= "S" in image_prompt_type_value) else: - image_prompt_type= ui_defaults.get("image_prompt_type","I") - if not isinstance(image_prompt_type, str): - 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_start = gr.Image(label= "Image as a starting point for a new video", type ="pil",value= ui_defaults.get("image_start", None), visible= "S" in image_prompt_type_value ) - 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 ) + if args.multiple_images: + image_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="E" in image_prompt_type_value, value= ui_defaults.get("image_end", None)) + else: + image_end = gr.Image(label= "Last Image for a new video", type ="pil", visible="E" in image_prompt_type_value, value= ui_defaults.get("image_end", None)) - image_source3 = gr.Video(label= "Video Mask (white pixels = Mask)", visible= "M" in image_prompt_type ) + with gr.Column(visible= "Vace" in model_filename ) as video_prompt_column: + video_prompt_type_value= ui_defaults.get("video_prompt_type","I") + video_prompt_type = gr.Radio( [("Use Images Ref", "I"),("a Video", "V"), ("Images + a Video", "IV"), ("Video + Video Mask", "VM"), ("Images + Video + Mask", "IVM")], value =video_prompt_type_value, label="Location", show_label= False, scale= 3) + image_refs = 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 video_prompt_type_value, value= ui_defaults.get("image_refs", None) ) - - 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]) + video_guide = gr.Video(label= "Reference Video", visible= "V" in video_prompt_type_value, value= ui_defaults.get("video_guide", None) ) + with gr.Row(): + max_frames = gr.Slider(0, 100, value=ui_defaults.get("max_frames",0), step=1, label="Nb of frames in Ref. Video (0 = as many as possible)", visible= "V" in video_prompt_type_value, 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 video_prompt_type_value, scale =1 ) + + video_mask = gr.Video(label= "Video Mask (white pixels = Mask)", visible= "M" in video_prompt_type_value, value= ui_defaults.get("video_mask", None) ) - advanced_prompt = advanced + advanced_prompt = advanced_ui prompt_vars=[] if advanced_prompt: @@ -2615,14 +2706,14 @@ def generate_video_tab(image2video=False): wizard_prompt_activated_var = gr.Text(wizard_prompt_activated, visible= False) wizard_variables_var = gr.Text(wizard_variables, visible = False) with gr.Row(): - if image2video: + if "image2video" in model_filename or "Fun_InP" in model_filename: resolution = gr.Dropdown( choices=[ # 720p ("720p", "1280x720"), ("480p", "832x480"), ], - value=ui_defaults["resolution"], + value=ui_defaults.get("resolution","480p"), label="Resolution (video will have the same height / width ratio than the original image)" ) else: @@ -2642,33 +2733,33 @@ def generate_video_tab(image2video=False): # ("624x832 (3:4, 540p)", "624x832"), # ("720x720 (1:1, 540p)", "720x720"), ], - value=ui_defaults["resolution"], + value=ui_defaults.get("resolution","832x480"), label="Resolution" ) with gr.Row(): with gr.Column(): - video_length = gr.Slider(5, 193, value=ui_defaults["video_length"], step=4, label="Number of frames (16 = 1s)") + video_length = gr.Slider(5, 193, value=ui_defaults.get("video_length", 81), 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") - show_advanced = gr.Checkbox(label="Advanced Mode", value=advanced) - with gr.Row(visible=advanced) as advanced_row: + num_inference_steps = gr.Slider(1, 100, value=ui_defaults.get("num_inference_steps",30), step=1, label="Number of Inference Steps") + show_advanced = gr.Checkbox(label="Advanced Mode", value=advanced_ui) + with gr.Row(visible=advanced_ui) as advanced_row: with gr.Column(): seed = gr.Slider(-1, 999999999, value=ui_defaults["seed"], step=1, label="Seed (-1 for random)") with gr.Row(): - repeat_generation = gr.Slider(1, 25.0, value=ui_defaults["repeat_generation"], step=1, label="Default Number of Generated Videos per Prompt") - multi_images_gen_type = gr.Dropdown( value=ui_defaults["multi_images_gen_type"], + repeat_generation = gr.Slider(1, 25.0, value=ui_defaults.get("repeat_generation",1), step=1, label="Default Number of Generated Videos per Prompt") + multi_images_gen_type = gr.Dropdown( value=ui_defaults.get("multi_images_gen_type",0), choices=[ ("Generate every combination of images and texts", 0), ("Match images and text prompts", 1), ], visible= args.multiple_images, label= "Multiple Images as Texts Prompts" ) with gr.Row(): - guidance_scale = gr.Slider(1.0, 20.0, value=ui_defaults["guidance_scale"], step=0.5, label="Guidance Scale", visible=True) + guidance_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("guidance_scale",5), step=0.5, label="Guidance Scale", visible=True) embedded_guidance_scale = gr.Slider(1.0, 20.0, value=6.0, step=0.5, label="Embedded Guidance Scale", visible=False) - flow_shift = gr.Slider(0.0, 25.0, value=ui_defaults["flow_shift"], step=0.1, label="Shift Scale") + flow_shift = gr.Slider(0.0, 25.0, value=ui_defaults.get("flow_shift",3), step=0.1, label="Shift Scale") with gr.Row(): - negative_prompt = gr.Textbox(label="Negative Prompt", value=ui_defaults["negative_prompt"]) - with gr.Column(visible = len(loras)>0) as loras_column: + negative_prompt = gr.Textbox(label="Negative Prompt", value=ui_defaults.get("negative_prompt", "") ) + with gr.Column(visible = True): #as loras_column: gr.Markdown("Loras can be used to create special effects on the video by mentioning a trigger word in the Prompt. You can save Loras combinations in presets.") loras_choices = gr.Dropdown( choices=[ @@ -2678,7 +2769,7 @@ def generate_video_tab(image2video=False): multiselect= True, label="Activated Loras" ) - loras_mult_choices = gr.Textbox(label="Loras Multipliers (1.0 by default) separated by space characters or carriage returns, line that starts with # are ignored", value=launch_multis_str) + loras_multipliers = gr.Textbox(label="Loras Multipliers (1.0 by default) separated by space characters or carriage returns, line that starts with # are ignored", value=launch_multis_str) with gr.Row(): gr.Markdown("Tea Cache accelerates by skipping intelligently some steps, the more steps are skipped the lower the quality of the video (Tea Cache consumes also VRAM)") with gr.Row(): @@ -2691,16 +2782,16 @@ def generate_video_tab(image2video=False): ("around x2.25 speed up", 2.25), ("around x2.5 speed up", 2.5), ], - value=float(ui_defaults["tea_cache"]), + value=float(ui_defaults.get("tea_cache_setting",0)), visible=True, label="Tea Cache Global Acceleration" ) - 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") + tea_cache_start_step_perc = gr.Slider(0, 100, value=ui_defaults.get("tea_cache_start_step_perc",0), step=1, label="Tea Cache starting moment in % of generation") with gr.Row(): gr.Markdown("Upsampling - postprocessing that may improve fluidity and the size of the video") with gr.Row(): - temporal_upsampling_choice = gr.Dropdown( + temporal_upsampling = gr.Dropdown( choices=[ ("Disabled", ""), ("Rife x2 (32 frames/s)", "rife2"), @@ -2711,7 +2802,7 @@ def generate_video_tab(image2video=False): scale = 1, label="Temporal Upsampling" ) - spatial_upsampling_choice = gr.Dropdown( + spatial_upsampling = gr.Dropdown( choices=[ ("Disabled", ""), ("Lanczos x1.5", "lanczos1.5"), @@ -2730,7 +2821,7 @@ def generate_video_tab(image2video=False): ("Always ON", 1), ("Always OFF", 2), ], - value=ui_defaults["RIFLEx_setting"], + value=ui_defaults.get("RIFLEx_setting",0), label="RIFLEx positional embedding to generate long video" ) with gr.Row(): @@ -2741,7 +2832,7 @@ def generate_video_tab(image2video=False): ("OFF", 0), ("ON", 1), ], - value=ui_defaults["slg_switch"], + value=ui_defaults.get("slg_switch",0), visible=True, scale = 1, label="Skip Layer guidance" @@ -2750,14 +2841,14 @@ def generate_video_tab(image2video=False): choices=[ (str(i), i ) for i in range(40) ], - value=ui_defaults["slg_layers"], + value=ui_defaults.get("slg_layers", ["9"]), multiselect= True, label="Skip Layers", scale= 3 ) with gr.Row(): - slg_start_perc = gr.Slider(0, 100, value=ui_defaults["slg_start_perc"], step=1, label="Denoising Steps % start") - slg_end_perc = gr.Slider(0, 100, value=ui_defaults["slg_end_perc"], step=1, label="Denoising Steps % end") + slg_start_perc = gr.Slider(0, 100, value=ui_defaults.get("slg_start_perc",10), step=1, label="Denoising Steps % start") + slg_end_perc = gr.Slider(0, 100, value=ui_defaults.get("slg_end_perc",90), step=1, label="Denoising Steps % end") with gr.Row(): gr.Markdown("Experimental: Classifier-Free Guidance Zero Star, better adherence to Text Prompt") @@ -2777,239 +2868,142 @@ def generate_video_tab(image2video=False): with gr.Row(): save_settings_btn = gr.Button("Set Settings as Default", visible = not args.lock_config) + + if not update_form: + with gr.Column(): + gen_status = gr.Text(interactive= False, label = "Status") + output = gr.Gallery( label="Generated videos", show_label=False, elem_id="gallery" , columns=[3], rows=[1], object_fit="contain", height=450, selected_index=0, interactive= False) + generate_btn = gr.Button("Generate") + add_to_queue_btn = gr.Button("Add New Prompt To Queue", visible = False) + + with gr.Column(visible= False) as current_gen_column: + with gr.Row(): + gen_info = gr.HTML(visible=False, min_height=1) + with gr.Row(): + onemore_btn = gr.Button("One More Sample Please !") + abort_btn = gr.Button("Abort") + + queue_df = gr.DataFrame( + 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"), + wrap=True, + value=[], + line_breaks= True, + visible= False, + elem_id="queue_df" + ) + + extra_inputs = prompt_vars + [wizard_prompt, wizard_variables_var, wizard_prompt_activated_var, video_prompt_column, image_prompt_column, + prompt_column_advanced, prompt_column_wizard_vars, prompt_column_wizard, lset_name, advanced_row] # show_advanced presets_column, + if update_form: + locals_dict = locals() + gen_inputs = [state_dict if k=="state" else locals_dict[k] for k in inputs_names] + [state_dict] + extra_inputs + return gen_inputs + else: + target_state = gr.Text(value = "state", interactive= False, visible= False) + target_settings = gr.Text(value = "settings", interactive= False, visible= False) + + image_prompt_type.change(fn=refresh_image_prompt_type, inputs=[state, image_prompt_type], outputs=[image_start, image_end]) + video_prompt_type.change(fn=refresh_video_prompt_type, inputs=[state, video_prompt_type], outputs=[image_refs, video_guide, video_mask, max_frames, remove_background_image_ref]) 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(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", - elem_id="generation_progress_bar_container", visible= False + queue_df.select( fn=handle_celll_selection, inputs=state, outputs=[queue_df, modal_image_display, modal_container]) + 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]) + confirm_save_lset_btn.click(fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , outputs= [prompt]).then( + save_lset, inputs=[state, lset_name, loras_choices, loras_multipliers, prompt, save_lset_prompt_drop], outputs=[lset_name, apply_lset_btn,refresh_lora_btn, delete_lset_btn, save_lset_btn, confirm_save_lset_btn, cancel_lset_btn, save_lset_prompt_drop]) + delete_lset_btn.click(validate_delete_lset, inputs=[lset_name], outputs=[apply_lset_btn, refresh_lora_btn, delete_lset_btn, save_lset_btn,confirm_delete_lset_btn, cancel_lset_btn ]) + confirm_delete_lset_btn.click(delete_lset, inputs=[state, lset_name], outputs=[lset_name, apply_lset_btn, refresh_lora_btn, delete_lset_btn, save_lset_btn,confirm_delete_lset_btn, cancel_lset_btn ]) + cancel_lset_btn.click(cancel_lset, inputs=[], outputs=[apply_lset_btn, refresh_lora_btn, delete_lset_btn, save_lset_btn, confirm_delete_lset_btn,confirm_save_lset_btn, cancel_lset_btn,save_lset_prompt_drop ]) + apply_lset_btn.click(apply_lset, inputs=[state, wizard_prompt_activated_var, lset_name,loras_choices, loras_multipliers, prompt], outputs=[wizard_prompt_activated_var, loras_choices, loras_multipliers, prompt]).then( + fn = fill_wizard_prompt, inputs = [state, wizard_prompt_activated_var, prompt, wizard_prompt], outputs = [ wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, prompt_column_advanced, prompt_column_wizard, prompt_column_wizard_vars, *prompt_vars] ) - output = gr.Gallery( - label="Generated videos", show_label=False, elem_id="gallery" - , columns=[3], rows=[1], object_fit="contain", height=450, selected_index=0, interactive= False) - generate_btn = gr.Button("Generate") - add_to_queue_btn = gr.Button("Add New Prompt To Queue", visible = False) - - with gr.Column(visible= False) as current_gen_column: - with gr.Row(): - gen_info = gr.HTML(visible=False, min_height=1) - with gr.Row(): - onemore_btn = gr.Button("One More Sample Please !") - abort_btn = gr.Button("Abort") - - queue_df = gr.DataFrame( - 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"), - wrap=True, - value=[], - line_breaks= True, - visible= False, - # every=1, - elem_id="queue_df" - ) - # queue_df = gr.HTML("", - # visible= False, - # elem_id="queue_df" - # ) - - def handle_selection(state, evt: gr.SelectData): - gen = get_gen_info(state) - queue = gen.get("queue", []) - - if evt.index is None: - return gr.update(), gr.update(), gr.update(visible=False) - row_index, col_index = evt.index - cell_value = None - if col_index in [6, 7, 8]: - if col_index == 6: cell_value = "↑" - elif col_index == 7: cell_value = "↓" - elif col_index == 8: cell_value = "✖" - if col_index == 6: - new_df_data = move_up(queue, [row_index]) - return new_df_data, gr.update(), gr.update(visible=False) - elif col_index == 7: - new_df_data = move_down(queue, [row_index]) - return new_df_data, gr.update(), gr.update(visible=False) - elif col_index == 8: - new_df_data = remove_task(queue, [row_index]) - gen["prompts_max"] = gen.get("prompts_max",0) - 1 - update_status(state) - return new_df_data, gr.update(), gr.update(visible=False) - start_img_col_idx = 4 - end_img_col_idx = 5 - image_data_to_show = None - if col_index == start_img_col_idx: - with lock: - if row_index < len(queue): - image_data_to_show = queue[row_index].get('start_image_data') - elif col_index == end_img_col_idx: - with lock: - if row_index < len(queue): - image_data_to_show = queue[row_index].get('end_image_data') - - if image_data_to_show: - return gr.update(), gr.update(value=image_data_to_show), gr.update(visible=True) - else: - return gr.update(), gr.update(), gr.update(visible=False) - selected_indices = gr.State([]) - queue_df.select( - fn=handle_selection, - inputs=state, - outputs=[queue_df, modal_image_display, modal_container], - ) - # gallery_update_trigger.change( - # fn=refresh_gallery_on_trigger, - # inputs=[state], - # outputs=[output] - # ) - # queue_df.change( - # fn=refresh_gallery, - # inputs=[state], - # outputs=[gallery_update_trigger] - # ).then( - # fn=refresh_progress, - # inputs=None, - # outputs=[progress_update_trigger] - # ) - progress_update_trigger.change( - fn=update_generation_status, - inputs=[progress_update_trigger], - outputs=[gen_progress_html], - 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, 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]) - confirm_save_lset_btn.click(fn=validate_wizard_prompt, inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , outputs= [prompt]).then( - save_lset, inputs=[state, lset_name, loras_choices, loras_mult_choices, prompt, save_lset_prompt_drop], outputs=[lset_name, apply_lset_btn,refresh_lora_btn, delete_lset_btn, save_lset_btn, confirm_save_lset_btn, cancel_lset_btn, save_lset_prompt_drop]) - delete_lset_btn.click(validate_delete_lset, inputs=[lset_name], outputs=[apply_lset_btn, refresh_lora_btn, delete_lset_btn, save_lset_btn,confirm_delete_lset_btn, cancel_lset_btn ]) - confirm_delete_lset_btn.click(delete_lset, inputs=[state, lset_name], outputs=[lset_name, apply_lset_btn, refresh_lora_btn, delete_lset_btn, save_lset_btn,confirm_delete_lset_btn, cancel_lset_btn ]) - cancel_lset_btn.click(cancel_lset, inputs=[], outputs=[apply_lset_btn, refresh_lora_btn, delete_lset_btn, save_lset_btn, confirm_delete_lset_btn,confirm_save_lset_btn, cancel_lset_btn,save_lset_prompt_drop ]) - apply_lset_btn.click(apply_lset, inputs=[state, wizard_prompt_activated_var, lset_name,loras_choices, loras_mult_choices, prompt], outputs=[wizard_prompt_activated_var, loras_choices, loras_mult_choices, prompt]).then( - fn = fill_wizard_prompt, inputs = [state, wizard_prompt_activated_var, prompt, wizard_prompt], outputs = [ wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, prompt_column_advanced, prompt_column_wizard, prompt_column_wizard_vars, *prompt_vars] - ) - refresh_lora_btn.click(refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) - refresh_lora_btn2.click(refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) - output.select(select_video, state, None ) - - + refresh_lora_btn.click(refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) + 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]) - 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(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, - 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(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] - ) + abort_btn.click(abort_generation, [state], [gen_status, abort_btn] ) #.then(refresh_gallery, inputs = [state, gen_info], outputs = [output, gen_info, queue_df] ) + onemore_btn.click(fn=one_more_sample,inputs=[state], outputs= [state]) - abort_btn.click(abort_generation, [state], [gen_status, abort_btn] ) #.then(refresh_gallery, inputs = [state, gen_info], outputs = [output, gen_info, queue_df] ) - onemore_btn.click(fn=one_more_sample,inputs=[state], outputs= [state]) + inputs_names= list(inspect.signature(save_inputs).parameters)[1:-1] + locals_dict = locals() + gen_inputs = [locals_dict[k] for k in inputs_names] + [state] + 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_inputs, inputs =[target_settings] + gen_inputs, outputs = []) - gen_inputs=[ - prompt, - negative_prompt, - resolution, - video_length, - seed, - num_inference_steps, - guidance_scale, - flow_shift, - embedded_guidance_scale, - repeat_generation, - multi_images_gen_type, - tea_cache_setting, - tea_cache_start_step_perc, - loras_choices, - loras_mult_choices, - image_prompt_type_radio, - image_source1, - image_source2, - image_source3, - max_frames, - remove_background_image_ref, - temporal_upsampling_choice, - spatial_upsampling_choice, - RIFLEx_setting, - slg_switch, - slg_layers, - slg_start_perc, - slg_end_perc, - cfg_star_switch, - cfg_zero_step, - state, - gr.State(image2video) - ] + model_choice.change(fn=validate_wizard_prompt, + inputs= [state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , + outputs= [prompt] + ).then(fn=save_inputs, + inputs =[target_state] + gen_inputs, + outputs= None + ).then(fn= change_model, + inputs=[state, model_choice], + outputs= [header] + ).then(fn= fill_inputs, + inputs=[state], + outputs=gen_inputs + extra_inputs + ).then(fn= preload_model, + inputs=[state], + outputs=[gen_status]) - generate_btn.click(fn=validate_wizard_prompt, - inputs= [state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , - outputs= [prompt] - ).then(fn=process_prompt_and_add_tasks, - inputs = gen_inputs, - outputs= queue_df - ).then(fn=prepare_generate_video, - inputs= [state], - outputs= [generate_btn, add_to_queue_btn, current_gen_column], - ).then(fn=process_tasks, - 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] - ) + generate_btn.click(fn=validate_wizard_prompt, + inputs= [state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , + outputs= [prompt] + ).then(fn=save_inputs, + inputs =[target_state] + gen_inputs, + outputs= None + ).then(fn=process_prompt_and_add_tasks, + inputs = [state, model_choice], + outputs= queue_df + ).then(fn=prepare_generate_video, + inputs= [state], + outputs= [generate_btn, add_to_queue_btn, current_gen_column], + ).then(fn=process_tasks, + 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] + ).then(unload_model_if_needed, + inputs= [state], + outputs= [] + ) - add_to_queue_btn.click(fn=validate_wizard_prompt, - inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , - outputs= [prompt] - ).then( - fn=process_prompt_and_add_tasks, - inputs = gen_inputs, - outputs=queue_df - ).then( - fn=update_status, - inputs = [state], - ) + add_to_queue_btn.click(fn=validate_wizard_prompt, + inputs =[state, wizard_prompt_activated_var, wizard_variables_var, prompt, wizard_prompt, *prompt_vars] , + outputs= [prompt] + ).then(fn=save_inputs, + inputs =[target_state] + gen_inputs, + outputs= None + ).then(fn=process_prompt_and_add_tasks, + inputs = [state, model_choice], + outputs=queue_df + ).then( + fn=update_status, + inputs = [state], + ) + close_modal_button.click( + lambda: gr.update(visible=False), + inputs=[], + outputs=[modal_container] + ) - close_modal_button.click( - lambda: gr.update(visible=False), - inputs=[], - outputs=[modal_container] - ) - return loras_column, loras_choices, presets_column, lset_name, header, light_sync, full_sync, state + return loras_choices, lset_name, state -def generate_download_tab(presets_column, loras_column, lset_name,loras_choices, state): +def generate_download_tab(lset_name,loras_choices, state): with gr.Row(): with gr.Row(scale =2): - gr.Markdown("Wan2GP's Lora Festival ! Press the following button to download i2v Remade Loras collection (and bonuses Loras).") + gr.Markdown("WanGP's Lora Festival ! Press the following button to download i2v Remade_AI Loras collection (and bonuses Loras).") with gr.Row(scale =1): download_loras_btn = gr.Button("---> Let the Lora's Festival Start !", scale =1) with gr.Row(scale =1): @@ -3017,7 +3011,7 @@ def generate_download_tab(presets_column, loras_column, lset_name,loras_choices, with gr.Row() as download_status_row: download_status = gr.Markdown() - download_loras_btn.click(fn=download_loras, inputs=[], outputs=[download_status_row, download_status, presets_column, loras_column]).then(fn=refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) + download_loras_btn.click(fn=download_loras, inputs=[], outputs=[download_status_row, download_status]).then(fn=refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) def generate_configuration_tab(): @@ -3025,37 +3019,50 @@ def generate_configuration_tab(): state = gr.State(state_dict) gr.Markdown("Please click Apply Changes at the bottom so that the changes are effective. Some choices below may be locked if the app has been launched by specifying a config preset.") with gr.Column(): - index = transformer_choices_t2v.index(transformer_filename_t2v) + index = transformer_choices.index(transformer_filename) index = 0 if index ==0 else index - transformer_t2v_choice = gr.Dropdown( + + model_list = [] + for model_type in model_types: + choice = get_model_filename(model_type, transformer_quantization) + model_list.append(choice) + dropdown_choices = [ ( get_model_name(choice), get_model_type(choice) ) for choice in model_list] + transformer_type_choice = gr.Dropdown( + choices= dropdown_choices, + value= get_model_type(transformer_filename), + label= "Default Wan Transformer Model", + scale= 2 + ) + + # transformer_choice = gr.Dropdown( + # choices=[ + # ("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), + # ("WAN 2.1 - 480p 14B Image to Video 16 bits - the default engine in its original glory, offers a slightly better image quality but slower and requires more RAM", 4), + # ("WAN 2.1 - 480p 14B Image to Video quantized to 8 bits (recommended) - the default engine but quantized", 5), + # ("WAN 2.1 - 720p 14B Image to Video 16 bits - the default engine in its original glory, offers a slightly better image quality but slower and requires more RAM", 6), + # ("WAN 2.1 - 720p 14B Image to Video quantized to 8 bits - the default engine but quantized", 7), + # ("WAN 2.1 - Fun InP 1.3B 16 bits - the small model for fast generations with low VRAM requirements", 8), + # ("WAN 2.1 - Fun InP 14B 16 bits - Fun InP version in its original glory, offers a slightly better image quality but slower and requires more RAM", 9), + # ("WAN 2.1 - Fun InP 14B quantized to 8 bits - quantized Fun InP version", 10), + # ], + # value= index, + # label="Transformer model for Image to Video", + # interactive= not lock_ui_transformer, + # visible = True, + # ) + + quantization_choice = gr.Dropdown( choices=[ - ("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), + ("Int8 Quantization (recommended)", "int8"), + ("BF16 (no quantization)", "bf16"), ], - value= index, - label="Transformer model for Text to Video", - interactive= not lock_ui_transformer, - visible=True - ) - index = transformer_choices_i2v.index(transformer_filename_i2v) - index = 0 if index ==0 else index - transformer_i2v_choice = gr.Dropdown( - choices=[ - ("WAN 2.1 - 480p 14B Image to Video 16 bits - the default engine in its original glory, offers a slightly better image quality but slower and requires more RAM", 0), - ("WAN 2.1 - 480p 14B Image to Video quantized to 8 bits (recommended) - the default engine but quantized", 1), - ("WAN 2.1 - 720p 14B Image to Video 16 bits - the default engine in its original glory, offers a slightly better image quality but slower and requires more RAM", 2), - ("WAN 2.1 - 720p 14B Image to Video quantized to 8 bits - the default engine but quantized", 3), - ("WAN 2.1 - Fun InP 1.3B 16 bits - the small model for fast generations with low VRAM requirements", 4), - ("WAN 2.1 - Fun InP 14B 16 bits - Fun InP version in its original glory, offers a slightly better image quality but slower and requires more RAM", 5), - ("WAN 2.1 - Fun InP 14B quantized to 8 bits - quantized Fun InP version", 6), - ], - value= index, - label="Transformer model for Image to Video", - interactive= not lock_ui_transformer, - visible = True, - ) + value= transformer_quantization, + label="Wan Transformer Model Quantization (if available)", + ) + index = text_encoder_choices.index(text_encoder_filename) index = 0 if index ==0 else index text_encoder_choice = gr.Dropdown( @@ -3149,11 +3156,12 @@ def generate_configuration_tab(): ) reload_choice = gr.Dropdown( choices=[ - ("When changing tabs", 1), - ("When pressing Generate", 2), + ("Load Model When Changing Model", 1), + ("Load Model When Pressing Generate", 2), + ("Load Model When Pressing Generate and Unload Model when Finished", 3), ], value=server_config.get("reload_model",2), - label="Reload model" + label="RAM Loading / Unloading Model Policy (in any case VRAM will be freed once the queue has been processed)" ) clear_file_list_choice = gr.Dropdown( @@ -3176,8 +3184,7 @@ def generate_configuration_tab(): fn=apply_changes, inputs=[ state, - transformer_t2v_choice, - transformer_i2v_choice, + transformer_type_choice, text_encoder_choice, save_path_choice, attention_choice, @@ -3185,7 +3192,7 @@ def generate_configuration_tab(): profile_choice, vae_config_choice, metadata_choice, - default_ui_choice, + quantization_choice, boost_choice, clear_file_list_choice, reload_choice, @@ -3194,7 +3201,7 @@ def generate_configuration_tab(): ) def generate_about_tab(): - gr.Markdown("

Wan2.1GP - Wan 2.1 model for the GPU Poor by DeepBeepMeep (GitHub)

") + gr.Markdown("

WanGP - Wan 2.1 model for the GPU Poor by DeepBeepMeep (GitHub)

") gr.Markdown("Original Wan 2.1 Model by Alibaba (GitHub)") gr.Markdown("Many thanks to:") gr.Markdown("- Alibaba Wan team for the best open source video generator") @@ -3202,114 +3209,19 @@ def generate_about_tab(): gr.Markdown("- Tophness : created multi tabs and queuing frameworks") gr.Markdown("- AmericanPresidentJimmyCarter : added original support for Skip Layer Guidance") gr.Markdown("- Remade_AI : for creating their awesome Loras collection") + +def generate_info_tab(): + gr.Markdown("Welcome to WanGP a super fast and low VRAM AI Video Generator !") - -def on_tab_select(global_state, t2v_state, i2v_state, evt: gr.SelectData): - t2v_header = generate_header(transformer_filename_t2v, compile, attention_mode) - i2v_header = generate_header(transformer_filename_i2v, compile, attention_mode) - - new_t2v = evt.index == 0 - new_i2v = evt.index == 1 - i2v_light_sync = gr.Text() - t2v_light_sync = gr.Text() - i2v_full_sync = gr.Text() - t2v_full_sync = gr.Text() - - 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 + gr.Markdown("The VRAM requirements will depend greatly of the resolution and the duration of the video, for instance :") + gr.Markdown("- 848 x 480 with a 14B model: 80 frames (5s) : 8 GB of VRAM") + gr.Markdown("- 848 x 480 with the 1.3B model: 80 frames (5s) : 5 GB of VRAM") + gr.Markdown("- 1280 x 720 with a 14B model: 80 frames (5s): 11 GB of VRAM") + gr.Markdown("It is not recommmended to generate a video longer than 8s (128 frames) even if there is still some VRAM left as some artifacts may appear") + gr.Markdown("Please note that if your turn on compilation, the first denoising step of the first video generation will be slow due to the compilation. Therefore all your tests should be done with compilation turned off.") - 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) : - if new_i2v: - i2v_full_sync = gr.Text(str(time.time())) - else: - t2v_full_sync = gr.Text(str(time.time())) - else: - if new_i2v: - i2v_light_sync = gr.Text(str(time.time())) - 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): - queue = gen.get("queue", []) - - queue_empty = len(queue) == 0 - if queue_empty: - global wan_model, offloadobj - if wan_model is not None: - if offloadobj is not None: - offloadobj.release() - offloadobj = None - wan_model = None - gc.collect() - torch.cuda.empty_cache() - wan_model, offloadobj, trans = load_models(new_i2v) - del trans - - if new_t2v or new_i2v: - state = i2v_state if new_i2v else t2v_state - lora_model_filename = state["loras_model"] - model_filename = model_needed(new_i2v) - if ("1.3B" in model_filename and not "1.3B" in lora_model_filename or "14B" in model_filename and not "14B" in lora_model_filename): - lora_dir = get_lora_dir(new_i2v) - loras, loras_names, loras_presets, _, _, _, _ = setup_loras(new_i2v, None, lora_dir, lora_preselected_preset, None) - state["loras"] = loras - state["loras_names"] = loras_names - state["loras_presets"] = loras_presets - state["loras_model"] = model_filename - - advanced = state["advanced"] - new_loras_choices = [(name, str(i)) for i, name in enumerate(loras_names)] - lset_choices = [(preset, preset) for preset in loras_presets] + [(get_new_preset_msg(advanced), "")] - visible = len(loras_names)>0 - if new_t2v: - return [ - gr.Column(visible= visible), - gr.Dropdown(choices=new_loras_choices, visible=visible, value=[]), - gr.Column(visible= visible), - gr.Dropdown(choices=lset_choices, value=get_new_preset_msg(advanced), visible=visible), - t2v_header, - t2v_light_sync, - t2v_full_sync, - gr.Column(), - gr.Dropdown(), - gr.Column(), - gr.Dropdown(), - gr.Markdown(), - gr.Text(), - gr.Text(), - ] - else: - return [ - gr.Column(), - gr.Dropdown(), - gr.Column(), - gr.Dropdown(), - gr.Markdown(), - gr.Text(), - gr.Text(), - gr.Text(), - gr.Column(visible= visible), - gr.Dropdown(choices=new_loras_choices, visible=visible, value=[]), - gr.Column(visible= visible), - gr.Dropdown(choices=lset_choices, value=get_new_preset_msg(advanced), visible=visible), - i2v_header, - i2v_light_sync, - i2v_full_sync, - ] - - return [gr.Column(), gr.Dropdown(), gr.Column(), gr.Dropdown(), t2v_header, t2v_light_sync, t2v_full_sync, - gr.Column(), gr.Dropdown(), gr.Column(), gr.Dropdown(), i2v_header, i2v_light_sync, i2v_full_sync] def create_demo(): @@ -3539,40 +3451,41 @@ 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("

Wan 2.1GP v4.0 by DeepBeepMeep (Updates)

") - gr.Markdown("Welcome to Wan 2.1GP a super fast and low VRAM AI Video Generator !") - - with gr.Accordion("Click here for some Info on how to use Wan2GP", open = False): - gr.Markdown("The VRAM requirements will depend greatly of the resolution and the duration of the video, for instance :") - gr.Markdown("- 848 x 480 with a 14B model: 80 frames (5s) : 8 GB of VRAM") - gr.Markdown("- 848 x 480 with the 1.3B model: 80 frames (5s) : 5 GB of VRAM") - gr.Markdown("- 1280 x 720 with a 14B model: 80 frames (5s): 11 GB of VRAM") - gr.Markdown("It is not recommmended to generate a video longer than 8s (128 frames) even if there is still some VRAM left as some artifacts may appear") - gr.Markdown("Please note that if your turn on compilation, the first denoising step of the first video generation will be slow due to the compilation. Therefore all your tests should be done with compilation turned off.") - global_dict = {} - global_dict["last_tab_was_image2video"] = use_image2video - global_state = gr.State(global_dict) + gr.Markdown("

WanGP v4.0 by DeepBeepMeep ") # (Updates)

") + global model_list - with gr.Tabs(selected="i2v" if use_image2video else "t2v") as main_tabs: - with gr.Tab("Text To Video", id="t2v") as t2v_tab: - t2v_loras_column, t2v_loras_choices, t2v_presets_column, t2v_lset_name, t2v_header, t2v_light_sync, t2v_full_sync, t2v_state = generate_video_tab(False) - with gr.Tab("Image To Video", id="i2v") as i2v_tab: - 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) + with gr.Tabs(selected="video_gen", ) as main_tabs: + with gr.Tab("Video Generator", id="video_gen") as t2v_tab: + with gr.Row(): + header = gr.Markdown(generate_header(compile, attention_mode), visible= True) + with gr.Row(): + gr.Markdown("
") + + model_list = [] + for model_type in model_types: + choice = get_model_filename(model_type, transformer_quantization) + model_list.append(choice) + dropdown_choices = [ ( get_model_name(choice), get_model_type(choice) ) for choice in model_list] + model_choice = gr.Dropdown( + choices= dropdown_choices, + value= get_model_type(transformer_filename), + show_label= False, + scale= 2 + ) + gr.Markdown("
") + with gr.Row(): + + loras_choices, lset_name, state = generate_video_tab(model_choice = model_choice, header = header) + with gr.Tab("Informations"): + generate_info_tab() if not args.lock_config: with gr.Tab("Downloads", id="downloads") as downloads_tab: - generate_download_tab(i2v_presets_column, i2v_loras_column, i2v_lset_name, i2v_loras_choices, i2v_state) + generate_download_tab(lset_name, loras_choices, state) with gr.Tab("Configuration"): generate_configuration_tab() with gr.Tab("About"): generate_about_tab() - main_tabs.select( - fn=on_tab_select, - inputs=[global_state, t2v_state, i2v_state], - outputs=[ - t2v_loras_column, t2v_loras_choices, t2v_presets_column, t2v_lset_name, t2v_header, t2v_light_sync, t2v_full_sync, - i2v_loras_column, i2v_loras_choices, i2v_presets_column, i2v_lset_name, i2v_header, i2v_light_sync, i2v_full_sync - ] - ) + return demo if __name__ == "__main__": From 5f280408f5faabc8a1b47bb90197dbe5806af736 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Mon, 7 Apr 2025 23:18:47 +0200 Subject: [PATCH 38/69] added default models selection option --- wgp.py | 162 ++++++++++++++++++++++++++++----------------------------- 1 file changed, 81 insertions(+), 81 deletions(-) diff --git a/wgp.py b/wgp.py index c81ab9a..db02e07 100644 --- a/wgp.py +++ b/wgp.py @@ -441,6 +441,12 @@ def _parse_args(): help="Prevent modifying the configuration from the web interface" ) + parser.add_argument( + "--lock-model", + action="store_true", + help="Prevent switch models" + ) + parser.add_argument( "--preload", type=str, @@ -700,11 +706,15 @@ transformer_choices_t2v=["ckpts/wan2.1_text2video_1.3B_bf16.safetensors", "ckpts 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", ] transformer_choices = transformer_choices_t2v + transformer_choices_i2v text_encoder_choices = ["ckpts/models_t5_umt5-xxl-enc-bf16.safetensors", "ckpts/models_t5_umt5-xxl-enc-quanto_int8.safetensors"] -server_config_filename = "gradio_config.json" +server_config_filename = "wgp_config.json" + +if not os.path.isfile(server_config_filename) and os.path.isfile("gradio_config.json"): + import shutil + shutil.move("gradio_config.json", server_config_filename) if not Path(server_config_filename).is_file(): server_config = {"attention_mode" : "auto", - "transformer_type": "t2v", + "transformer_types": [], "transformer_quantization": "int8", "text_encoder_filename" : text_encoder_choices[1], "save_path": os.path.join(os.getcwd(), "gradio_outputs"), @@ -826,7 +836,8 @@ def get_default_settings(filename): ui_defaults["num_inference_steps"] = default_number_steps return ui_defaults -transformer_type = server_config.get("transformer_type", "t2v") +transformer_types = server_config.get("transformer_types", []) +transformer_type = transformer_types[0] if len(transformer_types) > 0 else model_types[0] transformer_quantization =server_config.get("transformer_quantization", "int8") transformer_filename = get_model_filename(transformer_type, transformer_quantization) text_encoder_filename = server_config["text_encoder_filename"] @@ -1213,22 +1224,25 @@ def get_model_name(model_filename): # return header -def generate_header(compile, attention_mode): +def generate_header(model_filename, compile, attention_mode): - header = "
Attention mode: " + (attention_mode if attention_mode!="auto" else "auto/" + get_auto_attention() ) + header = "
Attention mode " + (attention_mode if attention_mode!="auto" else "auto/" + get_auto_attention() ) if attention_mode not in attention_modes_installed: header += " -NOT INSTALLED-" elif attention_mode not in attention_modes_supported: header += " -NOT SUPPORTED-" + header += "" if compile: - header += ", pytorch compilation ON" + header += ", Pytorch compilation ON" + if "int8" in model_filename: + header += ", Quantization Int8" header += "
" return header def apply_changes( state, - transformer_type_choice, + transformer_types_choices, text_encoder_choice, save_path_choice, attention_choice, @@ -1239,16 +1253,15 @@ def apply_changes( state, quantization_choice, boost_choice = 1, clear_file_list = 0, - reload_choice = 1 + reload_choice = 1, ): if args.lock_config: return if gen_in_progress: - yield "
Unable to change config when a generation is in progress
" - return + return "
Unable to change config when a generation is in progress
" global offloadobj, wan_model, server_config, loras, loras_names, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset, loras_presets server_config = {"attention_mode" : attention_choice, - "transformer_type": transformer_type_choice, + "transformer_types": transformer_types_choices, "text_encoder_filename" : text_encoder_choices[text_encoder_choice], "save_path" : save_path_choice, "compile" : compile_choice, @@ -1281,7 +1294,7 @@ def apply_changes( state, if v != v_old: changes.append(k) - global attention_mode, profile, compile, transformer_filename, text_encoder_filename, vae_config, boost, lora_dir, reload_needed, reload_model, transformer_quantization, transformer_type + global attention_mode, profile, compile, transformer_filename, text_encoder_filename, vae_config, boost, lora_dir, reload_needed, reload_model, transformer_quantization, transformer_types attention_mode = server_config["attention_mode"] profile = server_config["profile"] compile = server_config["compile"] @@ -1290,15 +1303,19 @@ def apply_changes( state, boost = server_config["boost"] reload_model = server_config["reload_model"] transformer_quantization = server_config["transformer_quantization"] - transformer_filename = get_model_filename(transformer_type, transformer_quantization) - - if all(change in ["attention_mode", "vae_config", "default_ui", "boost", "save_path", "metadata_choice", "clear_file_list"] for change in changes ): - pass + transformer_types = server_config["transformer_types"] + transformer_type = get_model_type(transformer_filename) + if not transformer_type in transformer_types: + transformer_type = transformer_types[0] if len(transformer_types) > 0 else model_types[0] + transformer_filename = get_model_filename(transformer_type, transformer_quantization) + if all(change in ["attention_mode", "vae_config", "boost", "save_path", "metadata_choice", "clear_file_list"] for change in changes ): + model_choice = gr.Dropdown() else: reload_needed = True + model_choice = generate_dropdown_model_list() - - yield "
The new configuration has been succesfully applied
" + header = generate_header(transformer_filename, compile=compile, attention_mode= attention_mode) + return "
The new configuration has been succesfully applied
", header, model_choice @@ -2505,16 +2522,11 @@ def handle_celll_selection(state, evt: gr.SelectData): def change_model(state, model_choice): - model_filename = "" - for filename in model_list: - if get_model_type(filename) == model_choice: - model_filename = filename - break - if len(model_filename) == 0: + if model_choice == None: return - + model_filename = get_model_filename(model_choice, transformer_quantization) state["model_filename"] = model_filename - header = generate_header(compile=compile, attention_mode=attention_mode) + header = generate_header(model_filename, compile=compile, attention_mode=attention_mode) return header def fill_inputs(state): @@ -3014,53 +3026,32 @@ def generate_download_tab(lset_name,loras_choices, state): download_loras_btn.click(fn=download_loras, inputs=[], outputs=[download_status_row, download_status]).then(fn=refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) -def generate_configuration_tab(): +def generate_configuration_tab(header, model_choice): state_dict = {} state = gr.State(state_dict) gr.Markdown("Please click Apply Changes at the bottom so that the changes are effective. Some choices below may be locked if the app has been launched by specifying a config preset.") with gr.Column(): - index = transformer_choices.index(transformer_filename) - index = 0 if index ==0 else index - model_list = [] + for model_type in model_types: choice = get_model_filename(model_type, transformer_quantization) model_list.append(choice) dropdown_choices = [ ( get_model_name(choice), get_model_type(choice) ) for choice in model_list] - transformer_type_choice = gr.Dropdown( + transformer_types_choices = gr.Dropdown( choices= dropdown_choices, - value= get_model_type(transformer_filename), - label= "Default Wan Transformer Model", - scale= 2 + value= transformer_types, + label= "Selectable Wan Transformer Models (keep empty to get All of them)", + scale= 2, + multiselect= True ) - # transformer_choice = gr.Dropdown( - # choices=[ - # ("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), - # ("WAN 2.1 - 480p 14B Image to Video 16 bits - the default engine in its original glory, offers a slightly better image quality but slower and requires more RAM", 4), - # ("WAN 2.1 - 480p 14B Image to Video quantized to 8 bits (recommended) - the default engine but quantized", 5), - # ("WAN 2.1 - 720p 14B Image to Video 16 bits - the default engine in its original glory, offers a slightly better image quality but slower and requires more RAM", 6), - # ("WAN 2.1 - 720p 14B Image to Video quantized to 8 bits - the default engine but quantized", 7), - # ("WAN 2.1 - Fun InP 1.3B 16 bits - the small model for fast generations with low VRAM requirements", 8), - # ("WAN 2.1 - Fun InP 14B 16 bits - Fun InP version in its original glory, offers a slightly better image quality but slower and requires more RAM", 9), - # ("WAN 2.1 - Fun InP 14B quantized to 8 bits - quantized Fun InP version", 10), - # ], - # value= index, - # label="Transformer model for Image to Video", - # interactive= not lock_ui_transformer, - # visible = True, - # ) - quantization_choice = gr.Dropdown( choices=[ ("Int8 Quantization (recommended)", "int8"), ("BF16 (no quantization)", "bf16"), ], value= transformer_quantization, - label="Wan Transformer Model Quantization (if available)", + label="Wan Transformer Model Quantization Type (if available)", ) index = text_encoder_choices.index(text_encoder_filename) @@ -3137,14 +3128,14 @@ def generate_configuration_tab(): value= profile, label="Profile (for power users only, not needed to change it)" ) - default_ui_choice = gr.Dropdown( - choices=[ - ("Text to Video", "t2v"), - ("Image to Video", "i2v"), - ], - value= default_ui, - label="Default mode when launching the App if not '--t2v' ot '--i2v' switch is specified when launching the server ", - ) + # default_ui_choice = gr.Dropdown( + # choices=[ + # ("Text to Video", "t2v"), + # ("Image to Video", "i2v"), + # ], + # value= default_ui, + # label="Default mode when launching the App if not '--t2v' ot '--i2v' switch is specified when launching the server ", + # ) metadata_choice = gr.Dropdown( choices=[ ("Export JSON files", "json"), @@ -3184,7 +3175,7 @@ def generate_configuration_tab(): fn=apply_changes, inputs=[ state, - transformer_type_choice, + transformer_types_choices, text_encoder_choice, save_path_choice, attention_choice, @@ -3197,7 +3188,7 @@ def generate_configuration_tab(): clear_file_list_choice, reload_choice, ], - outputs= msg + outputs= [msg , header, model_choice] ) def generate_about_tab(): @@ -3221,6 +3212,22 @@ def generate_info_tab(): gr.Markdown("Please note that if your turn on compilation, the first denoising step of the first video generation will be slow due to the compilation. Therefore all your tests should be done with compilation turned off.") +def generate_dropdown_model_list(): + dropdown_types= transformer_types if len(transformer_types) > 0 else model_types + current_model_type = get_model_type(transformer_filename) + if current_model_type not in dropdown_types: + dropdown_types.append(current_model_type) + model_list = [] + for model_type in dropdown_types: + choice = get_model_filename(model_type, transformer_quantization) + model_list.append(choice) + dropdown_choices = [ ( get_model_name(choice), get_model_type(choice) ) for choice in model_list] + return gr.Dropdown( + choices= dropdown_choices, + value= current_model_type, + show_label= False, + scale= 2 + ) @@ -3457,22 +3464,15 @@ def create_demo(): with gr.Tabs(selected="video_gen", ) as main_tabs: with gr.Tab("Video Generator", id="video_gen") as t2v_tab: with gr.Row(): - header = gr.Markdown(generate_header(compile, attention_mode), visible= True) + if args.lock_model: + gr.Markdown("

" + get_model_name(transformer_filename) + "

") + model_choice = gr.Dropdown(visible=False, value= get_model_type(transformer_filename)) + else: + gr.Markdown("
") + model_choice = generate_dropdown_model_list() + gr.Markdown("
") with gr.Row(): - gr.Markdown("
") - - model_list = [] - for model_type in model_types: - choice = get_model_filename(model_type, transformer_quantization) - model_list.append(choice) - dropdown_choices = [ ( get_model_name(choice), get_model_type(choice) ) for choice in model_list] - model_choice = gr.Dropdown( - choices= dropdown_choices, - value= get_model_type(transformer_filename), - show_label= False, - scale= 2 - ) - gr.Markdown("
") + header = gr.Markdown(generate_header(transformer_filename, compile, attention_mode), visible= True) with gr.Row(): loras_choices, lset_name, state = generate_video_tab(model_choice = model_choice, header = header) @@ -3482,7 +3482,7 @@ def create_demo(): with gr.Tab("Downloads", id="downloads") as downloads_tab: generate_download_tab(lset_name, loras_choices, state) with gr.Tab("Configuration"): - generate_configuration_tab() + generate_configuration_tab(header, model_choice) with gr.Tab("About"): generate_about_tab() From fea835f21f86ee33c414ce6b3562541e90b886a9 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Tue, 8 Apr 2025 01:14:47 +0200 Subject: [PATCH 39/69] Polishing and added new PreLoad Type --- wgp.py | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/wgp.py b/wgp.py index db02e07..083ee1e 100644 --- a/wgp.py +++ b/wgp.py @@ -96,6 +96,7 @@ def process_prompt_and_add_tasks(state, model_choice): inputs = state.get(get_model_type(model_filename), None) inputs["state"] = state + inputs.pop("lset_name") if inputs == None: return prompt = inputs["prompt"] @@ -1165,7 +1166,7 @@ def load_models(model_filename): return wan_model, offloadobj, pipe["transformer"] -if reload_model ==3: +if reload_model ==3 or reload_model ==4: wan_model, offloadobj, transformer = None, None, None reload_needed = True else: @@ -1550,7 +1551,7 @@ def generate_video( # gr.Info("Unable to generate a Video while a new configuration is being applied.") # return - if reload_model !=3 : + if reload_model !=3 and reload_model !=4 : while wan_model == None: time.sleep(1) @@ -2352,7 +2353,6 @@ def prepare_inputs_dict(target, inputs ): if target == "state": return inputs - unsaved_params = ["image_start", "image_end", "image_refs", "video_guide", "video_mask"] for k in unsaved_params: inputs.pop(k) @@ -2387,6 +2387,7 @@ def get_function_arguments(func, locals): def save_inputs( target, + lset_name, prompt, negative_prompt, resolution, @@ -2557,7 +2558,7 @@ def preload_model(state): def unload_model_if_needed(state): global reload_needed, wan_model, offloadobj - if reload_model == 3: + if reload_model == 4: if wan_model != None: wan_model = None if offloadobj is not None: @@ -2605,6 +2606,8 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non launch_loras = default_loras_choices launch_multis_str = default_loras_multis_str + if len(launch_preset) == 0: + launch_preset = ui_defaults.get("lset_name","") if len(launch_prompt) == 0: launch_prompt = ui_defaults.get("prompt","") if len(launch_loras) == 0: @@ -2673,18 +2676,19 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non image_end = gr.Image(label= "Last Image for a new video", type ="pil", visible="E" in image_prompt_type_value, value= ui_defaults.get("image_end", None)) with gr.Column(visible= "Vace" in model_filename ) as video_prompt_column: + gr.Markdown("Control conditions: Images References (custom Faces or Objects), Video (Open Pose, Depth maps), Mask (inpainting)") video_prompt_type_value= ui_defaults.get("video_prompt_type","I") - video_prompt_type = gr.Radio( [("Use Images Ref", "I"),("a Video", "V"), ("Images + a Video", "IV"), ("Video + Video Mask", "VM"), ("Images + Video + Mask", "IVM")], value =video_prompt_type_value, label="Location", show_label= False, scale= 3) + video_prompt_type = gr.Radio( [("Images Ref", "I"),("a Video", "V"), ("Images Refs + a Video", "IV"), ("Video + Video Mask", "VM"), ("Images + Video + Mask", "IVM")], value =video_prompt_type_value, label="Location", show_label= False, scale= 3) image_refs = gr.Gallery( - label="Reference Images of Faces and / or Object to be found in the Video", type ="pil", + label="Images Referencse (Custom faces and Objects 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 video_prompt_type_value, value= ui_defaults.get("image_refs", None) ) - video_guide = gr.Video(label= "Reference Video", visible= "V" in video_prompt_type_value, value= ui_defaults.get("video_guide", None) ) + video_guide = gr.Video(label= "Reference Video (an animated Video in the Open Pose format or Depth Map video)", visible= "V" in video_prompt_type_value, value= ui_defaults.get("video_guide", None) ) with gr.Row(): max_frames = gr.Slider(0, 100, value=ui_defaults.get("max_frames",0), step=1, label="Nb of frames in Ref. Video (0 = as many as possible)", visible= "V" in video_prompt_type_value, 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 video_prompt_type_value, scale =1 ) - video_mask = gr.Video(label= "Video Mask (white pixels = Mask)", visible= "M" in video_prompt_type_value, value= ui_defaults.get("video_mask", None) ) + video_mask = gr.Video(label= "Video Mask (for Inpainting or Outpaing, white pixels = Mask)", visible= "M" in video_prompt_type_value, value= ui_defaults.get("video_mask", None) ) advanced_prompt = advanced_ui @@ -3147,9 +3151,10 @@ def generate_configuration_tab(header, model_choice): ) reload_choice = gr.Dropdown( choices=[ - ("Load Model When Changing Model", 1), - ("Load Model When Pressing Generate", 2), - ("Load Model When Pressing Generate and Unload Model when Finished", 3), + ("Load Model When Starting the App and Changing Model if Model Changed", 1), + ("Load Model When Starting the App and Pressing Generate if Model Changed", 2), + ("Load Model When Pressing Generate if Model Changed", 3), + ("Load Model When Pressing Generate and Unload Model when Finished", 4), ], value=server_config.get("reload_model",2), label="RAM Loading / Unloading Model Policy (in any case VRAM will be freed once the queue has been processed)" From a15751892ed01a682f155c2bd242a2ee104ff2a3 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 8 Apr 2025 20:24:36 +1000 Subject: [PATCH 40/69] add queue saving/loading/clearing/autosaving/autoloading, fix empty prompt logic --- wgp.py | 593 +++++++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 581 insertions(+), 12 deletions(-) diff --git a/wgp.py b/wgp.py index 083ee1e..414cc80 100644 --- a/wgp.py +++ b/wgp.py @@ -28,6 +28,12 @@ from wan.utils import prompt_parser import base64 import io from PIL import Image +import zipfile +import tempfile +import shutil +import atexit +global_queue_ref = [] +AUTOSAVE_FILENAME = "queue.zip" PROMPT_VARS_MAX = 10 target_mmgp_version = "3.3.4" @@ -98,10 +104,14 @@ def process_prompt_and_add_tasks(state, model_choice): inputs["state"] = state inputs.pop("lset_name") if inputs == None: - return + gr.Warning("Internal state error: Could not retrieve inputs for the model.") + return update_queue_data(queue) prompt = inputs["prompt"] if len(prompt) ==0: - return + gr.Info("Prompt cannot be empty.") + gen = get_gen_info(state) + queue = gen.get("queue", []) + return get_queue_table(queue) prompt, errors = prompt_parser.process_template(prompt) if len(errors) > 0: gr.Info("Error processing prompt template: " + errors) @@ -111,7 +121,10 @@ def process_prompt_and_add_tasks(state, model_choice): prompts = prompt.replace("\r", "").split("\n") prompts = [prompt.strip() for prompt in prompts if len(prompt.strip())>0 and not prompt.startswith("#")] if len(prompts) ==0: - return + gr.Info("Prompt cannot be empty.") + gen = get_gen_info(state) + queue = gen.get("queue", []) + return get_queue_table(queue) resolution = inputs["resolution"] width, height = resolution.split("x") @@ -250,9 +263,6 @@ def process_prompt_and_add_tasks(state, model_choice): queue= gen.get("queue", []) return update_queue_data(queue) - - - def add_video_task(**inputs): global task_id state = inputs["state"] @@ -327,6 +337,444 @@ def remove_task(queue, selected_indices): del queue[idx] return update_queue_data(queue) +def update_global_queue_ref(queue): + global global_queue_ref + with lock: + global_queue_ref = queue[:] + +def save_queue_action(state): + gen = get_gen_info(state) + queue = gen.get("queue", []) + + if not queue or len(queue) <=1 : # Check if queue is empty or only has the placeholder + gr.Info("Queue is empty. Nothing to save.") + return None # Return None if nothing to save + + # Use an in-memory buffer for the zip file + zip_buffer = io.BytesIO() + + # Still use a temporary directory *only* for storing images before zipping + with tempfile.TemporaryDirectory() as tmpdir: + queue_manifest = [] + image_paths_in_zip = {} # Tracks image PIL object ID -> filename in zip + + for task_index, task in enumerate(queue): + # Skip the placeholder item if it exists + if task is None or not isinstance(task, dict) or task_index == 0: continue + + params_copy = task.get('params', {}).copy() + task_id_s = task.get('id', f"task_{task_index}") # Use a different var name + + image_keys = ["image_start", "image_end", "image_refs"] + for key in image_keys: + images_pil = params_copy.get(key) + if images_pil is None: + continue + + # Ensure images_pil is always a list for processing + is_originally_list = isinstance(images_pil, list) + if not is_originally_list: + images_pil = [images_pil] + + image_filenames_for_json = [] + for img_index, pil_image in enumerate(images_pil): + # Ensure it's actually a PIL Image object before proceeding + if not isinstance(pil_image, Image.Image): + print(f"Warning: Expected PIL Image for key '{key}' in task {task_id_s}, got {type(pil_image)}. Skipping image.") + continue + + # Use object ID to check if this specific image instance is already saved + img_id = id(pil_image) + if img_id in image_paths_in_zip: + # If already saved, just add its filename to the list + image_filenames_for_json.append(image_paths_in_zip[img_id]) + continue # Move to the next image in the list + + # Image not saved yet, create filename and save path + img_filename_in_zip = f"task{task_id_s}_{key}_{img_index}.png" + img_save_path = os.path.join(tmpdir, img_filename_in_zip) + + try: + # Save the image to the temporary directory + pil_image.save(img_save_path, "PNG") + image_filenames_for_json.append(img_filename_in_zip) + # Store the mapping from image ID to its filename in the zip + image_paths_in_zip[img_id] = img_filename_in_zip + except Exception as e: + print(f"Error saving image {img_filename_in_zip} for task {task_id_s}: {e}") + # Optionally decide if you want to continue or fail here + + # Update the params_copy with the list of filenames (or single filename) + if image_filenames_for_json: + params_copy[key] = image_filenames_for_json if is_originally_list else image_filenames_for_json[0] + else: + # If no images were successfully processed for this key, remove it + params_copy.pop(key, None) + + + # Clean up parameters before adding to manifest + params_copy.pop('state', None) + params_copy.pop('start_image_data_base64', None) # Don't need base64 in saved queue + params_copy.pop('end_image_data_base64', None) + # Also remove the actual PIL data if it somehow remained + params_copy.pop('start_image_data', None) + params_copy.pop('end_image_data', None) + + manifest_entry = { + "id": task.get('id'), + "params": params_copy, + # Keep other necessary top-level task info if needed, like repeats etc. + # Example: "repeats": task.get('repeats', 1) + } + queue_manifest.append(manifest_entry) + + # --- Create queue.json content --- + manifest_path = os.path.join(tmpdir, "queue.json") + try: + with open(manifest_path, 'w', encoding='utf-8') as f: + # Dump only the relevant manifest data + json.dump(queue_manifest, f, indent=4) + except Exception as e: + print(f"Error writing queue.json: {e}") + gr.Warning("Failed to create queue manifest.") + return None # Return None on failure + + # --- Create the zip file in memory --- + try: + with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zf: + # Add queue.json + zf.write(manifest_path, arcname="queue.json") + + # Add all unique images that were saved to the temp dir + for saved_img_rel_path in image_paths_in_zip.values(): + saved_img_abs_path = os.path.join(tmpdir, saved_img_rel_path) + if os.path.exists(saved_img_abs_path): + zf.write(saved_img_abs_path, arcname=saved_img_rel_path) + else: + # This shouldn't happen if saving was successful, but good to check + print(f"Warning: Image file {saved_img_rel_path} not found during zipping.") + + # --- Prepare for return --- + # Move buffer position to the beginning + zip_buffer.seek(0) + # Read the binary content + zip_binary_content = zip_buffer.getvalue() + # Encode as base64 string + zip_base64 = base64.b64encode(zip_binary_content).decode('utf-8') + print(f"Queue successfully prepared as base64 string ({len(zip_base64)} chars).") + return zip_base64 + + except Exception as e: + print(f"Error creating zip file in memory: {e}") + gr.Warning("Failed to create zip data for download.") + return None # Return None on failure + finally: + zip_buffer.close() + +def load_queue_action(filepath, state): + global task_id + gen = get_gen_info(state) + original_queue = gen.get("queue", []) # Store original queue for error case + + if not filepath or not hasattr(filepath, 'name') or not Path(filepath.name).is_file(): + print("[load_queue_action] Warning: No valid file selected or file not found.") + # Return the current state of the DataFrame + return update_queue_data(original_queue) + + newly_loaded_queue = [] + max_id_in_file = 0 + error_message = "" + local_queue_copy_for_global_ref = None + + try: + print(f"[load_queue_action] Attempting to load queue from: {filepath.name}") + with tempfile.TemporaryDirectory() as tmpdir: + with zipfile.ZipFile(filepath.name, 'r') as zf: + if "queue.json" not in zf.namelist(): raise ValueError("queue.json not found in zip file") + print(f"[load_queue_action] Extracting {filepath.name} to {tmpdir}") + zf.extractall(tmpdir) + print(f"[load_queue_action] Extraction complete.") + + manifest_path = os.path.join(tmpdir, "queue.json") + print(f"[load_queue_action] Reading manifest: {manifest_path}") + with open(manifest_path, 'r', encoding='utf-8') as f: + loaded_manifest = json.load(f) + print(f"[load_queue_action] Manifest loaded. Processing {len(loaded_manifest)} tasks.") + + for task_index, task_data in enumerate(loaded_manifest): + # (Keep the existing task processing logic here...) + if task_data is None or not isinstance(task_data, dict): + print(f"[load_queue_action] Skipping invalid task data at index {task_index}") + continue + + params = task_data.get('params', {}) + task_id_loaded = task_data.get('id', 0) + max_id_in_file = max(max_id_in_file, task_id_loaded) + loaded_pil_images = {} + image_keys = ["image_start", "image_end", "image_refs"] + params['state'] = state # Add state back temporarily for consistency if needed by internal logic, but it's removed before saving + + for key in image_keys: + image_filenames = params.get(key) + if image_filenames is None: continue + is_list = isinstance(image_filenames, list) + if not is_list: image_filenames = [image_filenames] + loaded_pils = [] + for img_filename_in_zip in image_filenames: + if not isinstance(img_filename_in_zip, str): continue + img_load_path = os.path.join(tmpdir, img_filename_in_zip) + if not os.path.exists(img_load_path): + print(f"[load_queue_action] Image file not found during load: {img_load_path}") + continue + try: + pil_image = Image.open(img_load_path) + # Ensure the image data is loaded into memory before the temp dir is cleaned up + pil_image.load() + # Convert image right after loading + converted_image = convert_image(pil_image) + loaded_pils.append(converted_image) + pil_image.close() # Close the file handle + except Exception as img_e: + print(f"[load_queue_action] Error loading image {img_filename_in_zip}: {img_e}") + if loaded_pils: + params[key] = loaded_pils if is_list else loaded_pils[0] + loaded_pil_images[key] = params[key] # Store loaded PILs for preview generation + else: params.pop(key, None) + + # Generate preview base64 strings + primary_preview_pil, secondary_preview_pil = None, None + start_prev_pil_list = loaded_pil_images.get("image_start") + end_prev_pil_list = loaded_pil_images.get("image_end") + ref_prev_pil_list = loaded_pil_images.get("image_refs") + + # Extract first image for preview if available + if start_prev_pil_list: + primary_preview_pil = start_prev_pil_list[0] if isinstance(start_prev_pil_list, list) and start_prev_pil_list else start_prev_pil_list if not isinstance(start_prev_pil_list, list) else None + if end_prev_pil_list: + secondary_preview_pil = end_prev_pil_list[0] if isinstance(end_prev_pil_list, list) and end_prev_pil_list else end_prev_pil_list if not isinstance(end_prev_pil_list, list) else None + elif ref_prev_pil_list and isinstance(ref_prev_pil_list, list) and ref_prev_pil_list: + primary_preview_pil = ref_prev_pil_list[0] + + # Generate base64 only if PIL image exists + start_b64 = [pil_to_base64_uri(primary_preview_pil, format="jpeg", quality=70)] if primary_preview_pil else None + end_b64 = [pil_to_base64_uri(secondary_preview_pil, format="jpeg", quality=70)] if secondary_preview_pil else None + + # Get top-level image data (PIL objects) for runtime task + top_level_start_image = loaded_pil_images.get("image_start") + top_level_end_image = loaded_pil_images.get("image_end") + + # Construct the runtime task dictionary + runtime_task = { + "id": task_id_loaded, + "params": params.copy(), # Use a copy of params + # Extract necessary params for top level if they exist + "repeats": params.get('repeat_generation', 1), + "length": params.get('video_length'), + "steps": params.get('num_inference_steps'), + "prompt": params.get('prompt'), + # Store the actual loaded PIL image data here + "start_image_data": top_level_start_image, + "end_image_data": top_level_end_image, + # Store base64 previews generated above + "start_image_data_base64": start_b64, + "end_image_data_base64": end_b64, + } + newly_loaded_queue.append(runtime_task) + print(f"[load_queue_action] Processed task {task_index+1}/{len(loaded_manifest)}, ID: {task_id_loaded}") + + # --- State Update --- + with lock: + print("[load_queue_action] Acquiring lock to update state...") + gen["queue"] = newly_loaded_queue[:] # Replace the queue in the state + local_queue_copy_for_global_ref = gen["queue"][:] # Copy for global ref update + current_max_id_in_new_queue = max([t['id'] for t in newly_loaded_queue if 'id' in t] + [0]) # Safer max ID calculation + + # Update global task ID only if the loaded max ID is higher + if current_max_id_in_new_queue > task_id: + print(f"[load_queue_action] Updating global task_id from {task_id} to {current_max_id_in_new_queue + 1}") + task_id = current_max_id_in_new_queue + 1 # Ensure next ID is unique + else: + print(f"[load_queue_action] Global task_id ({task_id}) is >= max in file ({current_max_id_in_new_queue}). Not changing task_id.") + + gen["prompts_max"] = len(newly_loaded_queue) + print("[load_queue_action] State update complete. Releasing lock.") + + # --- Global Reference Update --- + if local_queue_copy_for_global_ref is not None: + print("[load_queue_action] Updating global queue reference...") + update_global_queue_ref(local_queue_copy_for_global_ref) + else: + # This case should ideally not be reached if state update happens + print("[load_queue_action] Warning: Skipping global ref update as local copy is None.") + + print(f"[load_queue_action] Queue load successful. Returning DataFrame update for {len(newly_loaded_queue)} tasks.") + # *** Return the DataFrame update object *** + return update_queue_data(newly_loaded_queue) + + except (ValueError, zipfile.BadZipFile, FileNotFoundError, Exception) as e: + error_message = f"Error during queue load: {e}" + print(f"[load_queue_action] Caught error: {error_message}") + traceback.print_exc() + # Optionally show a Gradio warning/error to the user + gr.Warning(f"Failed to load queue: {error_message[:200]}") # Show truncated error + + # *** Return the DataFrame update for the original queue *** + print("[load_queue_action] Load failed. Returning DataFrame update for original queue.") + return update_queue_data(original_queue) + finally: + # Clean up the uploaded file object if it exists and has a path + if filepath and hasattr(filepath, 'name') and filepath.name and os.path.exists(filepath.name): + try: + # Gradio often uses temp files, attempting removal is good practice + # os.remove(filepath.name) + # print(f"[load_queue_action] Cleaned up temporary upload file: {filepath.name}") + pass # Let Gradio manage its temp files unless specifically needed + except OSError as e: + # Ignore errors like "file not found" if already cleaned up + print(f"[load_queue_action] Info: Could not remove temp file {filepath.name}: {e}") + pass + +def clear_queue_action(state): + gen = get_gen_info(state) + queue = gen.get("queue", []) + if not queue: + gr.Info("Queue is already empty.") + return update_queue_data([]) + + with lock: + queue.clear() + gen["prompts_max"] = 0 + + gr.Info("Queue cleared.") + return update_queue_data([]) + +def autosave_queue(): + global global_queue_ref + if not global_queue_ref: + print("Autosave: Queue is empty, nothing to save.") + return + + print(f"Autosaving queue ({len(global_queue_ref)} items) to {AUTOSAVE_FILENAME}...") + temp_state_for_save = {"gen": {"queue": global_queue_ref}} + zip_file_path = None + try: + + def _save_queue_to_file(queue_to_save, output_filename): + if not queue_to_save: return None + with tempfile.TemporaryDirectory() as tmpdir: + queue_manifest = [] + image_paths_in_zip = {} + for task_index, task in enumerate(queue_to_save): + if task is None or not isinstance(task, dict): continue + params_copy = task.get('params', {}).copy() + task_id_s = task.get('id', f"task_{task_index}") + image_keys = ["image_start", "image_end", "image_refs"] + for key in image_keys: + images_pil = params_copy.get(key) + if images_pil is None: continue + is_list = isinstance(images_pil, list) + if not is_list: images_pil = [images_pil] + image_filenames_for_json = [] + for img_index, pil_image in enumerate(images_pil): + if not isinstance(pil_image, Image.Image): continue + img_id = id(pil_image) + if img_id in image_paths_in_zip: + image_filenames_for_json.append(image_paths_in_zip[img_id]) + continue + img_filename_in_zip = f"task{task_id_s}_{key}_{img_index}.png" + img_save_path = os.path.join(tmpdir, img_filename_in_zip) + try: + pil_image.save(img_save_path, "PNG") + image_filenames_for_json.append(img_filename_in_zip) + image_paths_in_zip[img_id] = img_filename_in_zip + except Exception as e: + print(f"Autosave error saving image {img_filename_in_zip}: {e}") + if image_filenames_for_json: + params_copy[key] = image_filenames_for_json if is_list else image_filenames_for_json[0] + else: + params_copy.pop(key, None) + params_copy.pop('state', None) + params_copy.pop('start_image_data_base64', None) + params_copy.pop('end_image_data_base64', None) + manifest_entry = { + "id": task.get('id'), "params": params_copy, + } + queue_manifest.append(manifest_entry) + manifest_path = os.path.join(tmpdir, "queue.json") + with open(manifest_path, 'w', encoding='utf-8') as f: json.dump(queue_manifest, f, indent=4) + with zipfile.ZipFile(output_filename, 'w', zipfile.ZIP_DEFLATED) as zf: + zf.write(manifest_path, arcname="queue.json") + for saved_img_rel_path in image_paths_in_zip.values(): + saved_img_abs_path = os.path.join(tmpdir, saved_img_rel_path) + if os.path.exists(saved_img_abs_path): + zf.write(saved_img_abs_path, arcname=saved_img_rel_path) + return output_filename + return None # Should not happen if queue has items + + saved_path = _save_queue_to_file(global_queue_ref, AUTOSAVE_FILENAME) + + if saved_path: + print(f"Queue autosaved successfully to {saved_path}") + else: + print("Autosave failed.") + except Exception as e: + print(f"Error during autosave: {e}") + traceback.print_exc() + + +def autoload_queue(state): + global task_id + # Initial check using the original state + try: + gen = get_gen_info(state) # Make sure initial state is a dict + original_queue = gen.get("queue", []) + except AttributeError: + print("[autoload_queue] Error: Initial state is not a dictionary. Cannot autoload.") + # Return default values indicating no load occurred and the state is unchanged + return gr.update(visible=False), False, state # Return an empty DF update + + loaded_flag = False + dataframe_update = update_queue_data(original_queue) # Default update is the original queue + + if not original_queue and Path(AUTOSAVE_FILENAME).is_file(): + print(f"Autoloading queue from {AUTOSAVE_FILENAME}...") + class MockFile: + def __init__(self, name): + self.name = name + mock_filepath = MockFile(AUTOSAVE_FILENAME) + + # Call load_queue_action, it modifies 'state' internally and returns a DataFrame update + dataframe_update = load_queue_action(mock_filepath, state) + + # Now check the 'state' dictionary which should have been modified by load_queue_action + gen = get_gen_info(state) # Use the (potentially) modified state dictionary + loaded_queue_after_action = gen.get("queue", []) + + if loaded_queue_after_action: # Check if the queue in the state is now populated + print(f"Autoload successful. Loaded {len(loaded_queue_after_action)} tasks into state.") + loaded_flag = True + # Global ref update was already done inside load_queue_action if successful + else: + print("Autoload attempted but queue in state remains empty (file might be empty or invalid).") + # Ensure state reflects empty queue if load failed but file existed + with lock: + gen["queue"] = [] + gen["prompts_max"] = 0 + update_global_queue_ref([]) + dataframe_update = update_queue_data([]) # Ensure UI shows empty queue + + else: # Handle cases where autoload shouldn't happen + if original_queue: + print("Autoload skipped: Queue is not empty.") + update_global_queue_ref(original_queue) # Ensure global ref matches current state + dataframe_update = update_queue_data(original_queue) # UI should show current queue + else: + print(f"Autoload skipped: {AUTOSAVE_FILENAME} not found.") + update_global_queue_ref([]) # Ensure global ref is empty + dataframe_update = update_queue_data([]) # UI should show empty queue + + # Return the DataFrame update needed for the UI, the flag, and the final state dictionary + return dataframe_update, loaded_flag, state def get_queue_table(queue): @@ -390,7 +838,7 @@ def get_queue_table(queue): ]) return data def update_queue_data(queue): - + update_global_queue_ref(queue) data = get_queue_table(queue) # if len(data) == 0: @@ -1993,6 +2441,7 @@ def process_tasks(state, progress=gr.Progress()): yield status queue[:] = [item for item in queue if item['id'] != task['id']] + update_global_queue_ref(queue) gen["prompts_max"] = 0 gen["prompt"] = "" @@ -2716,7 +3165,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non wizard_variables = "\n".join(variables) for _ in range( PROMPT_VARS_MAX - len(prompt_vars)): prompt_vars.append(gr.Textbox(visible= False, min_width=80, show_label= False)) - + with gr.Column(not advanced_prompt) as prompt_column_wizard: 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) @@ -2902,7 +3351,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non queue_df = gr.DataFrame( 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"], + column_widths= ["5%", None, "7%", "7%", "10%", "10%", "3%", "3%", "3%"], interactive=False, col_count=(9, "fixed"), wrap=True, @@ -2911,6 +3360,72 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non visible= False, elem_id="queue_df" ) + with gr.Row(): + queue_zip_base64_output = gr.Text(visible=False) + save_queue_btn = gr.DownloadButton("Save Queue", size="sm") + load_queue_btn = gr.UploadButton("Load Queue", file_types=[".zip"], size="sm") + clear_queue_btn = gr.Button("Clear Queue", size="sm", variant="stop") + trigger_zip_download_js = """ + (base64String) => { + if (!base64String) { + console.log("No base64 zip data received, skipping download."); + return; + } + try { + const byteCharacters = atob(base64String); + const byteNumbers = new Array(byteCharacters.length); + for (let i = 0; i < byteCharacters.length; i++) { + byteNumbers[i] = byteCharacters.charCodeAt(i); + } + const byteArray = new Uint8Array(byteNumbers); + const blob = new Blob([byteArray], { type: 'application/zip' }); + + const url = URL.createObjectURL(blob); + const a = document.createElement('a'); + a.style.display = 'none'; + a.href = url; + a.download = 'queue.zip'; + document.body.appendChild(a); + a.click(); + + window.URL.revokeObjectURL(url); + document.body.removeChild(a); + console.log("Zip download triggered."); + } catch (e) { + console.error("Error processing base64 data or triggering download:", e); + } + } + """ + save_queue_btn.click( + fn=save_queue_action, + inputs=[state], + outputs=[queue_zip_base64_output] + ).then( + fn=None, + inputs=[queue_zip_base64_output], + outputs=None, + js=trigger_zip_download_js + ) + + load_queue_btn.upload( + fn=load_queue_action, + inputs=[load_queue_btn, state], + outputs=[queue_df] + ).then( + fn=lambda s: gr.update(visible=bool(get_gen_info(s).get("queue",[]))), + inputs=[state], + outputs=[current_gen_column] + ) + + clear_queue_btn.click( + fn=clear_queue_action, + inputs=[state], + outputs=[queue_df] + ).then( + fn=lambda: gr.update(visible=False), + inputs=None, + outputs=[current_gen_column] + ) extra_inputs = prompt_vars + [wizard_prompt, wizard_variables_var, wizard_prompt_activated_var, video_prompt_column, image_prompt_column, prompt_column_advanced, prompt_column_wizard_vars, prompt_column_wizard, lset_name, advanced_row] # show_advanced presets_column, @@ -3014,7 +3529,16 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non outputs=[modal_container] ) - return loras_choices, lset_name, state + return ( + loras_choices, lset_name, state, queue_df, current_gen_column, + gen_status, output, abort_btn, generate_btn, add_to_queue_btn, + gen_info, + prompt, wizard_prompt, wizard_prompt_activated_var, wizard_variables_var, + prompt_column_advanced, prompt_column_wizard, prompt_column_wizard_vars, + advanced_row, image_prompt_column, video_prompt_column, + *prompt_vars + ) + def generate_download_tab(lset_name,loras_choices, state): with gr.Row(): @@ -3479,8 +4003,15 @@ def create_demo(): with gr.Row(): header = gr.Markdown(generate_header(transformer_filename, compile, attention_mode), visible= True) with gr.Row(): - - loras_choices, lset_name, state = generate_video_tab(model_choice = model_choice, header = header) + ( + loras_choices, lset_name, state, queue_df, current_gen_column, + gen_status, output, abort_btn, generate_btn, add_to_queue_btn, + gen_info, + prompt, wizard_prompt, wizard_prompt_activated_var, wizard_variables_var, + prompt_column_advanced, prompt_column_wizard, prompt_column_wizard_vars, + advanced_row, image_prompt_column, video_prompt_column, + *prompt_vars_outputs + ) = generate_video_tab(model_choice=model_choice, header=header) with gr.Tab("Informations"): generate_info_tab() if not args.lock_config: @@ -3491,9 +4022,47 @@ def create_demo(): with gr.Tab("About"): generate_about_tab() + should_start_flag = gr.State(False) + def run_autoload_and_prepare_ui(current_state): + df_update, loaded_flag, modified_state = autoload_queue(current_state) + should_start_processing = loaded_flag + return df_update, gr.update(visible=loaded_flag), should_start_processing, modified_state + + def start_processing_if_needed(should_start, current_state): + if not isinstance(current_state, dict) or 'gen' not in current_state: + yield "Error: Invalid state received before processing." + return + if should_start: + yield from process_tasks(current_state) + else: + yield "Autoload complete. Processing not started." + + def finalize_generation_with_state(current_state): + if not isinstance(current_state, dict) or 'gen' not in current_state: + return gr.update(), gr.update(interactive=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False, value=""), current_state + gallery_update, abort_btn_update, gen_btn_update, add_queue_btn_update, current_gen_col_update, gen_info_update = finalize_generation(current_state) + return gallery_update, abort_btn_update, gen_btn_update, add_queue_btn_update, current_gen_col_update, gen_info_update, current_state + + demo.load( + fn=run_autoload_and_prepare_ui, + inputs=[state], + outputs=[queue_df, current_gen_column, should_start_flag, state] + ).then( + fn=start_processing_if_needed, + inputs=[should_start_flag, state], + outputs=[gen_status], + trigger_mode="once" + ).then( + fn=finalize_generation_with_state, + inputs=[state], + outputs=[output, abort_btn, generate_btn, add_to_queue_btn, current_gen_column, gen_info, state], + trigger_mode="always_last" + ) + return demo if __name__ == "__main__": + atexit.register(autosave_queue) # threading.Thread(target=runner, daemon=True).start() os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" server_port = int(args.server_port) From 1cad8d429ec18178c8abf502ef083ab45ea213ed Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 8 Apr 2025 21:51:36 +1000 Subject: [PATCH 41/69] fix no inputs warning --- wgp.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/wgp.py b/wgp.py index 414cc80..3068ff1 100644 --- a/wgp.py +++ b/wgp.py @@ -30,7 +30,6 @@ import io from PIL import Image import zipfile import tempfile -import shutil import atexit global_queue_ref = [] AUTOSAVE_FILENAME = "queue.zip" @@ -105,7 +104,7 @@ def process_prompt_and_add_tasks(state, model_choice): inputs.pop("lset_name") if inputs == None: gr.Warning("Internal state error: Could not retrieve inputs for the model.") - return update_queue_data(queue) + return get_queue_table(queue) prompt = inputs["prompt"] if len(prompt) ==0: gr.Info("Prompt cannot be empty.") From 6f4714ea198f1d2b5f08da12bd45f9f8d4eeba57 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Tue, 8 Apr 2025 21:52:25 +1000 Subject: [PATCH 42/69] fix no inputs warning (2) --- wgp.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/wgp.py b/wgp.py index 3068ff1..0f3b8e8 100644 --- a/wgp.py +++ b/wgp.py @@ -104,6 +104,8 @@ def process_prompt_and_add_tasks(state, model_choice): inputs.pop("lset_name") if inputs == None: gr.Warning("Internal state error: Could not retrieve inputs for the model.") + gen = get_gen_info(state) + queue = gen.get("queue", []) return get_queue_table(queue) prompt = inputs["prompt"] if len(prompt) ==0: From 9ac1674615fe218ef45951984d06594f66e760bb Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Wed, 9 Apr 2025 15:51:23 +0200 Subject: [PATCH 43/69] Simplified Vace, added auto open pose and depth extrators --- README.md | 17 +- preprocessing/dwpose/__init__.py | 2 + preprocessing/dwpose/onnxdet.py | 127 ++++++ preprocessing/dwpose/onnxpose.py | 362 +++++++++++++++++ preprocessing/dwpose/pose.py | 183 +++++++++ preprocessing/dwpose/util.py | 299 ++++++++++++++ preprocessing/dwpose/wholebody.py | 80 ++++ preprocessing/gray.py | 35 ++ preprocessing/midas/__init__.py | 2 + preprocessing/midas/api.py | 166 ++++++++ preprocessing/midas/base_model.py | 18 + preprocessing/midas/blocks.py | 391 ++++++++++++++++++ preprocessing/midas/depth.py | 84 ++++ preprocessing/midas/dpt_depth.py | 107 +++++ preprocessing/midas/midas_net.py | 80 ++++ preprocessing/midas/midas_net_custom.py | 167 ++++++++ preprocessing/midas/transforms.py | 231 +++++++++++ preprocessing/midas/utils.py | 193 +++++++++ preprocessing/midas/vit.py | 510 ++++++++++++++++++++++++ wan/text2video.py | 8 +- wan/utils/utils.py | 24 ++ wan/utils/vace_preprocessor.py | 71 ++-- wgp.py | 263 +++++++++--- 23 files changed, 3316 insertions(+), 104 deletions(-) create mode 100644 preprocessing/dwpose/__init__.py create mode 100644 preprocessing/dwpose/onnxdet.py create mode 100644 preprocessing/dwpose/onnxpose.py create mode 100644 preprocessing/dwpose/pose.py create mode 100644 preprocessing/dwpose/util.py create mode 100644 preprocessing/dwpose/wholebody.py create mode 100644 preprocessing/gray.py create mode 100644 preprocessing/midas/__init__.py create mode 100644 preprocessing/midas/api.py create mode 100644 preprocessing/midas/base_model.py create mode 100644 preprocessing/midas/blocks.py create mode 100644 preprocessing/midas/depth.py create mode 100644 preprocessing/midas/dpt_depth.py create mode 100644 preprocessing/midas/midas_net.py create mode 100644 preprocessing/midas/midas_net_custom.py create mode 100644 preprocessing/midas/transforms.py create mode 100644 preprocessing/midas/utils.py create mode 100644 preprocessing/midas/vit.py diff --git a/README.md b/README.md index c0380c0..8531620 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ ## 🔥 Latest News!! -* April 4 2025: 👋 Wan 2.1GP v4.0: lots of goodies for you ! +* April 9 2025: 👋 Wan 2.1GP v4.0: lots of goodies for you ! - A new queuing system that lets you stack in a queue as many text2video and imag2video tasks as you want. Each task can rely on complete different generation parameters (different number of frames, steps, loras, ...). - Temporal upsampling (Rife) and spatial upsampling (Lanczos) for a smoother video (32 fps or 64 fps) and to enlarge you video by x2 or x4. Check these new advanced options. - Wan Vace Control Net support : with Vace you can inject in the scene people or objects, animate a person, perform inpainting or outpainting, continue a video, ... I have provided an introduction guide below. @@ -272,23 +272,24 @@ You can define multiple lines of macros. If there is only one macro line, the ap ### VACE ControlNet introduction -Vace is a ControlNet 1.3B text2video model that allows you on top of a text prompt to provide visual hints to guide the generation. It can do more things than image2video although it is not as good for just starting a video with an image because it only a 1.3B model (in fact 3B) versus 14B and (it is not specialized for start frames). However, with Vace you can inject in the scene people or objects, animate a person, perform inpainting or outpainting, continue a video, ... +Vace is a ControlNet 1.3B text2video model that allows you to do Video to Video and Reference to Video (inject your own images into the output video). So with Vace you can inject in the scene people or objects of your choice, animate a person, perform inpainting or outpainting, continue a video, ... First you need to select the Vace 1.3B model in the Drop Down box at the top. Please note that Vace works well for the moment only with videos up to 5s (81 frames). Beside the usual Text Prompt, three new types of visual hints can be provided (and combined !): -- reference Images: use this to inject people or objects in the video. You can select multiple reference Images. The integration of the image is more efficient if the background is replaced by the full white color. You can do that with your preferred background remover or use the built in background remover by checking the box *Remove background* +- a Control Video: Based on your choice, you can decide to transfer the motion, the depth in a new Video. You can tell WanGP to use only the first n frames of Control Video and to extrapolate the rest. You can also do inpainting ). If the video contains area of grey color 127, they will be considered as masks and will be filled based on the Text prompt of the reference Images. -- a Video: this can be a video that contains a body pose (an animated wireframe that indicates the positions of limbs of a person), a greyed depth map video, a normal video combined with a masked video (see below),... The Vace model will detect automatically what to do depending on the video content. You can tell WanGP to use only the n first frames of this Video. All the frames beyond and up the number of requested frames will be generated by following the Text prompt and the other visual hints (for instance reference images). If the video contains area of grey color 127, they will be considered as masks and will be filled based on the Text prompt of the reference Images. There +- reference Images: Use this to inject people or objects of your choice in the video. You can select multiple reference Images. The integration of the image is more efficient if the background is replaced by the full white color. You can do that with your preferred background remover or use the built in background remover by checking the box *Remove background* - a Video Mask -This offers a stronger mechanism to tell Vace which parts should be kept (black) or replaced (white). You can do as well inpainting / outpainting, fill the missing part of a video more efficientlty with just the video hint. +This offers a stronger mechanism to tell Vace which parts should be kept (black) or replaced (white). You can do as well inpainting / outpainting, fill the missing part of a video more efficientlty with just the video hint. If a video mask is white, it will be generated so with black frames at the beginning and at the end and the rest white, you could generate the missing frames in between. Examples: -- Inject people and / objects into a scene describe by a text promtp: Ref. Images + text Prompt -- Animate a character described in a text prompt: Body Pose Video + text Prompt -- Animate a character of your choice : Ref Images + Body Pose Video + text Prompt +- Inject people and / objects into a scene describe by a text prompt: Ref. Images + text Prompt +- Animate a character described in a text prompt: a Video of person moving + text Prompt +- Animate a character of your choice (pose transfer) : Ref Images + a Video of person moving + text Prompt +- Change the style of a scene (depth transfer): a Video that contains objects / person at differen depths + text Prompt There are lots of possible combinations. Some of them require to prepare some materials (masks on top of video, full masks, etc...). diff --git a/preprocessing/dwpose/__init__.py b/preprocessing/dwpose/__init__.py new file mode 100644 index 0000000..cc26a06 --- /dev/null +++ b/preprocessing/dwpose/__init__.py @@ -0,0 +1,2 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. diff --git a/preprocessing/dwpose/onnxdet.py b/preprocessing/dwpose/onnxdet.py new file mode 100644 index 0000000..0bcebce --- /dev/null +++ b/preprocessing/dwpose/onnxdet.py @@ -0,0 +1,127 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +import cv2 +import numpy as np + +import onnxruntime + +def nms(boxes, scores, nms_thr): + """Single class NMS implemented in Numpy.""" + x1 = boxes[:, 0] + y1 = boxes[:, 1] + x2 = boxes[:, 2] + y2 = boxes[:, 3] + + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + order = scores.argsort()[::-1] + + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + + w = np.maximum(0.0, xx2 - xx1 + 1) + h = np.maximum(0.0, yy2 - yy1 + 1) + inter = w * h + ovr = inter / (areas[i] + areas[order[1:]] - inter) + + inds = np.where(ovr <= nms_thr)[0] + order = order[inds + 1] + + return keep + +def multiclass_nms(boxes, scores, nms_thr, score_thr): + """Multiclass NMS implemented in Numpy. Class-aware version.""" + final_dets = [] + num_classes = scores.shape[1] + for cls_ind in range(num_classes): + cls_scores = scores[:, cls_ind] + valid_score_mask = cls_scores > score_thr + if valid_score_mask.sum() == 0: + continue + else: + valid_scores = cls_scores[valid_score_mask] + valid_boxes = boxes[valid_score_mask] + keep = nms(valid_boxes, valid_scores, nms_thr) + if len(keep) > 0: + cls_inds = np.ones((len(keep), 1)) * cls_ind + dets = np.concatenate( + [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1 + ) + final_dets.append(dets) + if len(final_dets) == 0: + return None + return np.concatenate(final_dets, 0) + +def demo_postprocess(outputs, img_size, p6=False): + grids = [] + expanded_strides = [] + strides = [8, 16, 32] if not p6 else [8, 16, 32, 64] + + hsizes = [img_size[0] // stride for stride in strides] + wsizes = [img_size[1] // stride for stride in strides] + + for hsize, wsize, stride in zip(hsizes, wsizes, strides): + xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize)) + grid = np.stack((xv, yv), 2).reshape(1, -1, 2) + grids.append(grid) + shape = grid.shape[:2] + expanded_strides.append(np.full((*shape, 1), stride)) + + grids = np.concatenate(grids, 1) + expanded_strides = np.concatenate(expanded_strides, 1) + outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides + outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides + + return outputs + +def preprocess(img, input_size, swap=(2, 0, 1)): + if len(img.shape) == 3: + padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114 + else: + padded_img = np.ones(input_size, dtype=np.uint8) * 114 + + r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1]) + resized_img = cv2.resize( + img, + (int(img.shape[1] * r), int(img.shape[0] * r)), + interpolation=cv2.INTER_LINEAR, + ).astype(np.uint8) + padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img + + padded_img = padded_img.transpose(swap) + padded_img = np.ascontiguousarray(padded_img, dtype=np.float32) + return padded_img, r + +def inference_detector(session, oriImg): + input_shape = (640,640) + img, ratio = preprocess(oriImg, input_shape) + + ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]} + output = session.run(None, ort_inputs) + predictions = demo_postprocess(output[0], input_shape)[0] + + boxes = predictions[:, :4] + scores = predictions[:, 4:5] * predictions[:, 5:] + + boxes_xyxy = np.ones_like(boxes) + boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2. + boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2. + boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2. + boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2. + boxes_xyxy /= ratio + dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1) + if dets is not None: + final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5] + isscore = final_scores>0.3 + iscat = final_cls_inds == 0 + isbbox = [ i and j for (i, j) in zip(isscore, iscat)] + final_boxes = final_boxes[isbbox] + else: + final_boxes = np.array([]) + + return final_boxes diff --git a/preprocessing/dwpose/onnxpose.py b/preprocessing/dwpose/onnxpose.py new file mode 100644 index 0000000..16316ca --- /dev/null +++ b/preprocessing/dwpose/onnxpose.py @@ -0,0 +1,362 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import List, Tuple + +import cv2 +import numpy as np +import onnxruntime as ort + +def preprocess( + img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256) +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Do preprocessing for RTMPose model inference. + + Args: + img (np.ndarray): Input image in shape. + input_size (tuple): Input image size in shape (w, h). + + Returns: + tuple: + - resized_img (np.ndarray): Preprocessed image. + - center (np.ndarray): Center of image. + - scale (np.ndarray): Scale of image. + """ + # get shape of image + img_shape = img.shape[:2] + out_img, out_center, out_scale = [], [], [] + if len(out_bbox) == 0: + out_bbox = [[0, 0, img_shape[1], img_shape[0]]] + for i in range(len(out_bbox)): + x0 = out_bbox[i][0] + y0 = out_bbox[i][1] + x1 = out_bbox[i][2] + y1 = out_bbox[i][3] + bbox = np.array([x0, y0, x1, y1]) + + # get center and scale + center, scale = bbox_xyxy2cs(bbox, padding=1.25) + + # do affine transformation + resized_img, scale = top_down_affine(input_size, scale, center, img) + + # normalize image + mean = np.array([123.675, 116.28, 103.53]) + std = np.array([58.395, 57.12, 57.375]) + resized_img = (resized_img - mean) / std + + out_img.append(resized_img) + out_center.append(center) + out_scale.append(scale) + + return out_img, out_center, out_scale + + +def inference(sess: ort.InferenceSession, img: np.ndarray) -> np.ndarray: + """Inference RTMPose model. + + Args: + sess (ort.InferenceSession): ONNXRuntime session. + img (np.ndarray): Input image in shape. + + Returns: + outputs (np.ndarray): Output of RTMPose model. + """ + all_out = [] + # build input + for i in range(len(img)): + input = [img[i].transpose(2, 0, 1)] + + # build output + sess_input = {sess.get_inputs()[0].name: input} + sess_output = [] + for out in sess.get_outputs(): + sess_output.append(out.name) + + # run model + outputs = sess.run(sess_output, sess_input) + all_out.append(outputs) + + return all_out + + +def postprocess(outputs: List[np.ndarray], + model_input_size: Tuple[int, int], + center: Tuple[int, int], + scale: Tuple[int, int], + simcc_split_ratio: float = 2.0 + ) -> Tuple[np.ndarray, np.ndarray]: + """Postprocess for RTMPose model output. + + Args: + outputs (np.ndarray): Output of RTMPose model. + model_input_size (tuple): RTMPose model Input image size. + center (tuple): Center of bbox in shape (x, y). + scale (tuple): Scale of bbox in shape (w, h). + simcc_split_ratio (float): Split ratio of simcc. + + Returns: + tuple: + - keypoints (np.ndarray): Rescaled keypoints. + - scores (np.ndarray): Model predict scores. + """ + all_key = [] + all_score = [] + for i in range(len(outputs)): + # use simcc to decode + simcc_x, simcc_y = outputs[i] + keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio) + + # rescale keypoints + keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2 + all_key.append(keypoints[0]) + all_score.append(scores[0]) + + return np.array(all_key), np.array(all_score) + + +def bbox_xyxy2cs(bbox: np.ndarray, + padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]: + """Transform the bbox format from (x,y,w,h) into (center, scale) + + Args: + bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted + as (left, top, right, bottom) + padding (float): BBox padding factor that will be multilied to scale. + Default: 1.0 + + Returns: + tuple: A tuple containing center and scale. + - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or + (n, 2) + - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or + (n, 2) + """ + # convert single bbox from (4, ) to (1, 4) + dim = bbox.ndim + if dim == 1: + bbox = bbox[None, :] + + # get bbox center and scale + x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3]) + center = np.hstack([x1 + x2, y1 + y2]) * 0.5 + scale = np.hstack([x2 - x1, y2 - y1]) * padding + + if dim == 1: + center = center[0] + scale = scale[0] + + return center, scale + + +def _fix_aspect_ratio(bbox_scale: np.ndarray, + aspect_ratio: float) -> np.ndarray: + """Extend the scale to match the given aspect ratio. + + Args: + scale (np.ndarray): The image scale (w, h) in shape (2, ) + aspect_ratio (float): The ratio of ``w/h`` + + Returns: + np.ndarray: The reshaped image scale in (2, ) + """ + w, h = np.hsplit(bbox_scale, [1]) + bbox_scale = np.where(w > h * aspect_ratio, + np.hstack([w, w / aspect_ratio]), + np.hstack([h * aspect_ratio, h])) + return bbox_scale + + +def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray: + """Rotate a point by an angle. + + Args: + pt (np.ndarray): 2D point coordinates (x, y) in shape (2, ) + angle_rad (float): rotation angle in radian + + Returns: + np.ndarray: Rotated point in shape (2, ) + """ + sn, cs = np.sin(angle_rad), np.cos(angle_rad) + rot_mat = np.array([[cs, -sn], [sn, cs]]) + return rot_mat @ pt + + +def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray: + """To calculate the affine matrix, three pairs of points are required. This + function is used to get the 3rd point, given 2D points a & b. + + The 3rd point is defined by rotating vector `a - b` by 90 degrees + anticlockwise, using b as the rotation center. + + Args: + a (np.ndarray): The 1st point (x,y) in shape (2, ) + b (np.ndarray): The 2nd point (x,y) in shape (2, ) + + Returns: + np.ndarray: The 3rd point. + """ + direction = a - b + c = b + np.r_[-direction[1], direction[0]] + return c + + +def get_warp_matrix(center: np.ndarray, + scale: np.ndarray, + rot: float, + output_size: Tuple[int, int], + shift: Tuple[float, float] = (0., 0.), + inv: bool = False) -> np.ndarray: + """Calculate the affine transformation matrix that can warp the bbox area + in the input image to the output size. + + Args: + center (np.ndarray[2, ]): Center of the bounding box (x, y). + scale (np.ndarray[2, ]): Scale of the bounding box + wrt [width, height]. + rot (float): Rotation angle (degree). + output_size (np.ndarray[2, ] | list(2,)): Size of the + destination heatmaps. + shift (0-100%): Shift translation ratio wrt the width/height. + Default (0., 0.). + inv (bool): Option to inverse the affine transform direction. + (inv=False: src->dst or inv=True: dst->src) + + Returns: + np.ndarray: A 2x3 transformation matrix + """ + shift = np.array(shift) + src_w = scale[0] + dst_w = output_size[0] + dst_h = output_size[1] + + # compute transformation matrix + rot_rad = np.deg2rad(rot) + src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad) + dst_dir = np.array([0., dst_w * -0.5]) + + # get four corners of the src rectangle in the original image + src = np.zeros((3, 2), dtype=np.float32) + src[0, :] = center + scale * shift + src[1, :] = center + src_dir + scale * shift + src[2, :] = _get_3rd_point(src[0, :], src[1, :]) + + # get four corners of the dst rectangle in the input image + dst = np.zeros((3, 2), dtype=np.float32) + dst[0, :] = [dst_w * 0.5, dst_h * 0.5] + dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir + dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :]) + + if inv: + warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src)) + else: + warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst)) + + return warp_mat + + +def top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict, + img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Get the bbox image as the model input by affine transform. + + Args: + input_size (dict): The input size of the model. + bbox_scale (dict): The bbox scale of the img. + bbox_center (dict): The bbox center of the img. + img (np.ndarray): The original image. + + Returns: + tuple: A tuple containing center and scale. + - np.ndarray[float32]: img after affine transform. + - np.ndarray[float32]: bbox scale after affine transform. + """ + w, h = input_size + warp_size = (int(w), int(h)) + + # reshape bbox to fixed aspect ratio + bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h) + + # get the affine matrix + center = bbox_center + scale = bbox_scale + rot = 0 + warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h)) + + # do affine transform + img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR) + + return img, bbox_scale + + +def get_simcc_maximum(simcc_x: np.ndarray, + simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Get maximum response location and value from simcc representations. + + Note: + instance number: N + num_keypoints: K + heatmap height: H + heatmap width: W + + Args: + simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx) + simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy) + + Returns: + tuple: + - locs (np.ndarray): locations of maximum heatmap responses in shape + (K, 2) or (N, K, 2) + - vals (np.ndarray): values of maximum heatmap responses in shape + (K,) or (N, K) + """ + N, K, Wx = simcc_x.shape + simcc_x = simcc_x.reshape(N * K, -1) + simcc_y = simcc_y.reshape(N * K, -1) + + # get maximum value locations + x_locs = np.argmax(simcc_x, axis=1) + y_locs = np.argmax(simcc_y, axis=1) + locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) + max_val_x = np.amax(simcc_x, axis=1) + max_val_y = np.amax(simcc_y, axis=1) + + # get maximum value across x and y axis + mask = max_val_x > max_val_y + max_val_x[mask] = max_val_y[mask] + vals = max_val_x + locs[vals <= 0.] = -1 + + # reshape + locs = locs.reshape(N, K, 2) + vals = vals.reshape(N, K) + + return locs, vals + + +def decode(simcc_x: np.ndarray, simcc_y: np.ndarray, + simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]: + """Modulate simcc distribution with Gaussian. + + Args: + simcc_x (np.ndarray[K, Wx]): model predicted simcc in x. + simcc_y (np.ndarray[K, Wy]): model predicted simcc in y. + simcc_split_ratio (int): The split ratio of simcc. + + Returns: + tuple: A tuple containing center and scale. + - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2) + - np.ndarray[float32]: scores in shape (K,) or (n, K) + """ + keypoints, scores = get_simcc_maximum(simcc_x, simcc_y) + keypoints /= simcc_split_ratio + + return keypoints, scores + + +def inference_pose(session, out_bbox, oriImg): + h, w = session.get_inputs()[0].shape[2:] + model_input_size = (w, h) + resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size) + outputs = inference(session, resized_img) + keypoints, scores = postprocess(outputs, model_input_size, center, scale) + + return keypoints, scores \ No newline at end of file diff --git a/preprocessing/dwpose/pose.py b/preprocessing/dwpose/pose.py new file mode 100644 index 0000000..ff7b0ca --- /dev/null +++ b/preprocessing/dwpose/pose.py @@ -0,0 +1,183 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. + +import os + +import cv2 +import torch +import numpy as np +from . import util +from .wholebody import Wholebody, HWC3, resize_image +from PIL import Image + +os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" + +def convert_to_numpy(image): + if isinstance(image, Image.Image): + image = np.array(image) + elif isinstance(image, torch.Tensor): + image = image.detach().cpu().numpy() + elif isinstance(image, np.ndarray): + image = image.copy() + else: + raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.' + return image + + + +def draw_pose(pose, H, W, use_hand=False, use_body=False, use_face=False): + bodies = pose['bodies'] + faces = pose['faces'] + hands = pose['hands'] + candidate = bodies['candidate'] + subset = bodies['subset'] + canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8) + + if use_body: + canvas = util.draw_bodypose(canvas, candidate, subset) + if use_hand: + canvas = util.draw_handpose(canvas, hands) + if use_face: + canvas = util.draw_facepose(canvas, faces) + + return canvas + + +class PoseAnnotator: + def __init__(self, cfg, device=None): + onnx_det = cfg['DETECTION_MODEL'] + onnx_pose = cfg['POSE_MODEL'] + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device + self.pose_estimation = Wholebody(onnx_det, onnx_pose, device=self.device) + self.resize_size = cfg.get("RESIZE_SIZE", 1024) + self.use_body = cfg.get('USE_BODY', True) + self.use_face = cfg.get('USE_FACE', True) + self.use_hand = cfg.get('USE_HAND', True) + + @torch.no_grad() + @torch.inference_mode + def forward(self, image): + image = convert_to_numpy(image) + input_image = HWC3(image[..., ::-1]) + return self.process(resize_image(input_image, self.resize_size), image.shape[:2]) + + def process(self, ori_img, ori_shape): + ori_h, ori_w = ori_shape + ori_img = ori_img.copy() + H, W, C = ori_img.shape + with torch.no_grad(): + candidate, subset, det_result = self.pose_estimation(ori_img) + nums, keys, locs = candidate.shape + candidate[..., 0] /= float(W) + candidate[..., 1] /= float(H) + body = candidate[:, :18].copy() + body = body.reshape(nums * 18, locs) + score = subset[:, :18] + for i in range(len(score)): + for j in range(len(score[i])): + if score[i][j] > 0.3: + score[i][j] = int(18 * i + j) + else: + score[i][j] = -1 + + un_visible = subset < 0.3 + candidate[un_visible] = -1 + + foot = candidate[:, 18:24] + + faces = candidate[:, 24:92] + + hands = candidate[:, 92:113] + hands = np.vstack([hands, candidate[:, 113:]]) + + bodies = dict(candidate=body, subset=score) + pose = dict(bodies=bodies, hands=hands, faces=faces) + + ret_data = {} + if self.use_body: + detected_map_body = draw_pose(pose, H, W, use_body=True) + detected_map_body = cv2.resize(detected_map_body[..., ::-1], (ori_w, ori_h), + interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA) + ret_data["detected_map_body"] = detected_map_body + + if self.use_face: + detected_map_face = draw_pose(pose, H, W, use_face=True) + detected_map_face = cv2.resize(detected_map_face[..., ::-1], (ori_w, ori_h), + interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA) + ret_data["detected_map_face"] = detected_map_face + + if self.use_body and self.use_face: + detected_map_bodyface = draw_pose(pose, H, W, use_body=True, use_face=True) + detected_map_bodyface = cv2.resize(detected_map_bodyface[..., ::-1], (ori_w, ori_h), + interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA) + ret_data["detected_map_bodyface"] = detected_map_bodyface + + if self.use_hand and self.use_body and self.use_face: + detected_map_handbodyface = draw_pose(pose, H, W, use_hand=True, use_body=True, use_face=True) + detected_map_handbodyface = cv2.resize(detected_map_handbodyface[..., ::-1], (ori_w, ori_h), + interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA) + ret_data["detected_map_handbodyface"] = detected_map_handbodyface + + # convert_size + if det_result.shape[0] > 0: + w_ratio, h_ratio = ori_w / W, ori_h / H + det_result[..., ::2] *= h_ratio + det_result[..., 1::2] *= w_ratio + det_result = det_result.astype(np.int32) + return ret_data, det_result + + +class PoseBodyFaceAnnotator(PoseAnnotator): + def __init__(self, cfg): + super().__init__(cfg) + self.use_body, self.use_face, self.use_hand = True, True, False + @torch.no_grad() + @torch.inference_mode + def forward(self, image): + ret_data, det_result = super().forward(image) + return ret_data['detected_map_bodyface'] + + +class PoseBodyFaceVideoAnnotator(PoseBodyFaceAnnotator): + def forward(self, frames): + ret_frames = [] + for frame in frames: + anno_frame = super().forward(np.array(frame)) + ret_frames.append(anno_frame) + return ret_frames + +import imageio + +def save_one_video(file_path, videos, fps=8, quality=8, macro_block_size=None): + try: + video_writer = imageio.get_writer(file_path, fps=fps, codec='libx264', quality=quality, macro_block_size=macro_block_size) + for frame in videos: + video_writer.append_data(frame) + video_writer.close() + return True + except Exception as e: + print(f"Video save error: {e}") + return False + +def get_frames(video_path): + frames = [] + + + # Opens the Video file with CV2 + cap = cv2.VideoCapture(video_path) + + fps = cap.get(cv2.CAP_PROP_FPS) + print("video fps: " + str(fps)) + i = 0 + while cap.isOpened(): + ret, frame = cap.read() + if ret == False: + break + frames.append(frame) + i += 1 + + cap.release() + cv2.destroyAllWindows() + + return frames, fps + diff --git a/preprocessing/dwpose/util.py b/preprocessing/dwpose/util.py new file mode 100644 index 0000000..232de86 --- /dev/null +++ b/preprocessing/dwpose/util.py @@ -0,0 +1,299 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +import math +import numpy as np +import matplotlib +import cv2 + + +eps = 0.01 + + +def smart_resize(x, s): + Ht, Wt = s + if x.ndim == 2: + Ho, Wo = x.shape + Co = 1 + else: + Ho, Wo, Co = x.shape + if Co == 3 or Co == 1: + k = float(Ht + Wt) / float(Ho + Wo) + return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4) + else: + return np.stack([smart_resize(x[:, :, i], s) for i in range(Co)], axis=2) + + +def smart_resize_k(x, fx, fy): + if x.ndim == 2: + Ho, Wo = x.shape + Co = 1 + else: + Ho, Wo, Co = x.shape + Ht, Wt = Ho * fy, Wo * fx + if Co == 3 or Co == 1: + k = float(Ht + Wt) / float(Ho + Wo) + return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4) + else: + return np.stack([smart_resize_k(x[:, :, i], fx, fy) for i in range(Co)], axis=2) + + +def padRightDownCorner(img, stride, padValue): + h = img.shape[0] + w = img.shape[1] + + pad = 4 * [None] + pad[0] = 0 # up + pad[1] = 0 # left + pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down + pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right + + img_padded = img + pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1)) + img_padded = np.concatenate((pad_up, img_padded), axis=0) + pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1)) + img_padded = np.concatenate((pad_left, img_padded), axis=1) + pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1)) + img_padded = np.concatenate((img_padded, pad_down), axis=0) + pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1)) + img_padded = np.concatenate((img_padded, pad_right), axis=1) + + return img_padded, pad + + +def transfer(model, model_weights): + transfered_model_weights = {} + for weights_name in model.state_dict().keys(): + transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])] + return transfered_model_weights + + +def draw_bodypose(canvas, candidate, subset): + H, W, C = canvas.shape + candidate = np.array(candidate) + subset = np.array(subset) + + stickwidth = 4 + + limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \ + [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \ + [1, 16], [16, 18], [3, 17], [6, 18]] + + colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \ + [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \ + [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]] + + for i in range(17): + for n in range(len(subset)): + index = subset[n][np.array(limbSeq[i]) - 1] + if -1 in index: + continue + Y = candidate[index.astype(int), 0] * float(W) + X = candidate[index.astype(int), 1] * float(H) + mX = np.mean(X) + mY = np.mean(Y) + length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 + angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) + polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1) + cv2.fillConvexPoly(canvas, polygon, colors[i]) + + canvas = (canvas * 0.6).astype(np.uint8) + + for i in range(18): + for n in range(len(subset)): + index = int(subset[n][i]) + if index == -1: + continue + x, y = candidate[index][0:2] + x = int(x * W) + y = int(y * H) + cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1) + + return canvas + + +def draw_handpose(canvas, all_hand_peaks): + H, W, C = canvas.shape + + edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \ + [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] + + for peaks in all_hand_peaks: + peaks = np.array(peaks) + + for ie, e in enumerate(edges): + x1, y1 = peaks[e[0]] + x2, y2 = peaks[e[1]] + x1 = int(x1 * W) + y1 = int(y1 * H) + x2 = int(x2 * W) + y2 = int(y2 * H) + if x1 > eps and y1 > eps and x2 > eps and y2 > eps: + cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255, thickness=2) + + for i, keyponit in enumerate(peaks): + x, y = keyponit + x = int(x * W) + y = int(y * H) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1) + return canvas + + +def draw_facepose(canvas, all_lmks): + H, W, C = canvas.shape + for lmks in all_lmks: + lmks = np.array(lmks) + for lmk in lmks: + x, y = lmk + x = int(x * W) + y = int(y * H) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1) + return canvas + + +# detect hand according to body pose keypoints +# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp +def handDetect(candidate, subset, oriImg): + # right hand: wrist 4, elbow 3, shoulder 2 + # left hand: wrist 7, elbow 6, shoulder 5 + ratioWristElbow = 0.33 + detect_result = [] + image_height, image_width = oriImg.shape[0:2] + for person in subset.astype(int): + # if any of three not detected + has_left = np.sum(person[[5, 6, 7]] == -1) == 0 + has_right = np.sum(person[[2, 3, 4]] == -1) == 0 + if not (has_left or has_right): + continue + hands = [] + #left hand + if has_left: + left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]] + x1, y1 = candidate[left_shoulder_index][:2] + x2, y2 = candidate[left_elbow_index][:2] + x3, y3 = candidate[left_wrist_index][:2] + hands.append([x1, y1, x2, y2, x3, y3, True]) + # right hand + if has_right: + right_shoulder_index, right_elbow_index, right_wrist_index = person[[2, 3, 4]] + x1, y1 = candidate[right_shoulder_index][:2] + x2, y2 = candidate[right_elbow_index][:2] + x3, y3 = candidate[right_wrist_index][:2] + hands.append([x1, y1, x2, y2, x3, y3, False]) + + for x1, y1, x2, y2, x3, y3, is_left in hands: + # pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox + # handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]); + # handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]); + # const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow); + # const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder); + # handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder); + x = x3 + ratioWristElbow * (x3 - x2) + y = y3 + ratioWristElbow * (y3 - y2) + distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2) + distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) + width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder) + # x-y refers to the center --> offset to topLeft point + # handRectangle.x -= handRectangle.width / 2.f; + # handRectangle.y -= handRectangle.height / 2.f; + x -= width / 2 + y -= width / 2 # width = height + # overflow the image + if x < 0: x = 0 + if y < 0: y = 0 + width1 = width + width2 = width + if x + width > image_width: width1 = image_width - x + if y + width > image_height: width2 = image_height - y + width = min(width1, width2) + # the max hand box value is 20 pixels + if width >= 20: + detect_result.append([int(x), int(y), int(width), is_left]) + + ''' + return value: [[x, y, w, True if left hand else False]]. + width=height since the network require squared input. + x, y is the coordinate of top left + ''' + return detect_result + + +# Written by Lvmin +def faceDetect(candidate, subset, oriImg): + # left right eye ear 14 15 16 17 + detect_result = [] + image_height, image_width = oriImg.shape[0:2] + for person in subset.astype(int): + has_head = person[0] > -1 + if not has_head: + continue + + has_left_eye = person[14] > -1 + has_right_eye = person[15] > -1 + has_left_ear = person[16] > -1 + has_right_ear = person[17] > -1 + + if not (has_left_eye or has_right_eye or has_left_ear or has_right_ear): + continue + + head, left_eye, right_eye, left_ear, right_ear = person[[0, 14, 15, 16, 17]] + + width = 0.0 + x0, y0 = candidate[head][:2] + + if has_left_eye: + x1, y1 = candidate[left_eye][:2] + d = max(abs(x0 - x1), abs(y0 - y1)) + width = max(width, d * 3.0) + + if has_right_eye: + x1, y1 = candidate[right_eye][:2] + d = max(abs(x0 - x1), abs(y0 - y1)) + width = max(width, d * 3.0) + + if has_left_ear: + x1, y1 = candidate[left_ear][:2] + d = max(abs(x0 - x1), abs(y0 - y1)) + width = max(width, d * 1.5) + + if has_right_ear: + x1, y1 = candidate[right_ear][:2] + d = max(abs(x0 - x1), abs(y0 - y1)) + width = max(width, d * 1.5) + + x, y = x0, y0 + + x -= width + y -= width + + if x < 0: + x = 0 + + if y < 0: + y = 0 + + width1 = width * 2 + width2 = width * 2 + + if x + width > image_width: + width1 = image_width - x + + if y + width > image_height: + width2 = image_height - y + + width = min(width1, width2) + + if width >= 20: + detect_result.append([int(x), int(y), int(width)]) + + return detect_result + + +# get max index of 2d array +def npmax(array): + arrayindex = array.argmax(1) + arrayvalue = array.max(1) + i = arrayvalue.argmax() + j = arrayindex[i] + return i, j diff --git a/preprocessing/dwpose/wholebody.py b/preprocessing/dwpose/wholebody.py new file mode 100644 index 0000000..1ea43f3 --- /dev/null +++ b/preprocessing/dwpose/wholebody.py @@ -0,0 +1,80 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +import cv2 +import numpy as np +import onnxruntime as ort +from .onnxdet import inference_detector +from .onnxpose import inference_pose + +def HWC3(x): + assert x.dtype == np.uint8 + if x.ndim == 2: + x = x[:, :, None] + assert x.ndim == 3 + H, W, C = x.shape + assert C == 1 or C == 3 or C == 4 + if C == 3: + return x + if C == 1: + return np.concatenate([x, x, x], axis=2) + if C == 4: + color = x[:, :, 0:3].astype(np.float32) + alpha = x[:, :, 3:4].astype(np.float32) / 255.0 + y = color * alpha + 255.0 * (1.0 - alpha) + y = y.clip(0, 255).astype(np.uint8) + return y + + +def resize_image(input_image, resolution): + H, W, C = input_image.shape + H = float(H) + W = float(W) + k = float(resolution) / min(H, W) + H *= k + W *= k + H = int(np.round(H / 64.0)) * 64 + W = int(np.round(W / 64.0)) * 64 + img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA) + return img + +class Wholebody: + def __init__(self, onnx_det, onnx_pose, device = 'cuda:0'): + + providers = ['CPUExecutionProvider' + ] if device == 'cpu' else ['CUDAExecutionProvider'] + # onnx_det = 'annotator/ckpts/yolox_l.onnx' + # onnx_pose = 'annotator/ckpts/dw-ll_ucoco_384.onnx' + + self.session_det = ort.InferenceSession(path_or_bytes=onnx_det, providers=providers) + self.session_pose = ort.InferenceSession(path_or_bytes=onnx_pose, providers=providers) + + def __call__(self, ori_img): + det_result = inference_detector(self.session_det, ori_img) + keypoints, scores = inference_pose(self.session_pose, det_result, ori_img) + + keypoints_info = np.concatenate( + (keypoints, scores[..., None]), axis=-1) + # compute neck joint + neck = np.mean(keypoints_info[:, [5, 6]], axis=1) + # neck score when visualizing pred + neck[:, 2:4] = np.logical_and( + keypoints_info[:, 5, 2:4] > 0.3, + keypoints_info[:, 6, 2:4] > 0.3).astype(int) + new_keypoints_info = np.insert( + keypoints_info, 17, neck, axis=1) + mmpose_idx = [ + 17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3 + ] + openpose_idx = [ + 1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17 + ] + new_keypoints_info[:, openpose_idx] = \ + new_keypoints_info[:, mmpose_idx] + keypoints_info = new_keypoints_info + + keypoints, scores = keypoints_info[ + ..., :2], keypoints_info[..., 2] + + return keypoints, scores, det_result + + diff --git a/preprocessing/gray.py b/preprocessing/gray.py new file mode 100644 index 0000000..b1b35c7 --- /dev/null +++ b/preprocessing/gray.py @@ -0,0 +1,35 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. + +import cv2 +import numpy as np +from PIL import Image +import torch + +def convert_to_numpy(image): + if isinstance(image, Image.Image): + image = np.array(image) + elif isinstance(image, torch.Tensor): + image = image.detach().cpu().numpy() + elif isinstance(image, np.ndarray): + image = image.copy() + else: + raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.' + return image + +class GrayAnnotator: + def __init__(self, cfg): + pass + def forward(self, image): + image = convert_to_numpy(image) + gray_map = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + return gray_map[..., None].repeat(3, axis=2) + + +class GrayVideoAnnotator(GrayAnnotator): + def forward(self, frames): + ret_frames = [] + for frame in frames: + anno_frame = super().forward(np.array(frame)) + ret_frames.append(anno_frame) + return ret_frames diff --git a/preprocessing/midas/__init__.py b/preprocessing/midas/__init__.py new file mode 100644 index 0000000..cc26a06 --- /dev/null +++ b/preprocessing/midas/__init__.py @@ -0,0 +1,2 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. diff --git a/preprocessing/midas/api.py b/preprocessing/midas/api.py new file mode 100644 index 0000000..87beeb7 --- /dev/null +++ b/preprocessing/midas/api.py @@ -0,0 +1,166 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +# based on https://github.com/isl-org/MiDaS + +import cv2 +import torch +import torch.nn as nn +from torchvision.transforms import Compose + +from .dpt_depth import DPTDepthModel +from .midas_net import MidasNet +from .midas_net_custom import MidasNet_small +from .transforms import NormalizeImage, PrepareForNet, Resize + +# ISL_PATHS = { +# "dpt_large": "dpt_large-midas-2f21e586.pt", +# "dpt_hybrid": "dpt_hybrid-midas-501f0c75.pt", +# "midas_v21": "", +# "midas_v21_small": "", +# } + +# remote_model_path = +# "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt" + + +def disabled_train(self, mode=True): + """Overwrite model.train with this function to make sure train/eval mode + does not change anymore.""" + return self + + +def load_midas_transform(model_type): + # https://github.com/isl-org/MiDaS/blob/master/run.py + # load transform only + if model_type == 'dpt_large': # DPT-Large + net_w, net_h = 384, 384 + resize_mode = 'minimal' + normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], + std=[0.5, 0.5, 0.5]) + + elif model_type == 'dpt_hybrid': # DPT-Hybrid + net_w, net_h = 384, 384 + resize_mode = 'minimal' + normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], + std=[0.5, 0.5, 0.5]) + + elif model_type == 'midas_v21': + net_w, net_h = 384, 384 + resize_mode = 'upper_bound' + normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]) + + elif model_type == 'midas_v21_small': + net_w, net_h = 256, 256 + resize_mode = 'upper_bound' + normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]) + + else: + assert False, f"model_type '{model_type}' not implemented, use: --model_type large" + + transform = Compose([ + Resize( + net_w, + net_h, + resize_target=None, + keep_aspect_ratio=True, + ensure_multiple_of=32, + resize_method=resize_mode, + image_interpolation_method=cv2.INTER_CUBIC, + ), + normalization, + PrepareForNet(), + ]) + + return transform + + +def load_model(model_type, model_path): + # https://github.com/isl-org/MiDaS/blob/master/run.py + # load network + # model_path = ISL_PATHS[model_type] + if model_type == 'dpt_large': # DPT-Large + model = DPTDepthModel( + path=model_path, + backbone='vitl16_384', + non_negative=True, + ) + net_w, net_h = 384, 384 + resize_mode = 'minimal' + normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], + std=[0.5, 0.5, 0.5]) + + elif model_type == 'dpt_hybrid': # DPT-Hybrid + model = DPTDepthModel( + path=model_path, + backbone='vitb_rn50_384', + non_negative=True, + ) + net_w, net_h = 384, 384 + resize_mode = 'minimal' + normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], + std=[0.5, 0.5, 0.5]) + + elif model_type == 'midas_v21': + model = MidasNet(model_path, non_negative=True) + net_w, net_h = 384, 384 + resize_mode = 'upper_bound' + normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]) + + elif model_type == 'midas_v21_small': + model = MidasNet_small(model_path, + features=64, + backbone='efficientnet_lite3', + exportable=True, + non_negative=True, + blocks={'expand': True}) + net_w, net_h = 256, 256 + resize_mode = 'upper_bound' + normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]) + + else: + print( + f"model_type '{model_type}' not implemented, use: --model_type large" + ) + assert False + + transform = Compose([ + Resize( + net_w, + net_h, + resize_target=None, + keep_aspect_ratio=True, + ensure_multiple_of=32, + resize_method=resize_mode, + image_interpolation_method=cv2.INTER_CUBIC, + ), + normalization, + PrepareForNet(), + ]) + + return model.eval(), transform + + +class MiDaSInference(nn.Module): + MODEL_TYPES_TORCH_HUB = ['DPT_Large', 'DPT_Hybrid', 'MiDaS_small'] + MODEL_TYPES_ISL = [ + 'dpt_large', + 'dpt_hybrid', + 'midas_v21', + 'midas_v21_small', + ] + + def __init__(self, model_type, model_path): + super().__init__() + assert (model_type in self.MODEL_TYPES_ISL) + model, _ = load_model(model_type, model_path) + self.model = model + self.model.train = disabled_train + + def forward(self, x): + with torch.no_grad(): + prediction = self.model(x) + return prediction diff --git a/preprocessing/midas/base_model.py b/preprocessing/midas/base_model.py new file mode 100644 index 0000000..2f99b8e --- /dev/null +++ b/preprocessing/midas/base_model.py @@ -0,0 +1,18 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +import torch + + +class BaseModel(torch.nn.Module): + def load(self, path): + """Load model from file. + + Args: + path (str): file path + """ + parameters = torch.load(path, map_location=torch.device('cpu'), weights_only=True) + + if 'optimizer' in parameters: + parameters = parameters['model'] + + self.load_state_dict(parameters) diff --git a/preprocessing/midas/blocks.py b/preprocessing/midas/blocks.py new file mode 100644 index 0000000..8759490 --- /dev/null +++ b/preprocessing/midas/blocks.py @@ -0,0 +1,391 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +import torch +import torch.nn as nn + +from .vit import (_make_pretrained_vitb16_384, _make_pretrained_vitb_rn50_384, + _make_pretrained_vitl16_384) + + +def _make_encoder( + backbone, + features, + use_pretrained, + groups=1, + expand=False, + exportable=True, + hooks=None, + use_vit_only=False, + use_readout='ignore', +): + if backbone == 'vitl16_384': + pretrained = _make_pretrained_vitl16_384(use_pretrained, + hooks=hooks, + use_readout=use_readout) + scratch = _make_scratch( + [256, 512, 1024, 1024], features, groups=groups, + expand=expand) # ViT-L/16 - 85.0% Top1 (backbone) + elif backbone == 'vitb_rn50_384': + pretrained = _make_pretrained_vitb_rn50_384( + use_pretrained, + hooks=hooks, + use_vit_only=use_vit_only, + use_readout=use_readout, + ) + scratch = _make_scratch( + [256, 512, 768, 768], features, groups=groups, + expand=expand) # ViT-H/16 - 85.0% Top1 (backbone) + elif backbone == 'vitb16_384': + pretrained = _make_pretrained_vitb16_384(use_pretrained, + hooks=hooks, + use_readout=use_readout) + scratch = _make_scratch( + [96, 192, 384, 768], features, groups=groups, + expand=expand) # ViT-B/16 - 84.6% Top1 (backbone) + elif backbone == 'resnext101_wsl': + pretrained = _make_pretrained_resnext101_wsl(use_pretrained) + scratch = _make_scratch([256, 512, 1024, 2048], + features, + groups=groups, + expand=expand) # efficientnet_lite3 + elif backbone == 'efficientnet_lite3': + pretrained = _make_pretrained_efficientnet_lite3(use_pretrained, + exportable=exportable) + scratch = _make_scratch([32, 48, 136, 384], + features, + groups=groups, + expand=expand) # efficientnet_lite3 + else: + print(f"Backbone '{backbone}' not implemented") + assert False + + return pretrained, scratch + + +def _make_scratch(in_shape, out_shape, groups=1, expand=False): + scratch = nn.Module() + + out_shape1 = out_shape + out_shape2 = out_shape + out_shape3 = out_shape + out_shape4 = out_shape + if expand is True: + out_shape1 = out_shape + out_shape2 = out_shape * 2 + out_shape3 = out_shape * 4 + out_shape4 = out_shape * 8 + + scratch.layer1_rn = nn.Conv2d(in_shape[0], + out_shape1, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups) + scratch.layer2_rn = nn.Conv2d(in_shape[1], + out_shape2, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups) + scratch.layer3_rn = nn.Conv2d(in_shape[2], + out_shape3, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups) + scratch.layer4_rn = nn.Conv2d(in_shape[3], + out_shape4, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups) + + return scratch + + +def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False): + efficientnet = torch.hub.load('rwightman/gen-efficientnet-pytorch', + 'tf_efficientnet_lite3', + pretrained=use_pretrained, + exportable=exportable) + return _make_efficientnet_backbone(efficientnet) + + +def _make_efficientnet_backbone(effnet): + pretrained = nn.Module() + + pretrained.layer1 = nn.Sequential(effnet.conv_stem, effnet.bn1, + effnet.act1, *effnet.blocks[0:2]) + pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3]) + pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5]) + pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9]) + + return pretrained + + +def _make_resnet_backbone(resnet): + pretrained = nn.Module() + pretrained.layer1 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, + resnet.maxpool, resnet.layer1) + + pretrained.layer2 = resnet.layer2 + pretrained.layer3 = resnet.layer3 + pretrained.layer4 = resnet.layer4 + + return pretrained + + +def _make_pretrained_resnext101_wsl(use_pretrained): + resnet = torch.hub.load('facebookresearch/WSL-Images', + 'resnext101_32x8d_wsl') + return _make_resnet_backbone(resnet) + + +class Interpolate(nn.Module): + """Interpolation module. + """ + def __init__(self, scale_factor, mode, align_corners=False): + """Init. + + Args: + scale_factor (float): scaling + mode (str): interpolation mode + """ + super(Interpolate, self).__init__() + + self.interp = nn.functional.interpolate + self.scale_factor = scale_factor + self.mode = mode + self.align_corners = align_corners + + def forward(self, x): + """Forward pass. + + Args: + x (tensor): input + + Returns: + tensor: interpolated data + """ + + x = self.interp(x, + scale_factor=self.scale_factor, + mode=self.mode, + align_corners=self.align_corners) + + return x + + +class ResidualConvUnit(nn.Module): + """Residual convolution module. + """ + def __init__(self, features): + """Init. + + Args: + features (int): number of features + """ + super().__init__() + + self.conv1 = nn.Conv2d(features, + features, + kernel_size=3, + stride=1, + padding=1, + bias=True) + + self.conv2 = nn.Conv2d(features, + features, + kernel_size=3, + stride=1, + padding=1, + bias=True) + + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + """Forward pass. + + Args: + x (tensor): input + + Returns: + tensor: output + """ + out = self.relu(x) + out = self.conv1(out) + out = self.relu(out) + out = self.conv2(out) + + return out + x + + +class FeatureFusionBlock(nn.Module): + """Feature fusion block. + """ + def __init__(self, features): + """Init. + + Args: + features (int): number of features + """ + super(FeatureFusionBlock, self).__init__() + + self.resConfUnit1 = ResidualConvUnit(features) + self.resConfUnit2 = ResidualConvUnit(features) + + def forward(self, *xs): + """Forward pass. + + Returns: + tensor: output + """ + output = xs[0] + + if len(xs) == 2: + output += self.resConfUnit1(xs[1]) + + output = self.resConfUnit2(output) + + output = nn.functional.interpolate(output, + scale_factor=2, + mode='bilinear', + align_corners=True) + + return output + + +class ResidualConvUnit_custom(nn.Module): + """Residual convolution module. + """ + def __init__(self, features, activation, bn): + """Init. + + Args: + features (int): number of features + """ + super().__init__() + + self.bn = bn + + self.groups = 1 + + self.conv1 = nn.Conv2d(features, + features, + kernel_size=3, + stride=1, + padding=1, + bias=True, + groups=self.groups) + + self.conv2 = nn.Conv2d(features, + features, + kernel_size=3, + stride=1, + padding=1, + bias=True, + groups=self.groups) + + if self.bn is True: + self.bn1 = nn.BatchNorm2d(features) + self.bn2 = nn.BatchNorm2d(features) + + self.activation = activation + + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x): + """Forward pass. + + Args: + x (tensor): input + + Returns: + tensor: output + """ + + out = self.activation(x) + out = self.conv1(out) + if self.bn is True: + out = self.bn1(out) + + out = self.activation(out) + out = self.conv2(out) + if self.bn is True: + out = self.bn2(out) + + if self.groups > 1: + out = self.conv_merge(out) + + return self.skip_add.add(out, x) + + # return out + x + + +class FeatureFusionBlock_custom(nn.Module): + """Feature fusion block. + """ + def __init__(self, + features, + activation, + deconv=False, + bn=False, + expand=False, + align_corners=True): + """Init. + + Args: + features (int): number of features + """ + super(FeatureFusionBlock_custom, self).__init__() + + self.deconv = deconv + self.align_corners = align_corners + + self.groups = 1 + + self.expand = expand + out_features = features + if self.expand is True: + out_features = features // 2 + + self.out_conv = nn.Conv2d(features, + out_features, + kernel_size=1, + stride=1, + padding=0, + bias=True, + groups=1) + + self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn) + self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn) + + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, *xs): + """Forward pass. + + Returns: + tensor: output + """ + output = xs[0] + + if len(xs) == 2: + res = self.resConfUnit1(xs[1]) + output = self.skip_add.add(output, res) + # output += res + + output = self.resConfUnit2(output) + + output = nn.functional.interpolate(output, + scale_factor=2, + mode='bilinear', + align_corners=self.align_corners) + + output = self.out_conv(output) + + return output diff --git a/preprocessing/midas/depth.py b/preprocessing/midas/depth.py new file mode 100644 index 0000000..eb0f3d9 --- /dev/null +++ b/preprocessing/midas/depth.py @@ -0,0 +1,84 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. + +import numpy as np +import torch +from einops import rearrange +from PIL import Image +import cv2 + + + +def convert_to_numpy(image): + if isinstance(image, Image.Image): + image = np.array(image) + elif isinstance(image, torch.Tensor): + image = image.detach().cpu().numpy() + elif isinstance(image, np.ndarray): + image = image.copy() + else: + raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.' + return image + +def resize_image(input_image, resolution): + H, W, C = input_image.shape + H = float(H) + W = float(W) + k = float(resolution) / min(H, W) + H *= k + W *= k + H = int(np.round(H / 64.0)) * 64 + W = int(np.round(W / 64.0)) * 64 + img = cv2.resize( + input_image, (W, H), + interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA) + return img, k + + +def resize_image_ori(h, w, image, k): + img = cv2.resize( + image, (w, h), + interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA) + return img + +class DepthAnnotator: + def __init__(self, cfg, device=None): + from .api import MiDaSInference + pretrained_model = cfg['PRETRAINED_MODEL'] + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device + self.model = MiDaSInference(model_type='dpt_hybrid', model_path=pretrained_model).to(self.device) + self.a = cfg.get('A', np.pi * 2.0) + self.bg_th = cfg.get('BG_TH', 0.1) + + @torch.no_grad() + @torch.inference_mode() + @torch.autocast('cuda', enabled=False) + def forward(self, image): + image = convert_to_numpy(image) + image_depth = image + h, w, c = image.shape + image_depth, k = resize_image(image_depth, + 1024 if min(h, w) > 1024 else min(h, w)) + image_depth = torch.from_numpy(image_depth).float().to(self.device) + image_depth = image_depth / 127.5 - 1.0 + image_depth = rearrange(image_depth, 'h w c -> 1 c h w') + depth = self.model(image_depth)[0] + + depth_pt = depth.clone() + depth_pt -= torch.min(depth_pt) + depth_pt /= torch.max(depth_pt) + depth_pt = depth_pt.cpu().numpy() + depth_image = (depth_pt * 255.0).clip(0, 255).astype(np.uint8) + depth_image = depth_image[..., None].repeat(3, 2) + + depth_image = resize_image_ori(h, w, depth_image, k) + return depth_image + + +class DepthVideoAnnotator(DepthAnnotator): + def forward(self, frames): + ret_frames = [] + for frame in frames: + anno_frame = super().forward(np.array(frame)) + ret_frames.append(anno_frame) + return ret_frames \ No newline at end of file diff --git a/preprocessing/midas/dpt_depth.py b/preprocessing/midas/dpt_depth.py new file mode 100644 index 0000000..a2db4a9 --- /dev/null +++ b/preprocessing/midas/dpt_depth.py @@ -0,0 +1,107 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +import torch +import torch.nn as nn + +from .base_model import BaseModel +from .blocks import FeatureFusionBlock_custom, Interpolate, _make_encoder +from .vit import forward_vit + + +def _make_fusion_block(features, use_bn): + return FeatureFusionBlock_custom( + features, + nn.ReLU(False), + deconv=False, + bn=use_bn, + expand=False, + align_corners=True, + ) + + +class DPT(BaseModel): + def __init__( + self, + head, + features=256, + backbone='vitb_rn50_384', + readout='project', + channels_last=False, + use_bn=False, + ): + + super(DPT, self).__init__() + + self.channels_last = channels_last + + hooks = { + 'vitb_rn50_384': [0, 1, 8, 11], + 'vitb16_384': [2, 5, 8, 11], + 'vitl16_384': [5, 11, 17, 23], + } + + # Instantiate backbone and reassemble blocks + self.pretrained, self.scratch = _make_encoder( + backbone, + features, + False, # Set to true of you want to train from scratch, uses ImageNet weights + groups=1, + expand=False, + exportable=False, + hooks=hooks[backbone], + use_readout=readout, + ) + + self.scratch.refinenet1 = _make_fusion_block(features, use_bn) + self.scratch.refinenet2 = _make_fusion_block(features, use_bn) + self.scratch.refinenet3 = _make_fusion_block(features, use_bn) + self.scratch.refinenet4 = _make_fusion_block(features, use_bn) + + self.scratch.output_conv = head + + def forward(self, x): + if self.channels_last is True: + x.contiguous(memory_format=torch.channels_last) + + layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x) + + layer_1_rn = self.scratch.layer1_rn(layer_1) + layer_2_rn = self.scratch.layer2_rn(layer_2) + layer_3_rn = self.scratch.layer3_rn(layer_3) + layer_4_rn = self.scratch.layer4_rn(layer_4) + + path_4 = self.scratch.refinenet4(layer_4_rn) + path_3 = self.scratch.refinenet3(path_4, layer_3_rn) + path_2 = self.scratch.refinenet2(path_3, layer_2_rn) + path_1 = self.scratch.refinenet1(path_2, layer_1_rn) + + out = self.scratch.output_conv(path_1) + + return out + + +class DPTDepthModel(DPT): + def __init__(self, path=None, non_negative=True, **kwargs): + features = kwargs['features'] if 'features' in kwargs else 256 + + head = nn.Sequential( + nn.Conv2d(features, + features // 2, + kernel_size=3, + stride=1, + padding=1), + Interpolate(scale_factor=2, mode='bilinear', align_corners=True), + nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1), + nn.ReLU(True), + nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0), + nn.ReLU(True) if non_negative else nn.Identity(), + nn.Identity(), + ) + + super().__init__(head, **kwargs) + + if path is not None: + self.load(path) + + def forward(self, x): + return super().forward(x).squeeze(dim=1) diff --git a/preprocessing/midas/midas_net.py b/preprocessing/midas/midas_net.py new file mode 100644 index 0000000..04878f4 --- /dev/null +++ b/preprocessing/midas/midas_net.py @@ -0,0 +1,80 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +"""MidashNet: Network for monocular depth estimation trained by mixing several datasets. +This file contains code that is adapted from +https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py +""" +import torch +import torch.nn as nn + +from .base_model import BaseModel +from .blocks import FeatureFusionBlock, Interpolate, _make_encoder + + +class MidasNet(BaseModel): + """Network for monocular depth estimation. + """ + def __init__(self, path=None, features=256, non_negative=True): + """Init. + + Args: + path (str, optional): Path to saved model. Defaults to None. + features (int, optional): Number of features. Defaults to 256. + backbone (str, optional): Backbone network for encoder. Defaults to resnet50 + """ + print('Loading weights: ', path) + + super(MidasNet, self).__init__() + + use_pretrained = False if path is None else True + + self.pretrained, self.scratch = _make_encoder( + backbone='resnext101_wsl', + features=features, + use_pretrained=use_pretrained) + + self.scratch.refinenet4 = FeatureFusionBlock(features) + self.scratch.refinenet3 = FeatureFusionBlock(features) + self.scratch.refinenet2 = FeatureFusionBlock(features) + self.scratch.refinenet1 = FeatureFusionBlock(features) + + self.scratch.output_conv = nn.Sequential( + nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1), + Interpolate(scale_factor=2, mode='bilinear'), + nn.Conv2d(128, 32, kernel_size=3, stride=1, padding=1), + nn.ReLU(True), + nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0), + nn.ReLU(True) if non_negative else nn.Identity(), + ) + + if path: + self.load(path) + + def forward(self, x): + """Forward pass. + + Args: + x (tensor): input data (image) + + Returns: + tensor: depth + """ + + layer_1 = self.pretrained.layer1(x) + layer_2 = self.pretrained.layer2(layer_1) + layer_3 = self.pretrained.layer3(layer_2) + layer_4 = self.pretrained.layer4(layer_3) + + layer_1_rn = self.scratch.layer1_rn(layer_1) + layer_2_rn = self.scratch.layer2_rn(layer_2) + layer_3_rn = self.scratch.layer3_rn(layer_3) + layer_4_rn = self.scratch.layer4_rn(layer_4) + + path_4 = self.scratch.refinenet4(layer_4_rn) + path_3 = self.scratch.refinenet3(path_4, layer_3_rn) + path_2 = self.scratch.refinenet2(path_3, layer_2_rn) + path_1 = self.scratch.refinenet1(path_2, layer_1_rn) + + out = self.scratch.output_conv(path_1) + + return torch.squeeze(out, dim=1) diff --git a/preprocessing/midas/midas_net_custom.py b/preprocessing/midas/midas_net_custom.py new file mode 100644 index 0000000..7c5a354 --- /dev/null +++ b/preprocessing/midas/midas_net_custom.py @@ -0,0 +1,167 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +"""MidashNet: Network for monocular depth estimation trained by mixing several datasets. +This file contains code that is adapted from +https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py +""" +import torch +import torch.nn as nn + +from .base_model import BaseModel +from .blocks import FeatureFusionBlock_custom, Interpolate, _make_encoder + + +class MidasNet_small(BaseModel): + """Network for monocular depth estimation. + """ + def __init__(self, + path=None, + features=64, + backbone='efficientnet_lite3', + non_negative=True, + exportable=True, + channels_last=False, + align_corners=True, + blocks={'expand': True}): + """Init. + + Args: + path (str, optional): Path to saved model. Defaults to None. + features (int, optional): Number of features. Defaults to 256. + backbone (str, optional): Backbone network for encoder. Defaults to resnet50 + """ + print('Loading weights: ', path) + + super(MidasNet_small, self).__init__() + + use_pretrained = False if path else True + + self.channels_last = channels_last + self.blocks = blocks + self.backbone = backbone + + self.groups = 1 + + features1 = features + features2 = features + features3 = features + features4 = features + self.expand = False + if 'expand' in self.blocks and self.blocks['expand'] is True: + self.expand = True + features1 = features + features2 = features * 2 + features3 = features * 4 + features4 = features * 8 + + self.pretrained, self.scratch = _make_encoder(self.backbone, + features, + use_pretrained, + groups=self.groups, + expand=self.expand, + exportable=exportable) + + self.scratch.activation = nn.ReLU(False) + + self.scratch.refinenet4 = FeatureFusionBlock_custom( + features4, + self.scratch.activation, + deconv=False, + bn=False, + expand=self.expand, + align_corners=align_corners) + self.scratch.refinenet3 = FeatureFusionBlock_custom( + features3, + self.scratch.activation, + deconv=False, + bn=False, + expand=self.expand, + align_corners=align_corners) + self.scratch.refinenet2 = FeatureFusionBlock_custom( + features2, + self.scratch.activation, + deconv=False, + bn=False, + expand=self.expand, + align_corners=align_corners) + self.scratch.refinenet1 = FeatureFusionBlock_custom( + features1, + self.scratch.activation, + deconv=False, + bn=False, + align_corners=align_corners) + + self.scratch.output_conv = nn.Sequential( + nn.Conv2d(features, + features // 2, + kernel_size=3, + stride=1, + padding=1, + groups=self.groups), + Interpolate(scale_factor=2, mode='bilinear'), + nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1), + self.scratch.activation, + nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0), + nn.ReLU(True) if non_negative else nn.Identity(), + nn.Identity(), + ) + + if path: + self.load(path) + + def forward(self, x): + """Forward pass. + + Args: + x (tensor): input data (image) + + Returns: + tensor: depth + """ + if self.channels_last is True: + print('self.channels_last = ', self.channels_last) + x.contiguous(memory_format=torch.channels_last) + + layer_1 = self.pretrained.layer1(x) + layer_2 = self.pretrained.layer2(layer_1) + layer_3 = self.pretrained.layer3(layer_2) + layer_4 = self.pretrained.layer4(layer_3) + + layer_1_rn = self.scratch.layer1_rn(layer_1) + layer_2_rn = self.scratch.layer2_rn(layer_2) + layer_3_rn = self.scratch.layer3_rn(layer_3) + layer_4_rn = self.scratch.layer4_rn(layer_4) + + path_4 = self.scratch.refinenet4(layer_4_rn) + path_3 = self.scratch.refinenet3(path_4, layer_3_rn) + path_2 = self.scratch.refinenet2(path_3, layer_2_rn) + path_1 = self.scratch.refinenet1(path_2, layer_1_rn) + + out = self.scratch.output_conv(path_1) + + return torch.squeeze(out, dim=1) + + +def fuse_model(m): + prev_previous_type = nn.Identity() + prev_previous_name = '' + previous_type = nn.Identity() + previous_name = '' + for name, module in m.named_modules(): + if prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d and type( + module) == nn.ReLU: + # print("FUSED ", prev_previous_name, previous_name, name) + torch.quantization.fuse_modules( + m, [prev_previous_name, previous_name, name], inplace=True) + elif prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d: + # print("FUSED ", prev_previous_name, previous_name) + torch.quantization.fuse_modules( + m, [prev_previous_name, previous_name], inplace=True) + # elif previous_type == nn.Conv2d and type(module) == nn.ReLU: + # print("FUSED ", previous_name, name) + # torch.quantization.fuse_modules(m, [previous_name, name], inplace=True) + + prev_previous_type = previous_type + prev_previous_name = previous_name + previous_type = type(module) + previous_name = name diff --git a/preprocessing/midas/transforms.py b/preprocessing/midas/transforms.py new file mode 100644 index 0000000..5388362 --- /dev/null +++ b/preprocessing/midas/transforms.py @@ -0,0 +1,231 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +import math + +import cv2 +import numpy as np + + +def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA): + """Rezise the sample to ensure the given size. Keeps aspect ratio. + + Args: + sample (dict): sample + size (tuple): image size + + Returns: + tuple: new size + """ + shape = list(sample['disparity'].shape) + + if shape[0] >= size[0] and shape[1] >= size[1]: + return sample + + scale = [0, 0] + scale[0] = size[0] / shape[0] + scale[1] = size[1] / shape[1] + + scale = max(scale) + + shape[0] = math.ceil(scale * shape[0]) + shape[1] = math.ceil(scale * shape[1]) + + # resize + sample['image'] = cv2.resize(sample['image'], + tuple(shape[::-1]), + interpolation=image_interpolation_method) + + sample['disparity'] = cv2.resize(sample['disparity'], + tuple(shape[::-1]), + interpolation=cv2.INTER_NEAREST) + sample['mask'] = cv2.resize( + sample['mask'].astype(np.float32), + tuple(shape[::-1]), + interpolation=cv2.INTER_NEAREST, + ) + sample['mask'] = sample['mask'].astype(bool) + + return tuple(shape) + + +class Resize(object): + """Resize sample to given size (width, height). + """ + def __init__( + self, + width, + height, + resize_target=True, + keep_aspect_ratio=False, + ensure_multiple_of=1, + resize_method='lower_bound', + image_interpolation_method=cv2.INTER_AREA, + ): + """Init. + + Args: + width (int): desired output width + height (int): desired output height + resize_target (bool, optional): + True: Resize the full sample (image, mask, target). + False: Resize image only. + Defaults to True. + keep_aspect_ratio (bool, optional): + True: Keep the aspect ratio of the input sample. + Output sample might not have the given width and height, and + resize behaviour depends on the parameter 'resize_method'. + Defaults to False. + ensure_multiple_of (int, optional): + Output width and height is constrained to be multiple of this parameter. + Defaults to 1. + resize_method (str, optional): + "lower_bound": Output will be at least as large as the given size. + "upper_bound": Output will be at max as large as the given size. " + "(Output size might be smaller than given size.)" + "minimal": Scale as least as possible. (Output size might be smaller than given size.) + Defaults to "lower_bound". + """ + self.__width = width + self.__height = height + + self.__resize_target = resize_target + self.__keep_aspect_ratio = keep_aspect_ratio + self.__multiple_of = ensure_multiple_of + self.__resize_method = resize_method + self.__image_interpolation_method = image_interpolation_method + + def constrain_to_multiple_of(self, x, min_val=0, max_val=None): + y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int) + + if max_val is not None and y > max_val: + y = (np.floor(x / self.__multiple_of) * + self.__multiple_of).astype(int) + + if y < min_val: + y = (np.ceil(x / self.__multiple_of) * + self.__multiple_of).astype(int) + + return y + + def get_size(self, width, height): + # determine new height and width + scale_height = self.__height / height + scale_width = self.__width / width + + if self.__keep_aspect_ratio: + if self.__resize_method == 'lower_bound': + # scale such that output size is lower bound + if scale_width > scale_height: + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + elif self.__resize_method == 'upper_bound': + # scale such that output size is upper bound + if scale_width < scale_height: + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + elif self.__resize_method == 'minimal': + # scale as least as possbile + if abs(1 - scale_width) < abs(1 - scale_height): + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + else: + raise ValueError( + f'resize_method {self.__resize_method} not implemented') + + if self.__resize_method == 'lower_bound': + new_height = self.constrain_to_multiple_of(scale_height * height, + min_val=self.__height) + new_width = self.constrain_to_multiple_of(scale_width * width, + min_val=self.__width) + elif self.__resize_method == 'upper_bound': + new_height = self.constrain_to_multiple_of(scale_height * height, + max_val=self.__height) + new_width = self.constrain_to_multiple_of(scale_width * width, + max_val=self.__width) + elif self.__resize_method == 'minimal': + new_height = self.constrain_to_multiple_of(scale_height * height) + new_width = self.constrain_to_multiple_of(scale_width * width) + else: + raise ValueError( + f'resize_method {self.__resize_method} not implemented') + + return (new_width, new_height) + + def __call__(self, sample): + width, height = self.get_size(sample['image'].shape[1], + sample['image'].shape[0]) + + # resize sample + sample['image'] = cv2.resize( + sample['image'], + (width, height), + interpolation=self.__image_interpolation_method, + ) + + if self.__resize_target: + if 'disparity' in sample: + sample['disparity'] = cv2.resize( + sample['disparity'], + (width, height), + interpolation=cv2.INTER_NEAREST, + ) + + if 'depth' in sample: + sample['depth'] = cv2.resize(sample['depth'], (width, height), + interpolation=cv2.INTER_NEAREST) + + sample['mask'] = cv2.resize( + sample['mask'].astype(np.float32), + (width, height), + interpolation=cv2.INTER_NEAREST, + ) + sample['mask'] = sample['mask'].astype(bool) + + return sample + + +class NormalizeImage(object): + """Normlize image by given mean and std. + """ + def __init__(self, mean, std): + self.__mean = mean + self.__std = std + + def __call__(self, sample): + sample['image'] = (sample['image'] - self.__mean) / self.__std + + return sample + + +class PrepareForNet(object): + """Prepare sample for usage as network input. + """ + def __init__(self): + pass + + def __call__(self, sample): + image = np.transpose(sample['image'], (2, 0, 1)) + sample['image'] = np.ascontiguousarray(image).astype(np.float32) + + if 'mask' in sample: + sample['mask'] = sample['mask'].astype(np.float32) + sample['mask'] = np.ascontiguousarray(sample['mask']) + + if 'disparity' in sample: + disparity = sample['disparity'].astype(np.float32) + sample['disparity'] = np.ascontiguousarray(disparity) + + if 'depth' in sample: + depth = sample['depth'].astype(np.float32) + sample['depth'] = np.ascontiguousarray(depth) + + return sample diff --git a/preprocessing/midas/utils.py b/preprocessing/midas/utils.py new file mode 100644 index 0000000..8c703b1 --- /dev/null +++ b/preprocessing/midas/utils.py @@ -0,0 +1,193 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +"""Utils for monoDepth.""" +import re +import sys + +import cv2 +import numpy as np +import torch + + +def read_pfm(path): + """Read pfm file. + + Args: + path (str): path to file + + Returns: + tuple: (data, scale) + """ + with open(path, 'rb') as file: + + color = None + width = None + height = None + scale = None + endian = None + + header = file.readline().rstrip() + if header.decode('ascii') == 'PF': + color = True + elif header.decode('ascii') == 'Pf': + color = False + else: + raise Exception('Not a PFM file: ' + path) + + dim_match = re.match(r'^(\d+)\s(\d+)\s$', + file.readline().decode('ascii')) + if dim_match: + width, height = list(map(int, dim_match.groups())) + else: + raise Exception('Malformed PFM header.') + + scale = float(file.readline().decode('ascii').rstrip()) + if scale < 0: + # little-endian + endian = '<' + scale = -scale + else: + # big-endian + endian = '>' + + data = np.fromfile(file, endian + 'f') + shape = (height, width, 3) if color else (height, width) + + data = np.reshape(data, shape) + data = np.flipud(data) + + return data, scale + + +def write_pfm(path, image, scale=1): + """Write pfm file. + + Args: + path (str): pathto file + image (array): data + scale (int, optional): Scale. Defaults to 1. + """ + + with open(path, 'wb') as file: + color = None + + if image.dtype.name != 'float32': + raise Exception('Image dtype must be float32.') + + image = np.flipud(image) + + if len(image.shape) == 3 and image.shape[2] == 3: # color image + color = True + elif (len(image.shape) == 2 + or len(image.shape) == 3 and image.shape[2] == 1): # greyscale + color = False + else: + raise Exception( + 'Image must have H x W x 3, H x W x 1 or H x W dimensions.') + + file.write('PF\n' if color else 'Pf\n'.encode()) + file.write('%d %d\n'.encode() % (image.shape[1], image.shape[0])) + + endian = image.dtype.byteorder + + if endian == '<' or endian == '=' and sys.byteorder == 'little': + scale = -scale + + file.write('%f\n'.encode() % scale) + + image.tofile(file) + + +def read_image(path): + """Read image and output RGB image (0-1). + + Args: + path (str): path to file + + Returns: + array: RGB image (0-1) + """ + img = cv2.imread(path) + + if img.ndim == 2: + img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0 + + return img + + +def resize_image(img): + """Resize image and make it fit for network. + + Args: + img (array): image + + Returns: + tensor: data ready for network + """ + height_orig = img.shape[0] + width_orig = img.shape[1] + + if width_orig > height_orig: + scale = width_orig / 384 + else: + scale = height_orig / 384 + + height = (np.ceil(height_orig / scale / 32) * 32).astype(int) + width = (np.ceil(width_orig / scale / 32) * 32).astype(int) + + img_resized = cv2.resize(img, (width, height), + interpolation=cv2.INTER_AREA) + + img_resized = (torch.from_numpy(np.transpose( + img_resized, (2, 0, 1))).contiguous().float()) + img_resized = img_resized.unsqueeze(0) + + return img_resized + + +def resize_depth(depth, width, height): + """Resize depth map and bring to CPU (numpy). + + Args: + depth (tensor): depth + width (int): image width + height (int): image height + + Returns: + array: processed depth + """ + depth = torch.squeeze(depth[0, :, :, :]).to('cpu') + + depth_resized = cv2.resize(depth.numpy(), (width, height), + interpolation=cv2.INTER_CUBIC) + + return depth_resized + + +def write_depth(path, depth, bits=1): + """Write depth map to pfm and png file. + + Args: + path (str): filepath without extension + depth (array): depth + """ + write_pfm(path + '.pfm', depth.astype(np.float32)) + + depth_min = depth.min() + depth_max = depth.max() + + max_val = (2**(8 * bits)) - 1 + + if depth_max - depth_min > np.finfo('float').eps: + out = max_val * (depth - depth_min) / (depth_max - depth_min) + else: + out = np.zeros(depth.shape, dtype=depth.type) + + if bits == 1: + cv2.imwrite(path + '.png', out.astype('uint8')) + elif bits == 2: + cv2.imwrite(path + '.png', out.astype('uint16')) + + return diff --git a/preprocessing/midas/vit.py b/preprocessing/midas/vit.py new file mode 100644 index 0000000..a85488b --- /dev/null +++ b/preprocessing/midas/vit.py @@ -0,0 +1,510 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +import math +import types + +import timm +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class Slice(nn.Module): + def __init__(self, start_index=1): + super(Slice, self).__init__() + self.start_index = start_index + + def forward(self, x): + return x[:, self.start_index:] + + +class AddReadout(nn.Module): + def __init__(self, start_index=1): + super(AddReadout, self).__init__() + self.start_index = start_index + + def forward(self, x): + if self.start_index == 2: + readout = (x[:, 0] + x[:, 1]) / 2 + else: + readout = x[:, 0] + return x[:, self.start_index:] + readout.unsqueeze(1) + + +class ProjectReadout(nn.Module): + def __init__(self, in_features, start_index=1): + super(ProjectReadout, self).__init__() + self.start_index = start_index + + self.project = nn.Sequential(nn.Linear(2 * in_features, in_features), + nn.GELU()) + + def forward(self, x): + readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index:]) + features = torch.cat((x[:, self.start_index:], readout), -1) + + return self.project(features) + + +class Transpose(nn.Module): + def __init__(self, dim0, dim1): + super(Transpose, self).__init__() + self.dim0 = dim0 + self.dim1 = dim1 + + def forward(self, x): + x = x.transpose(self.dim0, self.dim1) + return x + + +def forward_vit(pretrained, x): + b, c, h, w = x.shape + + _ = pretrained.model.forward_flex(x) + + layer_1 = pretrained.activations['1'] + layer_2 = pretrained.activations['2'] + layer_3 = pretrained.activations['3'] + layer_4 = pretrained.activations['4'] + + layer_1 = pretrained.act_postprocess1[0:2](layer_1) + layer_2 = pretrained.act_postprocess2[0:2](layer_2) + layer_3 = pretrained.act_postprocess3[0:2](layer_3) + layer_4 = pretrained.act_postprocess4[0:2](layer_4) + + unflatten = nn.Sequential( + nn.Unflatten( + 2, + torch.Size([ + h // pretrained.model.patch_size[1], + w // pretrained.model.patch_size[0], + ]), + )) + + if layer_1.ndim == 3: + layer_1 = unflatten(layer_1) + if layer_2.ndim == 3: + layer_2 = unflatten(layer_2) + if layer_3.ndim == 3: + layer_3 = unflatten(layer_3) + if layer_4.ndim == 3: + layer_4 = unflatten(layer_4) + + layer_1 = pretrained.act_postprocess1[3:len(pretrained.act_postprocess1)]( + layer_1) + layer_2 = pretrained.act_postprocess2[3:len(pretrained.act_postprocess2)]( + layer_2) + layer_3 = pretrained.act_postprocess3[3:len(pretrained.act_postprocess3)]( + layer_3) + layer_4 = pretrained.act_postprocess4[3:len(pretrained.act_postprocess4)]( + layer_4) + + return layer_1, layer_2, layer_3, layer_4 + + +def _resize_pos_embed(self, posemb, gs_h, gs_w): + posemb_tok, posemb_grid = ( + posemb[:, :self.start_index], + posemb[0, self.start_index:], + ) + + gs_old = int(math.sqrt(len(posemb_grid))) + + posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, + -1).permute(0, 3, 1, 2) + posemb_grid = F.interpolate(posemb_grid, + size=(gs_h, gs_w), + mode='bilinear') + posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1) + + posemb = torch.cat([posemb_tok, posemb_grid], dim=1) + + return posemb + + +def forward_flex(self, x): + b, c, h, w = x.shape + + pos_embed = self._resize_pos_embed(self.pos_embed, h // self.patch_size[1], + w // self.patch_size[0]) + + B = x.shape[0] + + if hasattr(self.patch_embed, 'backbone'): + x = self.patch_embed.backbone(x) + if isinstance(x, (list, tuple)): + x = x[ + -1] # last feature if backbone outputs list/tuple of features + + x = self.patch_embed.proj(x).flatten(2).transpose(1, 2) + + if getattr(self, 'dist_token', None) is not None: + cls_tokens = self.cls_token.expand( + B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks + dist_token = self.dist_token.expand(B, -1, -1) + x = torch.cat((cls_tokens, dist_token, x), dim=1) + else: + cls_tokens = self.cls_token.expand( + B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks + x = torch.cat((cls_tokens, x), dim=1) + + x = x + pos_embed + x = self.pos_drop(x) + + for blk in self.blocks: + x = blk(x) + + x = self.norm(x) + + return x + + +activations = {} + + +def get_activation(name): + def hook(model, input, output): + activations[name] = output + + return hook + + +def get_readout_oper(vit_features, features, use_readout, start_index=1): + if use_readout == 'ignore': + readout_oper = [Slice(start_index)] * len(features) + elif use_readout == 'add': + readout_oper = [AddReadout(start_index)] * len(features) + elif use_readout == 'project': + readout_oper = [ + ProjectReadout(vit_features, start_index) for out_feat in features + ] + else: + assert ( + False + ), "wrong operation for readout token, use_readout can be 'ignore', 'add', or 'project'" + + return readout_oper + + +def _make_vit_b16_backbone( + model, + features=[96, 192, 384, 768], + size=[384, 384], + hooks=[2, 5, 8, 11], + vit_features=768, + use_readout='ignore', + start_index=1, +): + pretrained = nn.Module() + + pretrained.model = model + pretrained.model.blocks[hooks[0]].register_forward_hook( + get_activation('1')) + pretrained.model.blocks[hooks[1]].register_forward_hook( + get_activation('2')) + pretrained.model.blocks[hooks[2]].register_forward_hook( + get_activation('3')) + pretrained.model.blocks[hooks[3]].register_forward_hook( + get_activation('4')) + + pretrained.activations = activations + + readout_oper = get_readout_oper(vit_features, features, use_readout, + start_index) + + # 32, 48, 136, 384 + pretrained.act_postprocess1 = nn.Sequential( + readout_oper[0], + Transpose(1, 2), + nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), + nn.Conv2d( + in_channels=vit_features, + out_channels=features[0], + kernel_size=1, + stride=1, + padding=0, + ), + nn.ConvTranspose2d( + in_channels=features[0], + out_channels=features[0], + kernel_size=4, + stride=4, + padding=0, + bias=True, + dilation=1, + groups=1, + ), + ) + + pretrained.act_postprocess2 = nn.Sequential( + readout_oper[1], + Transpose(1, 2), + nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), + nn.Conv2d( + in_channels=vit_features, + out_channels=features[1], + kernel_size=1, + stride=1, + padding=0, + ), + nn.ConvTranspose2d( + in_channels=features[1], + out_channels=features[1], + kernel_size=2, + stride=2, + padding=0, + bias=True, + dilation=1, + groups=1, + ), + ) + + pretrained.act_postprocess3 = nn.Sequential( + readout_oper[2], + Transpose(1, 2), + nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), + nn.Conv2d( + in_channels=vit_features, + out_channels=features[2], + kernel_size=1, + stride=1, + padding=0, + ), + ) + + pretrained.act_postprocess4 = nn.Sequential( + readout_oper[3], + Transpose(1, 2), + nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), + nn.Conv2d( + in_channels=vit_features, + out_channels=features[3], + kernel_size=1, + stride=1, + padding=0, + ), + nn.Conv2d( + in_channels=features[3], + out_channels=features[3], + kernel_size=3, + stride=2, + padding=1, + ), + ) + + pretrained.model.start_index = start_index + pretrained.model.patch_size = [16, 16] + + # We inject this function into the VisionTransformer instances so that + # we can use it with interpolated position embeddings without modifying the library source. + pretrained.model.forward_flex = types.MethodType(forward_flex, + pretrained.model) + pretrained.model._resize_pos_embed = types.MethodType( + _resize_pos_embed, pretrained.model) + + return pretrained + + +def _make_pretrained_vitl16_384(pretrained, use_readout='ignore', hooks=None): + model = timm.create_model('vit_large_patch16_384', pretrained=pretrained) + + hooks = [5, 11, 17, 23] if hooks is None else hooks + return _make_vit_b16_backbone( + model, + features=[256, 512, 1024, 1024], + hooks=hooks, + vit_features=1024, + use_readout=use_readout, + ) + + +def _make_pretrained_vitb16_384(pretrained, use_readout='ignore', hooks=None): + model = timm.create_model('vit_base_patch16_384', pretrained=pretrained) + + hooks = [2, 5, 8, 11] if hooks is None else hooks + return _make_vit_b16_backbone(model, + features=[96, 192, 384, 768], + hooks=hooks, + use_readout=use_readout) + + +def _make_pretrained_deitb16_384(pretrained, use_readout='ignore', hooks=None): + model = timm.create_model('vit_deit_base_patch16_384', + pretrained=pretrained) + + hooks = [2, 5, 8, 11] if hooks is None else hooks + return _make_vit_b16_backbone(model, + features=[96, 192, 384, 768], + hooks=hooks, + use_readout=use_readout) + + +def _make_pretrained_deitb16_distil_384(pretrained, + use_readout='ignore', + hooks=None): + model = timm.create_model('vit_deit_base_distilled_patch16_384', + pretrained=pretrained) + + hooks = [2, 5, 8, 11] if hooks is None else hooks + return _make_vit_b16_backbone( + model, + features=[96, 192, 384, 768], + hooks=hooks, + use_readout=use_readout, + start_index=2, + ) + + +def _make_vit_b_rn50_backbone( + model, + features=[256, 512, 768, 768], + size=[384, 384], + hooks=[0, 1, 8, 11], + vit_features=768, + use_vit_only=False, + use_readout='ignore', + start_index=1, +): + pretrained = nn.Module() + + pretrained.model = model + + if use_vit_only is True: + pretrained.model.blocks[hooks[0]].register_forward_hook( + get_activation('1')) + pretrained.model.blocks[hooks[1]].register_forward_hook( + get_activation('2')) + else: + pretrained.model.patch_embed.backbone.stages[0].register_forward_hook( + get_activation('1')) + pretrained.model.patch_embed.backbone.stages[1].register_forward_hook( + get_activation('2')) + + pretrained.model.blocks[hooks[2]].register_forward_hook( + get_activation('3')) + pretrained.model.blocks[hooks[3]].register_forward_hook( + get_activation('4')) + + pretrained.activations = activations + + readout_oper = get_readout_oper(vit_features, features, use_readout, + start_index) + + if use_vit_only is True: + pretrained.act_postprocess1 = nn.Sequential( + readout_oper[0], + Transpose(1, 2), + nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), + nn.Conv2d( + in_channels=vit_features, + out_channels=features[0], + kernel_size=1, + stride=1, + padding=0, + ), + nn.ConvTranspose2d( + in_channels=features[0], + out_channels=features[0], + kernel_size=4, + stride=4, + padding=0, + bias=True, + dilation=1, + groups=1, + ), + ) + + pretrained.act_postprocess2 = nn.Sequential( + readout_oper[1], + Transpose(1, 2), + nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), + nn.Conv2d( + in_channels=vit_features, + out_channels=features[1], + kernel_size=1, + stride=1, + padding=0, + ), + nn.ConvTranspose2d( + in_channels=features[1], + out_channels=features[1], + kernel_size=2, + stride=2, + padding=0, + bias=True, + dilation=1, + groups=1, + ), + ) + else: + pretrained.act_postprocess1 = nn.Sequential(nn.Identity(), + nn.Identity(), + nn.Identity()) + pretrained.act_postprocess2 = nn.Sequential(nn.Identity(), + nn.Identity(), + nn.Identity()) + + pretrained.act_postprocess3 = nn.Sequential( + readout_oper[2], + Transpose(1, 2), + nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), + nn.Conv2d( + in_channels=vit_features, + out_channels=features[2], + kernel_size=1, + stride=1, + padding=0, + ), + ) + + pretrained.act_postprocess4 = nn.Sequential( + readout_oper[3], + Transpose(1, 2), + nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), + nn.Conv2d( + in_channels=vit_features, + out_channels=features[3], + kernel_size=1, + stride=1, + padding=0, + ), + nn.Conv2d( + in_channels=features[3], + out_channels=features[3], + kernel_size=3, + stride=2, + padding=1, + ), + ) + + pretrained.model.start_index = start_index + pretrained.model.patch_size = [16, 16] + + # We inject this function into the VisionTransformer instances so that + # we can use it with interpolated position embeddings without modifying the library source. + pretrained.model.forward_flex = types.MethodType(forward_flex, + pretrained.model) + + # We inject this function into the VisionTransformer instances so that + # we can use it with interpolated position embeddings without modifying the library source. + pretrained.model._resize_pos_embed = types.MethodType( + _resize_pos_embed, pretrained.model) + + return pretrained + + +def _make_pretrained_vitb_rn50_384(pretrained, + use_readout='ignore', + hooks=None, + use_vit_only=False): + model = timm.create_model('vit_base_resnet50_384', pretrained=pretrained) + + hooks = [0, 1, 8, 11] if hooks is None else hooks + return _make_vit_b_rn50_backbone( + model, + features=[256, 512, 768, 768], + size=[384, 384], + hooks=hooks, + use_vit_only=use_vit_only, + use_readout=use_readout, + ) diff --git a/wan/text2video.py b/wan/text2video.py index befe139..f86284e 100644 --- a/wan/text2video.py +++ b/wan/text2video.py @@ -207,18 +207,19 @@ class WanT2V: 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): + def prepare_source(self, src_video, src_mask, src_ref_images, num_frames, image_size, device, original_video = False, 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 is [-1, 1], 0 = inpainting area (in fact 127 in [0, 255]) + # src_mask is [-1, 1], 0 = preserve original video (in fact 127 in [0, 255]) and 1 = Inpainting (in fact 255 in [0, 255]) src_video[i] = src_video[i].to(device) src_mask[i] = src_mask[i].to(device) 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: @@ -228,10 +229,11 @@ class WanT2V: 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_mask[i] = torch.zeros_like(src_video[i], device=device) if original_video else torch.ones_like(src_video[i], device=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) + 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) image_sizes.append(src_video[i].shape[2:]) for i, ref_images in enumerate(src_ref_images): diff --git a/wan/utils/utils.py b/wan/utils/utils.py index d4e237d..ce4ecd5 100644 --- a/wan/utils/utils.py +++ b/wan/utils/utils.py @@ -21,6 +21,30 @@ __all__ = ['cache_video', 'cache_image', 'str2bool'] from PIL import Image + +def resample(video_fps, video_frames_count, max_frames, target_fps): + import math + + video_frame_duration = 1 /video_fps + target_frame_duration = 1 / target_fps + + cur_time = 0 + target_time = 0 + frame_no = 0 + frame_ids =[] + while True: + if max_frames != 0 and len(frame_ids) >= max_frames: + break + 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 frame_no >= video_frames_count -1: + break + frame_ids = frame_ids[:video_frames_count] + return frame_ids + def get_video_frame(file_name, frame_no): decord.bridge.set_bridge('torch') reader = decord.VideoReader(file_name) diff --git a/wan/utils/vace_preprocessor.py b/wan/utils/vace_preprocessor.py index 912ae39..3bfe885 100644 --- a/wan/utils/vace_preprocessor.py +++ b/wan/utils/vace_preprocessor.py @@ -180,26 +180,17 @@ class VaceVideoProcessor(object): ), 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_frames_count = len(frame_timestamps) - 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 frame_no >= video_frames_count -1: - break - frame_ids = frame_ids[:video_frames_count] + def _get_frameid_bbox_adjust_last(self, fps, video_frames_count, h, w, crop_box, rng, max_frames= 0): + from wan.utils.utils import resample + + target_fps = self.max_fps + + # video_frames_count = len(frame_timestamps) + + frame_ids= resample(fps, video_frames_count, max_frames, target_fps) + 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 @@ -235,11 +226,11 @@ class VaceVideoProcessor(object): 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): + def _get_frameid_bbox(self, fps, video_frames_count, 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) + return self._get_frameid_bbox_adjust_last(fps, video_frames_count, 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) + return self._get_frameid_bbox_default(fps, video_frames_count, 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) @@ -253,23 +244,37 @@ class VaceVideoProcessor(object): import decord decord.bridge.set_bridge('torch') readers = [] + src_video = None for data_k in data_key_batch: - reader = decord.VideoReader(data_k) - readers.append(reader) + if torch.is_tensor(data_k): + src_video = data_k + else: + 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] + if src_video != None: + fps = 16 + length = src_video.shape[1] + if len(readers) > 0: + min_readers = min([len(r) for r in readers]) + length = min(length, min_readers ) + else: + 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) 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) + if src_video != None: + src_video = src_video[:max_frames] + h, w = src_video.shape[1:3] + else: + h, w = readers[0].next().shape[:2] + frame_ids, (x1, x2, y1, y2), (oh, ow), fps = self._get_frameid_bbox(fps, length, 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] + if src_video != None: + videos = [src_video] + videos 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] diff --git a/wgp.py b/wgp.py index 083ee1e..1115bc9 100644 --- a/wgp.py +++ b/wgp.py @@ -141,12 +141,27 @@ def process_prompt_and_add_tasks(state, model_choice): 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 video_prompt_type: + if "I" in video_prompt_type: + if image_refs == None: + gr.Info("You must provide at one Refererence Image") + return + else: image_refs = None - if not "V" in video_prompt_type: + if "V" in video_prompt_type: + if video_guide == None: + gr.Info("You must provide a Control Video") + return + else: video_guide = None - if not "M" in video_prompt_type: + if "M" in video_prompt_type: + if video_mask == None: + gr.Info("You must provide a Video Mask ") + return + else: video_mask = None + if "O" in video_prompt_type and inputs["max_frames"]==0: + gr.Info(f"In order to extend a video, you need to indicate how many frames you want to reuse in the source video.") + return if isinstance(image_refs, list): image_refs = [ convert_image(tup[0]) for tup in image_refs ] @@ -260,7 +275,7 @@ def add_video_task(**inputs): queue = gen["queue"] task_id += 1 current_task_id = task_id - inputs_to_query = ["image_start", "image_end", "image_refs", "video_guide", "video_mask"] + inputs_to_query = ["image_start", "image_end", "video_guide", "image_refs","video_mask"] start_image_data = None end_image_data = None for name in inputs_to_query: @@ -718,7 +733,7 @@ if not Path(server_config_filename).is_file(): "transformer_types": [], "transformer_quantization": "int8", "text_encoder_filename" : text_encoder_choices[1], - "save_path": os.path.join(os.getcwd(), "gradio_outputs"), + "save_path": "outputs", #os.path.join(os.getcwd(), "compile" : "", "metadata_type": "metadata", "default_ui": "t2v", @@ -726,7 +741,7 @@ if not Path(server_config_filename).is_file(): "clear_file_list" : 0, "vae_config": 0, "profile" : profile_type.LowRAM_LowVRAM, - "reload_model": 2 } + "preload_model_policy": [] } with open(server_config_filename, "w", encoding="utf-8") as writer: writer.write(json.dumps(server_config)) @@ -860,7 +875,7 @@ if len(args.vae_config) > 0: reload_needed = False default_ui = server_config.get("default_ui", "t2v") save_path = server_config.get("save_path", os.path.join(os.getcwd(), "gradio_outputs")) -reload_model = server_config.get("reload_model", 2) +preload_model_policy = server_config.get("preload_model_policy", []) if args.t2v_14B or args.t2v: @@ -962,8 +977,8 @@ def download_models(transformer_filename, text_encoder_filename): from huggingface_hub import hf_hub_download, snapshot_download repoId = "DeepBeepMeep/Wan2.1" - sourceFolderList = ["xlm-roberta-large", "", ] - fileList = [ [], ["Wan2.1_VAE_bf16.safetensors", "models_clip_open-clip-xlm-roberta-large-vit-huge-14-bf16.safetensors", "flownet.pkl" ] + computeList(text_encoder_filename) + computeList(transformer_filename) ] + sourceFolderList = ["xlm-roberta-large", "pose", "depth", "", ] + fileList = [ [], [],[], ["Wan2.1_VAE_bf16.safetensors", "models_clip_open-clip-xlm-roberta-large-vit-huge-14-bf16.safetensors", "flownet.pkl" ] + computeList(text_encoder_filename) + computeList(transformer_filename) ] targetRoot = "ckpts/" for sourceFolder, files in zip(sourceFolderList,fileList ): if len(files)==0: @@ -1166,7 +1181,7 @@ def load_models(model_filename): return wan_model, offloadobj, pipe["transformer"] -if reload_model ==3 or reload_model ==4: +if not "P" in preload_model_policy: wan_model, offloadobj, transformer = None, None, None reload_needed = True else: @@ -1254,7 +1269,7 @@ def apply_changes( state, quantization_choice, boost_choice = 1, clear_file_list = 0, - reload_choice = 1, + preload_model_policy_choice = 1, ): if args.lock_config: return @@ -1272,7 +1287,7 @@ def apply_changes( state, "transformer_quantization" : quantization_choice, "boost" : boost_choice, "clear_file_list" : clear_file_list, - "reload_model" : reload_choice, + "preload_model_policy" : preload_model_policy_choice, } if Path(server_config_filename).is_file(): @@ -1295,14 +1310,14 @@ def apply_changes( state, if v != v_old: changes.append(k) - global attention_mode, profile, compile, transformer_filename, text_encoder_filename, vae_config, boost, lora_dir, reload_needed, reload_model, transformer_quantization, transformer_types + global attention_mode, profile, compile, transformer_filename, text_encoder_filename, vae_config, boost, lora_dir, reload_needed, preload_model_policy, transformer_quantization, transformer_types attention_mode = server_config["attention_mode"] profile = server_config["profile"] compile = server_config["compile"] text_encoder_filename = server_config["text_encoder_filename"] vae_config = server_config["vae_config"] boost = server_config["boost"] - reload_model = server_config["reload_model"] + preload_model_policy = server_config["preload_model_policy"] transformer_quantization = server_config["transformer_quantization"] transformer_types = server_config["transformer_types"] transformer_type = get_model_type(transformer_filename) @@ -1381,7 +1396,8 @@ def abort_generation(state): gen["abort"] = True gen["extra_orders"] = 0 - wan_model._interrupt= True + if wan_model != None: + wan_model._interrupt= True msg = "Processing Request to abort Current Generation" gr.Info(msg) return msg, gr.Button(interactive= False) @@ -1480,24 +1496,68 @@ def expand_slist(slist, num_inference_steps ): return new_slist def convert_image(image): - from PIL import ExifTags, ImageOps + from PIL import 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 preprocess_video(process_type, height, width, video_in, max_frames): + + from wan.utils.utils import resample + + import decord + decord.bridge.set_bridge('torch') + reader = decord.VideoReader(video_in) + + fps = reader.get_avg_fps() + + frame_nos = resample(fps, len(reader), max_frames= max_frames, target_fps=16) + frames_list = reader.get_batch(frame_nos) + frame_height, frame_width, _ = frames_list[0].shape + + scale = ((height * width ) / (frame_height * frame_width))**(1/2) + # scale = min(height / frame_height, width / frame_width) + + new_height = (int(frame_height * scale) // 16) * 16 + new_width = (int(frame_width * scale) // 16) * 16 + + processed_frames_list = [] + for frame in frames_list: + frame = Image.fromarray(np.clip(frame.cpu().numpy(), 0, 255).astype(np.uint8)) + frame = frame.resize((new_width,new_height), resample=Image.Resampling.LANCZOS) + processed_frames_list.append(frame) + + if process_type=="pose": + from preprocessing.dwpose.pose import PoseBodyFaceVideoAnnotator + cfg_dict = { + "DETECTION_MODEL": "ckpts/pose/yolox_l.onnx", + "POSE_MODEL": "ckpts/pose/dw-ll_ucoco_384.onnx", + "RESIZE_SIZE": 1024 + } + anno_ins = PoseBodyFaceVideoAnnotator(cfg_dict) + elif process_type=="depth": + from preprocessing.midas.depth import DepthVideoAnnotator + cfg_dict = { + "PRETRAINED_MODEL": "ckpts/depth/dpt_hybrid-midas-501f0c75.pt" + } + anno_ins = DepthVideoAnnotator(cfg_dict) + else: + from preprocessing.gray import GrayVideoAnnotator + cfg_dict = {} + anno_ins = GrayVideoAnnotator(cfg_dict) + + np_frames = anno_ins.forward(processed_frames_list) + + # from preprocessing.dwpose.pose import save_one_video + # save_one_video("test.mp4", np_frames, fps=8, quality=8, macro_block_size=None) + + torch_frames = [] + for np_frame in np_frames: + torch_frame = torch.from_numpy(np_frame) + torch_frames.append(torch_frame) + + return torch.stack(torch_frames) def generate_video( task_id, @@ -1551,7 +1611,7 @@ def generate_video( # gr.Info("Unable to generate a Video while a new configuration is being applied.") # return - if reload_model !=3 and reload_model !=4 : + if "P" in preload_model_policy: while wan_model == None: time.sleep(1) @@ -1681,10 +1741,32 @@ def generate_video( raise gr.Error("Teacache not supported for this model") if "Vace" in model_filename: + # video_prompt_type = video_prompt_type +"G" + if any(process in video_prompt_type for process in ("P", "D", "G")) : + prompts_max = gen["prompts_max"] + + status = get_generation_status(prompt_no, prompts_max, 1, 1) + preprocess_type = None + if "P" in video_prompt_type : + progress_args = [0, status + " - Extracting Open Pose Information"] + preprocess_type = "pose" + elif "D" in video_prompt_type : + progress_args = [0, status + " - Extracting Depth Information"] + preprocess_type = "depth" + elif "G" in video_prompt_type : + progress_args = [0, status + " - Extracting Gray Level Information"] + preprocess_type = "gray" + + if preprocess_type != None : + progress(*progress_args ) + gen["progress_args"] = progress_args + video_guide = preprocess_video(preprocess_type, width=width, height=height,video_in=video_guide, max_frames= video_length) + src_video, src_mask, src_ref_images = wan_model.prepare_source([video_guide], [video_mask], [image_refs], video_length, VACE_SIZE_CONFIGS[resolution_reformated], "cpu", + original_video= "O" in video_prompt_type, trim_video=max_frames) else: src_video, src_mask, src_ref_images = None, None, None @@ -2539,9 +2621,9 @@ def fill_inputs(state): return generate_video_tab(update_form = True, state_dict = state, ui_defaults = ui_defaults) -def preload_model(state): +def preload_model_when_switching(state): global reload_needed, wan_model, offloadobj - if reload_model == 1: + if "S" in preload_model_policy: model_filename = state["model_filename"] if state["model_filename"] != transformer_filename: wan_model = None @@ -2558,7 +2640,7 @@ def preload_model(state): def unload_model_if_needed(state): global reload_needed, wan_model, offloadobj - if reload_model == 4: + if "U" in preload_model_policy: if wan_model != None: wan_model = None if offloadobj is not None: @@ -2567,7 +2649,39 @@ def unload_model_if_needed(state): gc.collect() reload_needed= True +def filter_letters(source_str, letters): + ret = "" + for letter in letters: + if letter in source_str: + ret += letter + return ret +def add_to_sequence(source_str, letters): + ret = source_str + for letter in letters: + if not letter in source_str: + ret += letter + return ret + +def del_in_sequence(source_str, letters): + ret = source_str + for letter in letters: + if letter in source_str: + ret = ret.replace(letter, "") + return ret + + +def refresh_video_prompt_type_image_refs(video_prompt_type, video_prompt_type_image_refs): + video_prompt_type = add_to_sequence(video_prompt_type, "I") if video_prompt_type_image_refs else del_in_sequence(video_prompt_type, "I") + return video_prompt_type, gr.update(visible = video_prompt_type_image_refs),gr.update(visible = video_prompt_type_image_refs) + +def refresh_video_prompt_type_video_guide(video_prompt_type, video_prompt_type_video_guide): + video_prompt_type = del_in_sequence(video_prompt_type, "ODPCMV") + video_prompt_type = add_to_sequence(video_prompt_type, video_prompt_type_video_guide) + visible = "V" in video_prompt_type + return video_prompt_type, gr.update(visible = visible), gr.update(visible = visible), gr.update(visible= "M" in video_prompt_type ) + + def generate_video_tab(update_form = False, state_dict = None, ui_defaults = None, model_choice = None, header = None): global inputs_names #, advanced @@ -2676,19 +2790,36 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non image_end = gr.Image(label= "Last Image for a new video", type ="pil", visible="E" in image_prompt_type_value, value= ui_defaults.get("image_end", None)) with gr.Column(visible= "Vace" in model_filename ) as video_prompt_column: - gr.Markdown("Control conditions: Images References (custom Faces or Objects), Video (Open Pose, Depth maps), Mask (inpainting)") - video_prompt_type_value= ui_defaults.get("video_prompt_type","I") - video_prompt_type = gr.Radio( [("Images Ref", "I"),("a Video", "V"), ("Images Refs + a Video", "IV"), ("Video + Video Mask", "VM"), ("Images + Video + Mask", "IVM")], value =video_prompt_type_value, label="Location", show_label= False, scale= 3) - image_refs = gr.Gallery( - label="Images Referencse (Custom faces and Objects 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 video_prompt_type_value, value= ui_defaults.get("image_refs", None) ) + video_prompt_type_value= ui_defaults.get("video_prompt_type","") + video_prompt_type = gr.Text(value= video_prompt_type_value, visible= False) + video_prompt_type_video_guide = gr.Dropdown( + choices=[ + ("None, use only the Text Prompt", ""), + ("Extend the Control Video", "OV"), + ("Transfer Human Motion from the Control Video", "PV"), + ("Transfer Depth from the Control Video", "DV"), + ("Recolorize the Control Video", "CV"), + ("Control Video contains Open Pose, Depth or Black & White ", "V"), + ("Inpainting of Control Video using Mask Video ", "MV"), + ], + value=filter_letters(video_prompt_type_value, "ODPCMV"), + label="Video to Video" + ) + video_prompt_type_image_refs = gr.Checkbox(value="I" in video_prompt_type_value , label= "Use References Images (Faces, Objects) to customize New Video", scale =1 ) - video_guide = gr.Video(label= "Reference Video (an animated Video in the Open Pose format or Depth Map video)", visible= "V" in video_prompt_type_value, value= ui_defaults.get("video_guide", None) ) - with gr.Row(): - max_frames = gr.Slider(0, 100, value=ui_defaults.get("max_frames",0), step=1, label="Nb of frames in Ref. Video (0 = as many as possible)", visible= "V" in video_prompt_type_value, 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 video_prompt_type_value, scale =1 ) + video_guide = gr.Video(label= "Control Video", visible= "V" in video_prompt_type_value, value= ui_defaults.get("video_guide", None),) + max_frames = gr.Slider(0, 100, value=ui_defaults.get("max_frames",0), step=1, label="Nb of frames in Control Video to use (0 = max)", visible= "V" in video_prompt_type_value, scale = 2 ) - video_mask = gr.Video(label= "Video Mask (for Inpainting or Outpaing, white pixels = Mask)", visible= "M" in video_prompt_type_value, value= ui_defaults.get("video_mask", None) ) + image_refs = gr.Gallery( label ="Reference Images", + type ="pil", show_label= True, + columns=[3], rows=[1], object_fit="contain", height="auto", selected_index=0, interactive= True, visible= "I" in video_prompt_type_value, + value= ui_defaults.get("image_refs", None) ) + + # with gr.Row(): + remove_background_image_ref = gr.Checkbox(value=ui_defaults.get("remove_background_image_ref",1), label= "Remove Background of Images References", visible= "I" in video_prompt_type_value, scale =1 ) + + + video_mask = gr.Video(label= "Video Mask (for Inpainting or Outpaing, white pixels = Mask)", visible= "M" in video_prompt_type_value, value= ui_defaults.get("video_mask", None)) advanced_prompt = advanced_ui @@ -2923,7 +3054,10 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non target_settings = gr.Text(value = "settings", interactive= False, visible= False) image_prompt_type.change(fn=refresh_image_prompt_type, inputs=[state, image_prompt_type], outputs=[image_start, image_end]) - video_prompt_type.change(fn=refresh_video_prompt_type, inputs=[state, video_prompt_type], outputs=[image_refs, video_guide, video_mask, max_frames, remove_background_image_ref]) + # video_prompt_type.change(fn=refresh_video_prompt_type, inputs=[state, video_prompt_type], outputs=[image_refs, video_guide, video_mask, max_frames, remove_background_image_ref]) + video_prompt_type_image_refs.input(fn=refresh_video_prompt_type_image_refs, inputs = [video_prompt_type, video_prompt_type_image_refs], outputs = [video_prompt_type, image_refs, remove_background_image_ref ]) + video_prompt_type_video_guide.input(fn=refresh_video_prompt_type_video_guide, inputs = [video_prompt_type, video_prompt_type_video_guide], outputs = [video_prompt_type, video_guide, max_frames, video_mask]) + 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]) queue_df.select( fn=handle_celll_selection, inputs=state, outputs=[queue_df, modal_image_display, modal_container]) @@ -2967,7 +3101,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ).then(fn= fill_inputs, inputs=[state], outputs=gen_inputs + extra_inputs - ).then(fn= preload_model, + ).then(fn= preload_model_when_switching, inputs=[state], outputs=[gen_status]) @@ -3149,14 +3283,8 @@ def generate_configuration_tab(header, model_choice): value=server_config.get("metadata_type", "metadata"), label="Metadata Handling" ) - reload_choice = gr.Dropdown( - choices=[ - ("Load Model When Starting the App and Changing Model if Model Changed", 1), - ("Load Model When Starting the App and Pressing Generate if Model Changed", 2), - ("Load Model When Pressing Generate if Model Changed", 3), - ("Load Model When Pressing Generate and Unload Model when Finished", 4), - ], - value=server_config.get("reload_model",2), + preload_model_policy_choice = gr.CheckboxGroup([("Preload Model while Launching the App","P"), ("Preload Model while Switching Model", "S"), ("Unload Model when Queue is Done", "U")], + value=server_config.get("preload_model_policy",[]), label="RAM Loading / Unloading Model Policy (in any case VRAM will be freed once the queue has been processed)" ) @@ -3191,7 +3319,7 @@ def generate_configuration_tab(header, model_choice): quantization_choice, boost_choice, clear_file_list_choice, - reload_choice, + preload_model_policy_choice, ], outputs= [msg , header, model_choice] ) @@ -3201,10 +3329,16 @@ def generate_about_tab(): gr.Markdown("Original Wan 2.1 Model by Alibaba (GitHub)") gr.Markdown("Many thanks to:") gr.Markdown("- Alibaba Wan team for the best open source video generator") + gr.Markdown("- Alibaba Vace and Fun Teams for their incredible control net models") gr.Markdown("- Cocktail Peanuts : QA and simple installation via Pinokio.computer") gr.Markdown("- Tophness : created multi tabs and queuing frameworks") gr.Markdown("- AmericanPresidentJimmyCarter : added original support for Skip Layer Guidance") - gr.Markdown("- Remade_AI : for creating their awesome Loras collection") + gr.Markdown("- Remade_AI : for their awesome Loras collection") + gr.Markdown("
Huge acknowlegments to these great open source projects used in WanGP:") + gr.Markdown("- Rife: temporal upsampler (https://github.com/hzwer/ECCV2022-RIFE)") + gr.Markdown("- DwPose: Open Pose extractor (https://github.com/IDEA-Research/DWPose)") + gr.Markdown("- Midas: Depth extractor (https://github.com/isl-org/MiDaS") + def generate_info_tab(): gr.Markdown("Welcome to WanGP a super fast and low VRAM AI Video Generator !") @@ -3231,17 +3365,26 @@ def generate_dropdown_model_list(): choices= dropdown_choices, value= current_model_type, show_label= False, - scale= 2 + scale= 2, + elem_id="model_list", + elem_classes="model_list_class", ) def create_demo(): css = """ + #model_list{ + background-color:black; + padding:1px} + + #model_list input { + font-size:25px} + .title-with-lines { display: flex; align-items: center; - margin: 30px 0; + margin: 25px 0; } .line { flex-grow: 1; @@ -3462,7 +3605,7 @@ def create_demo(): pointer-events: none; } """ - with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="sky", neutral_hue="slate", text_size="md"), title= "Wan2GP") as demo: + with gr.Blocks(css=css, theme=gr.themes.Soft(font=["Verdana"], primary_hue="sky", neutral_hue="slate", text_size="md"), title= "Wan2GP") as demo: gr.Markdown("

WanGP v4.0 by DeepBeepMeep ") # (Updates)

") global model_list From 016e57c2536ed03b635d2ea8aa3a752390a3c342 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Wed, 9 Apr 2025 16:23:26 +0200 Subject: [PATCH 44/69] Fixed bug related to preload policy --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 8531620..fa9d69f 100644 --- a/README.md +++ b/README.md @@ -288,7 +288,7 @@ This offers a stronger mechanism to tell Vace which parts should be kept (black) Examples: - Inject people and / objects into a scene describe by a text prompt: Ref. Images + text Prompt - Animate a character described in a text prompt: a Video of person moving + text Prompt -- Animate a character of your choice (pose transfer) : Ref Images + a Video of person moving + text Prompt +- Animate a character of your choice (motion transfer) : Ref Images + a Video of person moving + text Prompt - Change the style of a scene (depth transfer): a Video that contains objects / person at differen depths + text Prompt From 452d246c8819c4859f1de94148a96de5c51c86d8 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Thu, 10 Apr 2025 04:22:35 +1000 Subject: [PATCH 45/69] move load/save/clear buttons to accordion --- wgp.py | 49 ++++++++++++++++++++++++++++++------------------- 1 file changed, 30 insertions(+), 19 deletions(-) diff --git a/wgp.py b/wgp.py index 97cd830..785fff2 100644 --- a/wgp.py +++ b/wgp.py @@ -3479,7 +3479,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non with gr.Row(): onemore_btn = gr.Button("One More Sample Please !") abort_btn = gr.Button("Abort") - + with gr.Accordion("Queue Management", open=False) as queue_accordion: queue_df = gr.DataFrame( headers=["Qty","Prompt", "Length","Steps","", "", "", "", ""], datatype=[ "str","markdown","str", "markdown", "markdown", "markdown", "str", "str", "str"], @@ -3489,14 +3489,14 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non wrap=True, value=[], line_breaks= True, - visible= False, + visible= True, elem_id="queue_df" ) - with gr.Row(): - queue_zip_base64_output = gr.Text(visible=False) - save_queue_btn = gr.DownloadButton("Save Queue", size="sm") - load_queue_btn = gr.UploadButton("Load Queue", file_types=[".zip"], size="sm") - clear_queue_btn = gr.Button("Clear Queue", size="sm", variant="stop") + with gr.Row(): + queue_zip_base64_output = gr.Text(visible=False) + save_queue_btn = gr.DownloadButton("Save Queue", size="sm") + load_queue_btn = gr.UploadButton("Load Queue", file_types=[".zip"], size="sm") + clear_queue_btn = gr.Button("Clear Queue", size="sm", variant="stop") trigger_zip_download_js = """ (base64String) => { if (!base64String) { @@ -3544,9 +3544,9 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non inputs=[load_queue_btn, state], outputs=[queue_df] ).then( - fn=lambda s: gr.update(visible=bool(get_gen_info(s).get("queue",[]))), + fn=lambda s: (gr.update(visible=bool(get_gen_info(s).get("queue",[]))), gr.Accordion(open=True)) if bool(get_gen_info(s).get("queue",[])) else (gr.update(visible=False), gr.update()), inputs=[state], - outputs=[current_gen_column] + outputs=[current_gen_column, queue_accordion] ) clear_queue_btn.click( @@ -3554,9 +3554,9 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non inputs=[state], outputs=[queue_df] ).then( - fn=lambda: gr.update(visible=False), + fn=lambda: (gr.update(visible=False), gr.Accordion(open=False)), inputs=None, - outputs=[current_gen_column] + outputs=[current_gen_column, queue_accordion] ) extra_inputs = prompt_vars + [wizard_prompt, wizard_variables_var, wizard_prompt_activated_var, video_prompt_column, image_prompt_column, @@ -3639,6 +3639,10 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ).then(finalize_generation, inputs= [state], outputs= [output, abort_btn, generate_btn, add_to_queue_btn, current_gen_column, gen_info] + ).then( + fn=lambda s: gr.Accordion(open=False) if len(get_gen_info(s).get("queue", [])) <= 1 else gr.update(), + inputs=[state], + outputs=[queue_accordion] ).then(unload_model_if_needed, inputs= [state], outputs= [] @@ -3653,6 +3657,10 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ).then(fn=process_prompt_and_add_tasks, inputs = [state, model_choice], outputs=queue_df + ).then( + fn=lambda s: gr.Accordion(open=True) if len(get_gen_info(s).get("queue", [])) > 1 else gr.update(), # Expand if queue has items (len > 1 assumes placeholder) + inputs=[state], + outputs=[queue_accordion] ).then( fn=update_status, inputs = [state], @@ -3670,7 +3678,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non gen_info, prompt, wizard_prompt, wizard_prompt_activated_var, wizard_variables_var, prompt_column_advanced, prompt_column_wizard, prompt_column_wizard_vars, - advanced_row, image_prompt_column, video_prompt_column, + advanced_row, image_prompt_column, video_prompt_column, queue_accordion, *prompt_vars ) @@ -4153,7 +4161,7 @@ def create_demo(): gen_info, prompt, wizard_prompt, wizard_prompt_activated_var, wizard_variables_var, prompt_column_advanced, prompt_column_wizard, prompt_column_wizard_vars, - advanced_row, image_prompt_column, video_prompt_column, + advanced_row, image_prompt_column, video_prompt_column, queue_accordion, *prompt_vars_outputs ) = generate_video_tab(model_choice=model_choice, header=header) with gr.Tab("Informations"): @@ -4170,7 +4178,8 @@ def create_demo(): def run_autoload_and_prepare_ui(current_state): df_update, loaded_flag, modified_state = autoload_queue(current_state) should_start_processing = loaded_flag - return df_update, gr.update(visible=loaded_flag), should_start_processing, modified_state + accordion_update = gr.Accordion(open=True) if loaded_flag else gr.update() + return df_update, gr.update(visible=loaded_flag), accordion_update, should_start_processing, modified_state def start_processing_if_needed(should_start, current_state): if not isinstance(current_state, dict) or 'gen' not in current_state: @@ -4179,18 +4188,20 @@ def create_demo(): if should_start: yield from process_tasks(current_state) else: - yield "Autoload complete. Processing not started." + yield None def finalize_generation_with_state(current_state): if not isinstance(current_state, dict) or 'gen' not in current_state: - return gr.update(), gr.update(interactive=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False, value=""), current_state + return gr.update(), gr.update(interactive=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False, value=""), gr.update(), current_state + gallery_update, abort_btn_update, gen_btn_update, add_queue_btn_update, current_gen_col_update, gen_info_update = finalize_generation(current_state) - return gallery_update, abort_btn_update, gen_btn_update, add_queue_btn_update, current_gen_col_update, gen_info_update, current_state + accordion_update = gr.Accordion(open=False) if len(get_gen_info(current_state).get("queue", [])) <= 1 else gr.update() + return gallery_update, abort_btn_update, gen_btn_update, add_queue_btn_update, current_gen_col_update, gen_info_update, accordion_update, current_state demo.load( fn=run_autoload_and_prepare_ui, inputs=[state], - outputs=[queue_df, current_gen_column, should_start_flag, state] + outputs=[queue_df, current_gen_column, queue_accordion, should_start_flag, state] ).then( fn=start_processing_if_needed, inputs=[should_start_flag, state], @@ -4199,7 +4210,7 @@ def create_demo(): ).then( fn=finalize_generation_with_state, inputs=[state], - outputs=[output, abort_btn, generate_btn, add_to_queue_btn, current_gen_column, gen_info, state], + outputs=[output, abort_btn, generate_btn, add_to_queue_btn, current_gen_column, gen_info, queue_accordion, state], trigger_mode="always_last" ) From 965267d74004d71bd2b23a0ec87ce7aaff81b79c Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Thu, 10 Apr 2025 04:37:27 +1000 Subject: [PATCH 46/69] clear queue aborts currently generating item --- wgp.py | 156 ++++++++++++++++++++++----------------------------------- 1 file changed, 60 insertions(+), 96 deletions(-) diff --git a/wgp.py b/wgp.py index 785fff2..a2a468f 100644 --- a/wgp.py +++ b/wgp.py @@ -362,24 +362,21 @@ def save_queue_action(state): gen = get_gen_info(state) queue = gen.get("queue", []) - if not queue or len(queue) <=1 : # Check if queue is empty or only has the placeholder + if not queue or len(queue) <=1 : gr.Info("Queue is empty. Nothing to save.") - return None # Return None if nothing to save + return None - # Use an in-memory buffer for the zip file zip_buffer = io.BytesIO() - # Still use a temporary directory *only* for storing images before zipping with tempfile.TemporaryDirectory() as tmpdir: queue_manifest = [] - image_paths_in_zip = {} # Tracks image PIL object ID -> filename in zip + image_paths_in_zip = {} for task_index, task in enumerate(queue): - # Skip the placeholder item if it exists if task is None or not isinstance(task, dict) or task_index == 0: continue params_copy = task.get('params', {}).copy() - task_id_s = task.get('id', f"task_{task_index}") # Use a different var name + task_id_s = task.get('id', f"task_{task_index}") image_keys = ["image_start", "image_end", "image_refs"] for key in image_keys: @@ -387,95 +384,71 @@ def save_queue_action(state): if images_pil is None: continue - # Ensure images_pil is always a list for processing is_originally_list = isinstance(images_pil, list) if not is_originally_list: images_pil = [images_pil] image_filenames_for_json = [] for img_index, pil_image in enumerate(images_pil): - # Ensure it's actually a PIL Image object before proceeding if not isinstance(pil_image, Image.Image): print(f"Warning: Expected PIL Image for key '{key}' in task {task_id_s}, got {type(pil_image)}. Skipping image.") continue - # Use object ID to check if this specific image instance is already saved img_id = id(pil_image) if img_id in image_paths_in_zip: - # If already saved, just add its filename to the list image_filenames_for_json.append(image_paths_in_zip[img_id]) - continue # Move to the next image in the list + continue - # Image not saved yet, create filename and save path img_filename_in_zip = f"task{task_id_s}_{key}_{img_index}.png" img_save_path = os.path.join(tmpdir, img_filename_in_zip) try: - # Save the image to the temporary directory pil_image.save(img_save_path, "PNG") image_filenames_for_json.append(img_filename_in_zip) - # Store the mapping from image ID to its filename in the zip image_paths_in_zip[img_id] = img_filename_in_zip except Exception as e: print(f"Error saving image {img_filename_in_zip} for task {task_id_s}: {e}") - # Optionally decide if you want to continue or fail here - # Update the params_copy with the list of filenames (or single filename) if image_filenames_for_json: params_copy[key] = image_filenames_for_json if is_originally_list else image_filenames_for_json[0] else: - # If no images were successfully processed for this key, remove it params_copy.pop(key, None) - # Clean up parameters before adding to manifest params_copy.pop('state', None) - params_copy.pop('start_image_data_base64', None) # Don't need base64 in saved queue + params_copy.pop('start_image_data_base64', None) params_copy.pop('end_image_data_base64', None) - # Also remove the actual PIL data if it somehow remained params_copy.pop('start_image_data', None) params_copy.pop('end_image_data', None) manifest_entry = { "id": task.get('id'), "params": params_copy, - # Keep other necessary top-level task info if needed, like repeats etc. - # Example: "repeats": task.get('repeats', 1) } queue_manifest.append(manifest_entry) - # --- Create queue.json content --- manifest_path = os.path.join(tmpdir, "queue.json") try: with open(manifest_path, 'w', encoding='utf-8') as f: - # Dump only the relevant manifest data json.dump(queue_manifest, f, indent=4) except Exception as e: print(f"Error writing queue.json: {e}") gr.Warning("Failed to create queue manifest.") - return None # Return None on failure + return None - # --- Create the zip file in memory --- try: with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zf: - # Add queue.json zf.write(manifest_path, arcname="queue.json") - # Add all unique images that were saved to the temp dir for saved_img_rel_path in image_paths_in_zip.values(): saved_img_abs_path = os.path.join(tmpdir, saved_img_rel_path) if os.path.exists(saved_img_abs_path): zf.write(saved_img_abs_path, arcname=saved_img_rel_path) else: - # This shouldn't happen if saving was successful, but good to check print(f"Warning: Image file {saved_img_rel_path} not found during zipping.") - # --- Prepare for return --- - # Move buffer position to the beginning zip_buffer.seek(0) - # Read the binary content zip_binary_content = zip_buffer.getvalue() - # Encode as base64 string zip_base64 = base64.b64encode(zip_binary_content).decode('utf-8') print(f"Queue successfully prepared as base64 string ({len(zip_base64)} chars).") return zip_base64 @@ -483,18 +456,17 @@ def save_queue_action(state): except Exception as e: print(f"Error creating zip file in memory: {e}") gr.Warning("Failed to create zip data for download.") - return None # Return None on failure + return None finally: zip_buffer.close() def load_queue_action(filepath, state): global task_id gen = get_gen_info(state) - original_queue = gen.get("queue", []) # Store original queue for error case + original_queue = gen.get("queue", []) if not filepath or not hasattr(filepath, 'name') or not Path(filepath.name).is_file(): print("[load_queue_action] Warning: No valid file selected or file not found.") - # Return the current state of the DataFrame return update_queue_data(original_queue) newly_loaded_queue = [] @@ -518,7 +490,6 @@ def load_queue_action(filepath, state): print(f"[load_queue_action] Manifest loaded. Processing {len(loaded_manifest)} tasks.") for task_index, task_data in enumerate(loaded_manifest): - # (Keep the existing task processing logic here...) if task_data is None or not isinstance(task_data, dict): print(f"[load_queue_action] Skipping invalid task data at index {task_index}") continue @@ -528,7 +499,7 @@ def load_queue_action(filepath, state): max_id_in_file = max(max_id_in_file, task_id_loaded) loaded_pil_images = {} image_keys = ["image_start", "image_end", "image_refs"] - params['state'] = state # Add state back temporarily for consistency if needed by internal logic, but it's removed before saving + params['state'] = state for key in image_keys: image_filenames = params.get(key) @@ -544,26 +515,22 @@ def load_queue_action(filepath, state): continue try: pil_image = Image.open(img_load_path) - # Ensure the image data is loaded into memory before the temp dir is cleaned up pil_image.load() - # Convert image right after loading converted_image = convert_image(pil_image) loaded_pils.append(converted_image) - pil_image.close() # Close the file handle + pil_image.close() except Exception as img_e: print(f"[load_queue_action] Error loading image {img_filename_in_zip}: {img_e}") if loaded_pils: params[key] = loaded_pils if is_list else loaded_pils[0] - loaded_pil_images[key] = params[key] # Store loaded PILs for preview generation + loaded_pil_images[key] = params[key] else: params.pop(key, None) - # Generate preview base64 strings primary_preview_pil, secondary_preview_pil = None, None start_prev_pil_list = loaded_pil_images.get("image_start") end_prev_pil_list = loaded_pil_images.get("image_end") ref_prev_pil_list = loaded_pil_images.get("image_refs") - # Extract first image for preview if available if start_prev_pil_list: primary_preview_pil = start_prev_pil_list[0] if isinstance(start_prev_pil_list, list) and start_prev_pil_list else start_prev_pil_list if not isinstance(start_prev_pil_list, list) else None if end_prev_pil_list: @@ -571,97 +538,102 @@ def load_queue_action(filepath, state): elif ref_prev_pil_list and isinstance(ref_prev_pil_list, list) and ref_prev_pil_list: primary_preview_pil = ref_prev_pil_list[0] - # Generate base64 only if PIL image exists start_b64 = [pil_to_base64_uri(primary_preview_pil, format="jpeg", quality=70)] if primary_preview_pil else None end_b64 = [pil_to_base64_uri(secondary_preview_pil, format="jpeg", quality=70)] if secondary_preview_pil else None - # Get top-level image data (PIL objects) for runtime task top_level_start_image = loaded_pil_images.get("image_start") top_level_end_image = loaded_pil_images.get("image_end") - # Construct the runtime task dictionary runtime_task = { "id": task_id_loaded, - "params": params.copy(), # Use a copy of params - # Extract necessary params for top level if they exist + "params": params.copy(), "repeats": params.get('repeat_generation', 1), "length": params.get('video_length'), "steps": params.get('num_inference_steps'), "prompt": params.get('prompt'), - # Store the actual loaded PIL image data here "start_image_data": top_level_start_image, "end_image_data": top_level_end_image, - # Store base64 previews generated above "start_image_data_base64": start_b64, "end_image_data_base64": end_b64, } newly_loaded_queue.append(runtime_task) print(f"[load_queue_action] Processed task {task_index+1}/{len(loaded_manifest)}, ID: {task_id_loaded}") - # --- State Update --- with lock: print("[load_queue_action] Acquiring lock to update state...") - gen["queue"] = newly_loaded_queue[:] # Replace the queue in the state - local_queue_copy_for_global_ref = gen["queue"][:] # Copy for global ref update - current_max_id_in_new_queue = max([t['id'] for t in newly_loaded_queue if 'id' in t] + [0]) # Safer max ID calculation + gen["queue"] = newly_loaded_queue[:] + local_queue_copy_for_global_ref = gen["queue"][:] + current_max_id_in_new_queue = max([t['id'] for t in newly_loaded_queue if 'id' in t] + [0]) - # Update global task ID only if the loaded max ID is higher if current_max_id_in_new_queue > task_id: print(f"[load_queue_action] Updating global task_id from {task_id} to {current_max_id_in_new_queue + 1}") - task_id = current_max_id_in_new_queue + 1 # Ensure next ID is unique + task_id = current_max_id_in_new_queue + 1 else: print(f"[load_queue_action] Global task_id ({task_id}) is >= max in file ({current_max_id_in_new_queue}). Not changing task_id.") gen["prompts_max"] = len(newly_loaded_queue) print("[load_queue_action] State update complete. Releasing lock.") - # --- Global Reference Update --- if local_queue_copy_for_global_ref is not None: print("[load_queue_action] Updating global queue reference...") update_global_queue_ref(local_queue_copy_for_global_ref) else: - # This case should ideally not be reached if state update happens print("[load_queue_action] Warning: Skipping global ref update as local copy is None.") print(f"[load_queue_action] Queue load successful. Returning DataFrame update for {len(newly_loaded_queue)} tasks.") - # *** Return the DataFrame update object *** return update_queue_data(newly_loaded_queue) except (ValueError, zipfile.BadZipFile, FileNotFoundError, Exception) as e: error_message = f"Error during queue load: {e}" print(f"[load_queue_action] Caught error: {error_message}") traceback.print_exc() - # Optionally show a Gradio warning/error to the user - gr.Warning(f"Failed to load queue: {error_message[:200]}") # Show truncated error + gr.Warning(f"Failed to load queue: {error_message[:200]}") - # *** Return the DataFrame update for the original queue *** print("[load_queue_action] Load failed. Returning DataFrame update for original queue.") return update_queue_data(original_queue) finally: - # Clean up the uploaded file object if it exists and has a path if filepath and hasattr(filepath, 'name') and filepath.name and os.path.exists(filepath.name): try: - # Gradio often uses temp files, attempting removal is good practice - # os.remove(filepath.name) - # print(f"[load_queue_action] Cleaned up temporary upload file: {filepath.name}") - pass # Let Gradio manage its temp files unless specifically needed + pass except OSError as e: - # Ignore errors like "file not found" if already cleaned up print(f"[load_queue_action] Info: Could not remove temp file {filepath.name}: {e}") pass def clear_queue_action(state): gen = get_gen_info(state) queue = gen.get("queue", []) - if not queue: - gr.Info("Queue is already empty.") - return update_queue_data([]) + aborted_current = False + cleared_pending = False with lock: - queue.clear() - gen["prompts_max"] = 0 + if "in_progress" in gen and gen["in_progress"]: + print("Clear Queue: Signalling abort for in-progress task.") + gen["abort"] = True + gen["extra_orders"] = 0 + if wan_model is not None: + wan_model._interrupt = True + aborted_current = True + + if queue: + if len(queue) > 1 or (len(queue) == 1 and queue[0] is not None and queue[0].get('id') is not None): + print(f"Clear Queue: Clearing {len(queue)} tasks from queue.") + queue.clear() + cleared_pending = True + else: + pass + + if aborted_current or cleared_pending: + gen["prompts_max"] = 0 + + if aborted_current and cleared_pending: + gr.Info("Queue cleared and current generation aborted.") + elif aborted_current: + gr.Info("Current generation aborted.") + elif cleared_pending: + gr.Info("Queue cleared.") + else: + gr.Info("Queue is already empty or only contains the active task (which wasn't aborted now).") - gr.Info("Queue cleared.") return update_queue_data([]) def autosave_queue(): @@ -725,7 +697,7 @@ def autosave_queue(): if os.path.exists(saved_img_abs_path): zf.write(saved_img_abs_path, arcname=saved_img_rel_path) return output_filename - return None # Should not happen if queue has items + return None saved_path = _save_queue_to_file(global_queue_ref, AUTOSAVE_FILENAME) @@ -740,17 +712,15 @@ def autosave_queue(): def autoload_queue(state): global task_id - # Initial check using the original state try: - gen = get_gen_info(state) # Make sure initial state is a dict + gen = get_gen_info(state) original_queue = gen.get("queue", []) except AttributeError: print("[autoload_queue] Error: Initial state is not a dictionary. Cannot autoload.") - # Return default values indicating no load occurred and the state is unchanged - return gr.update(visible=False), False, state # Return an empty DF update + return gr.update(visible=False), False, state loaded_flag = False - dataframe_update = update_queue_data(original_queue) # Default update is the original queue + dataframe_update = update_queue_data(original_queue) if not original_queue and Path(AUTOSAVE_FILENAME).is_file(): print(f"Autoloading queue from {AUTOSAVE_FILENAME}...") @@ -758,38 +728,32 @@ def autoload_queue(state): def __init__(self, name): self.name = name mock_filepath = MockFile(AUTOSAVE_FILENAME) - - # Call load_queue_action, it modifies 'state' internally and returns a DataFrame update dataframe_update = load_queue_action(mock_filepath, state) - # Now check the 'state' dictionary which should have been modified by load_queue_action - gen = get_gen_info(state) # Use the (potentially) modified state dictionary + gen = get_gen_info(state) loaded_queue_after_action = gen.get("queue", []) - if loaded_queue_after_action: # Check if the queue in the state is now populated + if loaded_queue_after_action: print(f"Autoload successful. Loaded {len(loaded_queue_after_action)} tasks into state.") loaded_flag = True - # Global ref update was already done inside load_queue_action if successful else: print("Autoload attempted but queue in state remains empty (file might be empty or invalid).") - # Ensure state reflects empty queue if load failed but file existed with lock: gen["queue"] = [] gen["prompts_max"] = 0 update_global_queue_ref([]) - dataframe_update = update_queue_data([]) # Ensure UI shows empty queue + dataframe_update = update_queue_data([]) - else: # Handle cases where autoload shouldn't happen + else: if original_queue: print("Autoload skipped: Queue is not empty.") - update_global_queue_ref(original_queue) # Ensure global ref matches current state - dataframe_update = update_queue_data(original_queue) # UI should show current queue + update_global_queue_ref(original_queue) + dataframe_update = update_queue_data(original_queue) else: print(f"Autoload skipped: {AUTOSAVE_FILENAME} not found.") - update_global_queue_ref([]) # Ensure global ref is empty - dataframe_update = update_queue_data([]) # UI should show empty queue + update_global_queue_ref([]) + dataframe_update = update_queue_data([]) - # Return the DataFrame update needed for the UI, the flag, and the final state dictionary return dataframe_update, loaded_flag, state From 87614b8216173caac19b54a674fc66d85b153ede Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Thu, 10 Apr 2025 04:53:18 +1000 Subject: [PATCH 47/69] fix manual load queue --- wgp.py | 67 +++++++++++++++++++++++++++++++++++++++------------------- 1 file changed, 45 insertions(+), 22 deletions(-) diff --git a/wgp.py b/wgp.py index a2a468f..5dd3633 100644 --- a/wgp.py +++ b/wgp.py @@ -756,6 +756,28 @@ def autoload_queue(state): return dataframe_update, loaded_flag, state +def run_autoload_and_prepare_ui(current_state): + df_update, loaded_flag, modified_state = autoload_queue(current_state) + should_start_processing = loaded_flag + accordion_update = gr.Accordion(open=True) if loaded_flag else gr.update() + return df_update, gr.update(visible=loaded_flag), accordion_update, should_start_processing, modified_state + +def start_processing_if_needed(should_start, current_state): + if not isinstance(current_state, dict) or 'gen' not in current_state: + yield "Error: Invalid state received before processing." + return + if should_start: + yield from process_tasks(current_state) + else: + yield None + +def finalize_generation_with_state(current_state): + if not isinstance(current_state, dict) or 'gen' not in current_state: + return gr.update(), gr.update(interactive=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False, value=""), gr.update(), current_state + + gallery_update, abort_btn_update, gen_btn_update, add_queue_btn_update, current_gen_col_update, gen_info_update = finalize_generation(current_state) + accordion_update = gr.Accordion(open=False) if len(get_gen_info(current_state).get("queue", [])) <= 1 else gr.update() + return gallery_update, abort_btn_update, gen_btn_update, add_queue_btn_update, current_gen_col_update, gen_info_update, accordion_update, current_state def get_queue_table(queue): data = [] @@ -3503,6 +3525,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non js=trigger_zip_download_js ) + should_start_flag = gr.State(False) load_queue_btn.upload( fn=load_queue_action, inputs=[load_queue_btn, state], @@ -3511,6 +3534,28 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non fn=lambda s: (gr.update(visible=bool(get_gen_info(s).get("queue",[]))), gr.Accordion(open=True)) if bool(get_gen_info(s).get("queue",[])) else (gr.update(visible=False), gr.update()), inputs=[state], outputs=[current_gen_column, queue_accordion] + ).then( + fn=lambda s: ( + (gr.Button(visible=False), gr.Button(visible=True), gr.Column(visible=True), True) + if bool(get_gen_info(s).get("queue",[])) + else (gr.Button(visible=True), gr.Button(visible=False), gr.Column(visible=False), False) + ), + inputs=[state], + outputs=[generate_btn, add_to_queue_btn, current_gen_column, should_start_flag] + ).then( + fn=start_processing_if_needed, + inputs=[should_start_flag, state], + outputs=[gen_status], + trigger_mode="once" + ).then( + fn=finalize_generation_with_state, + inputs=[state], + outputs=[output, abort_btn, generate_btn, add_to_queue_btn, current_gen_column, gen_info, queue_accordion, state], + trigger_mode="always_last" + ).then( + unload_model_if_needed, + inputs= [state], + outputs= [] ) clear_queue_btn.click( @@ -4139,28 +4184,6 @@ def create_demo(): generate_about_tab() should_start_flag = gr.State(False) - def run_autoload_and_prepare_ui(current_state): - df_update, loaded_flag, modified_state = autoload_queue(current_state) - should_start_processing = loaded_flag - accordion_update = gr.Accordion(open=True) if loaded_flag else gr.update() - return df_update, gr.update(visible=loaded_flag), accordion_update, should_start_processing, modified_state - - def start_processing_if_needed(should_start, current_state): - if not isinstance(current_state, dict) or 'gen' not in current_state: - yield "Error: Invalid state received before processing." - return - if should_start: - yield from process_tasks(current_state) - else: - yield None - - def finalize_generation_with_state(current_state): - if not isinstance(current_state, dict) or 'gen' not in current_state: - return gr.update(), gr.update(interactive=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False, value=""), gr.update(), current_state - - gallery_update, abort_btn_update, gen_btn_update, add_queue_btn_update, current_gen_col_update, gen_info_update = finalize_generation(current_state) - accordion_update = gr.Accordion(open=False) if len(get_gen_info(current_state).get("queue", [])) <= 1 else gr.update() - return gallery_update, abort_btn_update, gen_btn_update, add_queue_btn_update, current_gen_col_update, gen_info_update, accordion_update, current_state demo.load( fn=run_autoload_and_prepare_ui, From 2e026c67c26eca0adb34bb28ae0a9f1d3af30766 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Thu, 10 Apr 2025 06:05:45 +1000 Subject: [PATCH 48/69] add save and exit, fix vace attachments --- wgp.py | 227 ++++++++++++++++++++++++++++++++++++++++++++------------- 1 file changed, 176 insertions(+), 51 deletions(-) diff --git a/wgp.py b/wgp.py index 5dd3633..5c4643d 100644 --- a/wgp.py +++ b/wgp.py @@ -31,6 +31,7 @@ from PIL import Image import zipfile import tempfile import atexit +import shutil global_queue_ref = [] AUTOSAVE_FILENAME = "queue.zip" PROMPT_VARS_MAX = 10 @@ -364,21 +365,23 @@ def save_queue_action(state): if not queue or len(queue) <=1 : gr.Info("Queue is empty. Nothing to save.") - return None + return "" zip_buffer = io.BytesIO() with tempfile.TemporaryDirectory() as tmpdir: queue_manifest = [] - image_paths_in_zip = {} + file_paths_in_zip = {} for task_index, task in enumerate(queue): - if task is None or not isinstance(task, dict) or task_index == 0: continue + if task is None or not isinstance(task, dict) or task.get('id') is None: continue params_copy = task.get('params', {}).copy() task_id_s = task.get('id', f"task_{task_index}") image_keys = ["image_start", "image_end", "image_refs"] + video_keys = ["video_guide", "video_mask"] + for key in image_keys: images_pil = params_copy.get(key) if images_pil is None: @@ -395,8 +398,8 @@ def save_queue_action(state): continue img_id = id(pil_image) - if img_id in image_paths_in_zip: - image_filenames_for_json.append(image_paths_in_zip[img_id]) + if img_id in file_paths_in_zip: + image_filenames_for_json.append(file_paths_in_zip[img_id]) continue img_filename_in_zip = f"task{task_id_s}_{key}_{img_index}.png" @@ -405,7 +408,8 @@ def save_queue_action(state): try: pil_image.save(img_save_path, "PNG") image_filenames_for_json.append(img_filename_in_zip) - image_paths_in_zip[img_id] = img_filename_in_zip + file_paths_in_zip[img_id] = img_filename_in_zip + print(f"Saved image: {img_filename_in_zip}") except Exception as e: print(f"Error saving image {img_filename_in_zip} for task {task_id_s}: {e}") @@ -414,17 +418,47 @@ def save_queue_action(state): else: params_copy.pop(key, None) + for key in video_keys: + video_path_orig = params_copy.get(key) + if video_path_orig is None or not isinstance(video_path_orig, str): + continue + + if video_path_orig in file_paths_in_zip: + params_copy[key] = file_paths_in_zip[video_path_orig] + continue + + if not os.path.isfile(video_path_orig): + print(f"Warning: Video file not found for key '{key}' in task {task_id_s}: {video_path_orig}. Skipping video.") + params_copy.pop(key, None) + continue + + _, extension = os.path.splitext(video_path_orig) + vid_filename_in_zip = f"task{task_id_s}_{key}{extension if extension else '.mp4'}" + vid_save_path = os.path.join(tmpdir, vid_filename_in_zip) + + try: + shutil.copy2(video_path_orig, vid_save_path) + params_copy[key] = vid_filename_in_zip + file_paths_in_zip[video_path_orig] = vid_filename_in_zip + print(f"Copied video: {video_path_orig} -> {vid_filename_in_zip}") + except Exception as e: + print(f"Error copying video {video_path_orig} to {vid_filename_in_zip} for task {task_id_s}: {e}") + params_copy.pop(key, None) + params_copy.pop('state', None) params_copy.pop('start_image_data_base64', None) params_copy.pop('end_image_data_base64', None) params_copy.pop('start_image_data', None) params_copy.pop('end_image_data', None) + task.pop('start_image_data', None) + task.pop('end_image_data', None) manifest_entry = { "id": task.get('id'), "params": params_copy, } + manifest_entry = {k: v for k, v in manifest_entry.items() if v is not None} queue_manifest.append(manifest_entry) manifest_path = os.path.join(tmpdir, "queue.json") @@ -440,12 +474,13 @@ def save_queue_action(state): with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zf: zf.write(manifest_path, arcname="queue.json") - for saved_img_rel_path in image_paths_in_zip.values(): - saved_img_abs_path = os.path.join(tmpdir, saved_img_rel_path) - if os.path.exists(saved_img_abs_path): - zf.write(saved_img_abs_path, arcname=saved_img_rel_path) + for file_id, saved_file_rel_path in file_paths_in_zip.items(): + saved_file_abs_path = os.path.join(tmpdir, saved_file_rel_path) + if os.path.exists(saved_file_abs_path): + zf.write(saved_file_abs_path, arcname=saved_file_rel_path) + print(f"Adding to zip: {saved_file_rel_path}") else: - print(f"Warning: Image file {saved_img_rel_path} not found during zipping.") + print(f"Warning: File {saved_file_rel_path} (ID: {file_id}) not found during zipping.") zip_buffer.seek(0) zip_binary_content = zip_buffer.getvalue() @@ -464,6 +499,8 @@ def load_queue_action(filepath, state): global task_id gen = get_gen_info(state) original_queue = gen.get("queue", []) + save_path_base = server_config.get("save_path", "outputs") + loaded_cache_dir = os.path.join(save_path_base, "_loaded_queue_cache") if not filepath or not hasattr(filepath, 'name') or not Path(filepath.name).is_file(): print("[load_queue_action] Warning: No valid file selected or file not found.") @@ -476,6 +513,9 @@ def load_queue_action(filepath, state): try: print(f"[load_queue_action] Attempting to load queue from: {filepath.name}") + os.makedirs(loaded_cache_dir, exist_ok=True) + print(f"[load_queue_action] Using cache directory: {loaded_cache_dir}") + with tempfile.TemporaryDirectory() as tmpdir: with zipfile.ZipFile(filepath.name, 'r') as zf: if "queue.json" not in zf.namelist(): raise ValueError("queue.json not found in zip file") @@ -497,21 +537,29 @@ def load_queue_action(filepath, state): params = task_data.get('params', {}) task_id_loaded = task_data.get('id', 0) max_id_in_file = max(max_id_in_file, task_id_loaded) - loaded_pil_images = {} - image_keys = ["image_start", "image_end", "image_refs"] params['state'] = state + image_keys = ["image_start", "image_end", "image_refs"] + video_keys = ["video_guide", "video_mask"] + + loaded_pil_images = {} + loaded_video_paths = {} + for key in image_keys: image_filenames = params.get(key) if image_filenames is None: continue + is_list = isinstance(image_filenames, list) if not is_list: image_filenames = [image_filenames] + loaded_pils = [] for img_filename_in_zip in image_filenames: - if not isinstance(img_filename_in_zip, str): continue + if not isinstance(img_filename_in_zip, str): + print(f"[load_queue_action] Warning: Non-string filename found for image key '{key}'. Skipping.") + continue img_load_path = os.path.join(tmpdir, img_filename_in_zip) if not os.path.exists(img_load_path): - print(f"[load_queue_action] Image file not found during load: {img_load_path}") + print(f"[load_queue_action] Image file not found in extracted data: {img_load_path}. Skipping.") continue try: pil_image = Image.open(img_load_path) @@ -519,30 +567,53 @@ def load_queue_action(filepath, state): converted_image = convert_image(pil_image) loaded_pils.append(converted_image) pil_image.close() + print(f"Loaded image: {img_filename_in_zip} for key {key}") except Exception as img_e: print(f"[load_queue_action] Error loading image {img_filename_in_zip}: {img_e}") if loaded_pils: params[key] = loaded_pils if is_list else loaded_pils[0] loaded_pil_images[key] = params[key] - else: params.pop(key, None) + else: + params.pop(key, None) - primary_preview_pil, secondary_preview_pil = None, None - start_prev_pil_list = loaded_pil_images.get("image_start") - end_prev_pil_list = loaded_pil_images.get("image_end") - ref_prev_pil_list = loaded_pil_images.get("image_refs") + for key in video_keys: + video_filename_in_zip = params.get(key) + if video_filename_in_zip is None or not isinstance(video_filename_in_zip, str): + continue - if start_prev_pil_list: - primary_preview_pil = start_prev_pil_list[0] if isinstance(start_prev_pil_list, list) and start_prev_pil_list else start_prev_pil_list if not isinstance(start_prev_pil_list, list) else None - if end_prev_pil_list: - secondary_preview_pil = end_prev_pil_list[0] if isinstance(end_prev_pil_list, list) and end_prev_pil_list else end_prev_pil_list if not isinstance(end_prev_pil_list, list) else None - elif ref_prev_pil_list and isinstance(ref_prev_pil_list, list) and ref_prev_pil_list: - primary_preview_pil = ref_prev_pil_list[0] + video_load_path = os.path.join(tmpdir, video_filename_in_zip) + if not os.path.exists(video_load_path): + print(f"[load_queue_action] Video file not found in extracted data: {video_load_path}. Skipping.") + params.pop(key, None) + continue + + persistent_video_path = os.path.join(loaded_cache_dir, video_filename_in_zip) + try: + shutil.copy2(video_load_path, persistent_video_path) + params[key] = persistent_video_path + loaded_video_paths[key] = persistent_video_path + print(f"Loaded video: {video_filename_in_zip} -> {persistent_video_path}") + except Exception as vid_e: + print(f"[load_queue_action] Error copying video {video_filename_in_zip} to cache: {vid_e}") + params.pop(key, None) + + + primary_preview_pil_list = loaded_pil_images.get("image_start") or loaded_pil_images.get("image_refs") + secondary_preview_pil_list = loaded_pil_images.get("image_end") + + primary_preview_pil = None + if primary_preview_pil_list: + primary_preview_pil = primary_preview_pil_list[0] if isinstance(primary_preview_pil_list, list) else primary_preview_pil_list + + secondary_preview_pil = None + if secondary_preview_pil_list: + secondary_preview_pil = secondary_preview_pil_list[0] if isinstance(secondary_preview_pil_list, list) else secondary_preview_pil_list start_b64 = [pil_to_base64_uri(primary_preview_pil, format="jpeg", quality=70)] if primary_preview_pil else None end_b64 = [pil_to_base64_uri(secondary_preview_pil, format="jpeg", quality=70)] if secondary_preview_pil else None - top_level_start_image = loaded_pil_images.get("image_start") - top_level_end_image = loaded_pil_images.get("image_end") + top_level_start_image = params.get("image_start") or params.get("image_refs") + top_level_end_image = params.get("image_end") runtime_task = { "id": task_id_loaded, @@ -557,19 +628,20 @@ def load_queue_action(filepath, state): "end_image_data_base64": end_b64, } newly_loaded_queue.append(runtime_task) - print(f"[load_queue_action] Processed task {task_index+1}/{len(loaded_manifest)}, ID: {task_id_loaded}") + print(f"[load_queue_action] Reconstructed task {task_index+1}/{len(loaded_manifest)}, ID: {task_id_loaded}") with lock: print("[load_queue_action] Acquiring lock to update state...") gen["queue"] = newly_loaded_queue[:] local_queue_copy_for_global_ref = gen["queue"][:] - current_max_id_in_new_queue = max([t['id'] for t in newly_loaded_queue if 'id' in t] + [0]) - if current_max_id_in_new_queue > task_id: - print(f"[load_queue_action] Updating global task_id from {task_id} to {current_max_id_in_new_queue + 1}") - task_id = current_max_id_in_new_queue + 1 + current_max_id_in_new_queue = max([t['id'] for t in newly_loaded_queue if 'id' in t] + [0]) + if current_max_id_in_new_queue >= task_id: + new_task_id = current_max_id_in_new_queue + 1 + print(f"[load_queue_action] Updating global task_id from {task_id} to {new_task_id}") + task_id = new_task_id else: - print(f"[load_queue_action] Global task_id ({task_id}) is >= max in file ({current_max_id_in_new_queue}). Not changing task_id.") + print(f"[load_queue_action] Global task_id ({task_id}) is > max in file ({current_max_id_in_new_queue}). Not changing task_id.") gen["prompts_max"] = len(newly_loaded_queue) print("[load_queue_action] State update complete. Releasing lock.") @@ -593,11 +665,14 @@ def load_queue_action(filepath, state): return update_queue_data(original_queue) finally: if filepath and hasattr(filepath, 'name') and filepath.name and os.path.exists(filepath.name): - try: - pass - except OSError as e: - print(f"[load_queue_action] Info: Could not remove temp file {filepath.name}: {e}") - pass + if tempfile.gettempdir() in os.path.abspath(filepath.name): + try: + os.remove(filepath.name) + print(f"[load_queue_action] Removed temporary upload file: {filepath.name}") + except OSError as e: + print(f"[load_queue_action] Info: Could not remove temp file {filepath.name}: {e}") + else: + print(f"[load_queue_action] Info: Did not remove non-temporary file: {filepath.name}") def clear_queue_action(state): gen = get_gen_info(state) @@ -636,6 +711,12 @@ def clear_queue_action(state): return update_queue_data([]) +def quit_application(): + print("Save and Quit requested...") + autosave_queue() + import signal + os.kill(os.getpid(), signal.SIGINT) + def autosave_queue(): global global_queue_ref if not global_queue_ref: @@ -649,14 +730,20 @@ def autosave_queue(): def _save_queue_to_file(queue_to_save, output_filename): if not queue_to_save: return None + with tempfile.TemporaryDirectory() as tmpdir: queue_manifest = [] - image_paths_in_zip = {} + file_paths_in_zip = {} + for task_index, task in enumerate(queue_to_save): - if task is None or not isinstance(task, dict): continue + if task is None or not isinstance(task, dict) or task.get('id') is None: continue + params_copy = task.get('params', {}).copy() task_id_s = task.get('id', f"task_{task_index}") + image_keys = ["image_start", "image_end", "image_refs"] + video_keys = ["video_guide", "video_mask"] + for key in image_keys: images_pil = params_copy.get(key) if images_pil is None: continue @@ -666,36 +753,70 @@ def autosave_queue(): for img_index, pil_image in enumerate(images_pil): if not isinstance(pil_image, Image.Image): continue img_id = id(pil_image) - if img_id in image_paths_in_zip: - image_filenames_for_json.append(image_paths_in_zip[img_id]) + if img_id in file_paths_in_zip: + image_filenames_for_json.append(file_paths_in_zip[img_id]) continue img_filename_in_zip = f"task{task_id_s}_{key}_{img_index}.png" img_save_path = os.path.join(tmpdir, img_filename_in_zip) try: pil_image.save(img_save_path, "PNG") image_filenames_for_json.append(img_filename_in_zip) - image_paths_in_zip[img_id] = img_filename_in_zip + file_paths_in_zip[img_id] = img_filename_in_zip except Exception as e: print(f"Autosave error saving image {img_filename_in_zip}: {e}") if image_filenames_for_json: params_copy[key] = image_filenames_for_json if is_list else image_filenames_for_json[0] else: params_copy.pop(key, None) + + for key in video_keys: + video_path_orig = params_copy.get(key) + if video_path_orig is None or not isinstance(video_path_orig, str): + continue + + if video_path_orig in file_paths_in_zip: + params_copy[key] = file_paths_in_zip[video_path_orig] + continue + + if not os.path.isfile(video_path_orig): + print(f"Warning (Autosave): Video file not found for key '{key}' in task {task_id_s}: {video_path_orig}. Skipping.") + params_copy.pop(key, None) + continue + + _, extension = os.path.splitext(video_path_orig) + vid_filename_in_zip = f"task{task_id_s}_{key}{extension if extension else '.mp4'}" + vid_save_path = os.path.join(tmpdir, vid_filename_in_zip) + + try: + shutil.copy2(video_path_orig, vid_save_path) + params_copy[key] = vid_filename_in_zip + file_paths_in_zip[video_path_orig] = vid_filename_in_zip + except Exception as e: + print(f"Error (Autosave) copying video {video_path_orig} to {vid_filename_in_zip} for task {task_id_s}: {e}") + params_copy.pop(key, None) params_copy.pop('state', None) params_copy.pop('start_image_data_base64', None) params_copy.pop('end_image_data_base64', None) + params_copy.pop('start_image_data', None) + params_copy.pop('end_image_data', None) + manifest_entry = { - "id": task.get('id'), "params": params_copy, + "id": task.get('id'), + "params": params_copy, } + manifest_entry = {k: v for k, v in manifest_entry.items() if v is not None} queue_manifest.append(manifest_entry) + manifest_path = os.path.join(tmpdir, "queue.json") with open(manifest_path, 'w', encoding='utf-8') as f: json.dump(queue_manifest, f, indent=4) with zipfile.ZipFile(output_filename, 'w', zipfile.ZIP_DEFLATED) as zf: zf.write(manifest_path, arcname="queue.json") - for saved_img_rel_path in image_paths_in_zip.values(): - saved_img_abs_path = os.path.join(tmpdir, saved_img_rel_path) - if os.path.exists(saved_img_abs_path): - zf.write(saved_img_abs_path, arcname=saved_img_rel_path) + for saved_file_rel_path in file_paths_in_zip.values(): + saved_file_abs_path = os.path.join(tmpdir, saved_file_rel_path) + if os.path.exists(saved_file_abs_path): + zf.write(saved_file_abs_path, arcname=saved_file_rel_path) + else: + print(f"Warning (Autosave): File {saved_file_rel_path} not found during zipping.") return output_filename return None @@ -1160,7 +1281,6 @@ text_encoder_choices = ["ckpts/models_t5_umt5-xxl-enc-bf16.safetensors", "ckpts/ server_config_filename = "wgp_config.json" if not os.path.isfile(server_config_filename) and os.path.isfile("gradio_config.json"): - import shutil shutil.move("gradio_config.json", server_config_filename) if not Path(server_config_filename).is_file(): @@ -2965,7 +3085,6 @@ def download_loras(): lora_dir = get_lora_dir(get_model_filename("i2v"), quantizeTransformer) log_path = os.path.join(lora_dir, "log.txt") if not os.path.isfile(log_path): - import shutil tmp_path = os.path.join(lora_dir, "tmp_lora_dowload") import glob snapshot_download(repo_id="DeepBeepMeep/Wan2.1", allow_patterns="loras_i2v/*", local_dir= tmp_path) @@ -3483,6 +3602,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non save_queue_btn = gr.DownloadButton("Save Queue", size="sm") load_queue_btn = gr.UploadButton("Load Queue", file_types=[".zip"], size="sm") clear_queue_btn = gr.Button("Clear Queue", size="sm", variant="stop") + quit_button = gr.Button("Save and Quit", size="sm", variant="secondary") trigger_zip_download_js = """ (base64String) => { if (!base64String) { @@ -3567,6 +3687,11 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non inputs=None, outputs=[current_gen_column, queue_accordion] ) + quit_button.click( + fn=quit_application, + inputs=[], + outputs=[] + ) extra_inputs = prompt_vars + [wizard_prompt, wizard_variables_var, wizard_prompt_activated_var, video_prompt_column, image_prompt_column, prompt_column_advanced, prompt_column_wizard_vars, prompt_column_wizard, lset_name, advanced_row] # show_advanced presets_column, From eba935515a4e3a3d7092cbfd3e36121310d56a1e Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Wed, 9 Apr 2025 22:11:04 +0200 Subject: [PATCH 49/69] Added missing requirements --- requirements.txt | 1 + wgp.py | 13 +++++++++---- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/requirements.txt b/requirements.txt index a4cd3b3..babec2a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,4 +22,5 @@ mutagen decord onnxruntime-gpu rembg[gpu]==2.0.65 +matplotlib # rembg==2.0.65 \ No newline at end of file diff --git a/wgp.py b/wgp.py index 97cd830..35200d4 100644 --- a/wgp.py +++ b/wgp.py @@ -173,9 +173,14 @@ def process_prompt_and_add_tasks(state, model_choice): return else: video_mask = None - if "O" in video_prompt_type and inputs["max_frames"]==0: - gr.Info(f"In order to extend a video, you need to indicate how many frames you want to reuse in the source video.") - return + if "O" in video_prompt_type : + max_frames= inputs["max_frames"] + video_length = inputs["video_length"] + if max_frames ==0: + gr.Info(f"Warning : you have asked to reuse all the frames of the control Video before extending it. Please make sure the number of frames of the control Video is lower than the total number of frames to generate otherwise it won't make a difference.") + elif max_frames >= video_length: + gr.Info(f"The number of frames in the control Video to reuse ({max_frames}) before extending the Video can not be bigger than the total number of frames ({video_length}) to generate.") + return if isinstance(image_refs, list): image_refs = [ convert_image(tup[0]) for tup in image_refs ] @@ -2060,7 +2065,7 @@ def generate_video( # gr.Info("Unable to generate a Video while a new configuration is being applied.") # return - if "P" in preload_model_policy: + if "P" in preload_model_policy and not "U" in preload_model_policy: while wan_model == None: time.sleep(1) From cd787eedcacc16f814e47fc0cfabf0a9576a557d Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Wed, 9 Apr 2025 22:40:18 +0200 Subject: [PATCH 50/69] Added missing requirements --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index babec2a..f254df2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -23,4 +23,5 @@ decord onnxruntime-gpu rembg[gpu]==2.0.65 matplotlib +timm # rembg==2.0.65 \ No newline at end of file From 8d26bbc2c63a48f0771611cf9ce0b718dd46dc17 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Thu, 10 Apr 2025 00:18:44 +0200 Subject: [PATCH 51/69] Tweaked Vace UI --- requirements.txt | 2 +- wgp.py | 47 ++++++++++++++++++++++++++++++----------------- 2 files changed, 31 insertions(+), 18 deletions(-) diff --git a/requirements.txt b/requirements.txt index f254df2..70126cd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,7 +12,7 @@ ftfy dashscope imageio-ffmpeg # flash_attn -gradio>=5.0.0 +gradio>=5.0.0 numpy>=1.23.5,<2 einops moviepy==1.0.3 diff --git a/wgp.py b/wgp.py index 191040c..77182ab 100644 --- a/wgp.py +++ b/wgp.py @@ -177,10 +177,10 @@ def process_prompt_and_add_tasks(state, model_choice): if "O" in video_prompt_type : max_frames= inputs["max_frames"] video_length = inputs["video_length"] - if max_frames ==0: - gr.Info(f"Warning : you have asked to reuse all the frames of the control Video before extending it. Please make sure the number of frames of the control Video is lower than the total number of frames to generate otherwise it won't make a difference.") + if max_frames ==0: + gr.Info(f"Warning : you have asked to reuse all the frames of the control Video in the Alternate Video End it. Please make sure the number of frames of the control Video is lower than the total number of frames to generate otherwise it won't make a difference.") elif max_frames >= video_length: - gr.Info(f"The number of frames in the control Video to reuse ({max_frames}) before extending the Video can not be bigger than the total number of frames ({video_length}) to generate.") + gr.Info(f"The number of frames in the control Video to reuse ({max_frames}) in Alternate Video End can not be bigger than the total number of frames ({video_length}) to generate.") return if isinstance(image_refs, list): @@ -3232,6 +3232,8 @@ def del_in_sequence(source_str, letters): def refresh_video_prompt_type_image_refs(video_prompt_type, video_prompt_type_image_refs): + # video_prompt_type = add_to_sequence(video_prompt_type, "I") if video_prompt_type_image_refs else del_in_sequence(video_prompt_type, "I") + video_prompt_type_image_refs = "I" in video_prompt_type_image_refs video_prompt_type = add_to_sequence(video_prompt_type, "I") if video_prompt_type_image_refs else del_in_sequence(video_prompt_type, "I") return video_prompt_type, gr.update(visible = video_prompt_type_image_refs),gr.update(visible = video_prompt_type_image_refs) @@ -3352,20 +3354,31 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non with gr.Column(visible= "Vace" in model_filename ) as video_prompt_column: video_prompt_type_value= ui_defaults.get("video_prompt_type","") video_prompt_type = gr.Text(value= video_prompt_type_value, visible= False) - video_prompt_type_video_guide = gr.Dropdown( - choices=[ - ("None, use only the Text Prompt", ""), - ("Extend the Control Video", "OV"), - ("Transfer Human Motion from the Control Video", "PV"), - ("Transfer Depth from the Control Video", "DV"), - ("Recolorize the Control Video", "CV"), - ("Control Video contains Open Pose, Depth or Black & White ", "V"), - ("Inpainting of Control Video using Mask Video ", "MV"), - ], - value=filter_letters(video_prompt_type_value, "ODPCMV"), - label="Video to Video" - ) - video_prompt_type_image_refs = gr.Checkbox(value="I" in video_prompt_type_value , label= "Use References Images (Faces, Objects) to customize New Video", scale =1 ) + with gr.Row(): + video_prompt_type_video_guide = gr.Dropdown( + choices=[ + ("None", ""), + ("Transfer Human Motion from the Control Video", "PV"), + ("Transfer Depth from the Control Video", "DV"), + ("Recolorize the Control Video", "CV"), + ("Alternate Video End", "OV"), + ("(adv) Video contains Open Pose, Depth or Black & White ", "V"), + ("(adv) Inpainting of Control Video using Mask Video ", "MV"), + ], + value=filter_letters(video_prompt_type_value, "ODPCMV"), + label="Video to Video", scale = 3 + ) + + video_prompt_type_image_refs = gr.Dropdown( + choices=[ + ("None", ""), + ("Inject custom Faces / Objects", "I"), + ], + value="I" if "I" in video_prompt_type_value else "", + label="Reference Images", scale = 2 + ) + + # video_prompt_type_image_refs = gr.Checkbox(value="I" in video_prompt_type_value , label= "Use References Images (Faces, Objects) to customize New Video", scale =1 ) video_guide = gr.Video(label= "Control Video", visible= "V" in video_prompt_type_value, value= ui_defaults.get("video_guide", None),) max_frames = gr.Slider(0, 100, value=ui_defaults.get("max_frames",0), step=1, label="Nb of frames in Control Video to use (0 = max)", visible= "V" in video_prompt_type_value, scale = 2 ) From b4ac6a8e40a50cab9e490fc40772296718c16b5e Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Thu, 10 Apr 2025 00:45:17 +0200 Subject: [PATCH 52/69] Fixed spelling --- wgp.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/wgp.py b/wgp.py index 77182ab..6681097 100644 --- a/wgp.py +++ b/wgp.py @@ -178,9 +178,9 @@ def process_prompt_and_add_tasks(state, model_choice): max_frames= inputs["max_frames"] video_length = inputs["video_length"] if max_frames ==0: - gr.Info(f"Warning : you have asked to reuse all the frames of the control Video in the Alternate Video End it. Please make sure the number of frames of the control Video is lower than the total number of frames to generate otherwise it won't make a difference.") + gr.Info(f"Warning : you have asked to reuse all the frames of the control Video in the Alternate Video Ending it. Please make sure the number of frames of the control Video is lower than the total number of frames to generate otherwise it won't make a difference.") elif max_frames >= video_length: - gr.Info(f"The number of frames in the control Video to reuse ({max_frames}) in Alternate Video End can not be bigger than the total number of frames ({video_length}) to generate.") + gr.Info(f"The number of frames in the control Video to reuse ({max_frames}) in Alternate Video Ending can not be bigger than the total number of frames ({video_length}) to generate.") return if isinstance(image_refs, list): @@ -3361,7 +3361,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ("Transfer Human Motion from the Control Video", "PV"), ("Transfer Depth from the Control Video", "DV"), ("Recolorize the Control Video", "CV"), - ("Alternate Video End", "OV"), + ("Alternate Video Ending", "OV"), ("(adv) Video contains Open Pose, Depth or Black & White ", "V"), ("(adv) Inpainting of Control Video using Mask Video ", "MV"), ], From f5bb9d597298ffd671e777ef26cae73940cf92d1 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Thu, 10 Apr 2025 01:00:19 +0200 Subject: [PATCH 53/69] disabled Vaced alternate ending for the moment --- wgp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wgp.py b/wgp.py index 6681097..53e6b5c 100644 --- a/wgp.py +++ b/wgp.py @@ -3361,7 +3361,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ("Transfer Human Motion from the Control Video", "PV"), ("Transfer Depth from the Control Video", "DV"), ("Recolorize the Control Video", "CV"), - ("Alternate Video Ending", "OV"), + # ("Alternate Video Ending", "OV"), ("(adv) Video contains Open Pose, Depth or Black & White ", "V"), ("(adv) Inpainting of Control Video using Mask Video ", "MV"), ], From e39800a661e0cbb2c3e913d5eee6d3aa8cacdb24 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Thu, 10 Apr 2025 14:13:48 +1000 Subject: [PATCH 54/69] remove previous queue autosave if queue cleared --- wgp.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/wgp.py b/wgp.py index 5c4643d..4f7d491 100644 --- a/wgp.py +++ b/wgp.py @@ -700,6 +700,15 @@ def clear_queue_action(state): if aborted_current or cleared_pending: gen["prompts_max"] = 0 + if cleared_pending: + try: + if os.path.isfile(AUTOSAVE_FILENAME): + os.remove(AUTOSAVE_FILENAME) + print(f"Clear Queue: Deleted autosave file '{AUTOSAVE_FILENAME}'.") + except OSError as e: + print(f"Clear Queue: Error deleting autosave file '{AUTOSAVE_FILENAME}': {e}") + gr.Warning(f"Could not delete the autosave file '{AUTOSAVE_FILENAME}'. You may need to remove it manually.") + if aborted_current and cleared_pending: gr.Info("Queue cleared and current generation aborted.") elif aborted_current: @@ -3792,7 +3801,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non inputs = [state, model_choice], outputs=queue_df ).then( - fn=lambda s: gr.Accordion(open=True) if len(get_gen_info(s).get("queue", [])) > 1 else gr.update(), # Expand if queue has items (len > 1 assumes placeholder) + fn=lambda s: gr.Accordion(open=True) if len(get_gen_info(s).get("queue", [])) > 1 else gr.update(), inputs=[state], outputs=[queue_accordion] ).then( From e3d8acde83790e471451beb879af9e8d460a3b8d Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Thu, 10 Apr 2025 17:34:19 +1000 Subject: [PATCH 55/69] 5 sec confirmation window for quit button --- wgp.py | 89 ++++++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 84 insertions(+), 5 deletions(-) diff --git a/wgp.py b/wgp.py index 704df0c..48ca7b8 100644 --- a/wgp.py +++ b/wgp.py @@ -731,6 +731,12 @@ def quit_application(): import signal os.kill(os.getpid(), signal.SIGINT) +def request_quit_confirmation(): + return gr.update(visible=False), gr.update(visible=True) + +def cancel_quit_confirmation(): + return gr.update(visible=True), gr.update(visible=False) + def autosave_queue(): global global_queue_ref if not global_queue_ref: @@ -3630,6 +3636,55 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non load_queue_btn = gr.UploadButton("Load Queue", file_types=[".zip"], size="sm") clear_queue_btn = gr.Button("Clear Queue", size="sm", variant="stop") quit_button = gr.Button("Save and Quit", size="sm", variant="secondary") + with gr.Row(visible=False) as quit_confirmation_row: + gr.Markdown("Quitting in 5 seconds...", elem_id="quit_timer_label") + confirm_quit_button = gr.Button("Confirm Quit Now", elem_id="comfirm_quit_btn_hidden", size="sm", variant="stop") + cancel_quit_button = gr.Button("Cancel Quit", size="sm", variant="secondary") + hidden_force_quit_trigger = gr.Button("force_quit", visible=False, elem_id="force_quit_btn_hidden") + + start_quit_timer_js = """ + () => { + function findAndClickGradioButton(elemId) { + const gradioApp = document.querySelector('gradio-app') || document; + const button = gradioApp.querySelector(`#${elemId}`); + if (button) { + button.click(); + } + } + window.quitTimerId = setTimeout(() => { + }, 5000); + let countdown = 5; + const label = document.getElementById('quit_timer_label'); + if (label) { + label.innerText = `Quitting in ${countdown}...`; + window.quitCountdownInterval = setInterval(() => { + countdown--; + if (countdown > 0) { + label.innerText = `Quitting in ${countdown}...`; + } else { + clearInterval(window.quitCountdownInterval); + findAndClickGradioButton('comfirm_quit_btn_hidden'); + } + }, 1000); + } + } + """ + + cancel_quit_timer_js = """ + () => { + if (window.quitTimerId) { + clearTimeout(window.quitTimerId); + window.quitTimerId = null; + } + if(window.quitCountdownInterval) { + clearInterval(window.quitCountdownInterval); + window.quitCountdownInterval = null; + } + const label = document.getElementById('quit_timer_label'); + if(label) { label.innerText = 'Quit cancelled.'; } + } + """ + trigger_zip_download_js = """ (base64String) => { if (!base64String) { @@ -3661,6 +3716,35 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non } } """ + + quit_button.click( + fn=request_quit_confirmation, + inputs=[], + outputs=[quit_button, quit_confirmation_row] + ).then( + fn=None, inputs=None, outputs=None, js=start_quit_timer_js + ) + + confirm_quit_button.click( + fn=quit_application, + inputs=[], + outputs=[] + ) + + cancel_quit_button.click( + fn=cancel_quit_confirmation, + inputs=[], + outputs=[quit_button, quit_confirmation_row] + ).then( + fn=None, inputs=None, outputs=None, js=cancel_quit_timer_js + ) + + hidden_force_quit_trigger.click( + fn=quit_application, + inputs=[], + outputs=[] + ) + save_queue_btn.click( fn=save_queue_action, inputs=[state], @@ -3714,11 +3798,6 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non inputs=None, outputs=[current_gen_column, queue_accordion] ) - quit_button.click( - fn=quit_application, - inputs=[], - outputs=[] - ) extra_inputs = prompt_vars + [wizard_prompt, wizard_variables_var, wizard_prompt_activated_var, video_prompt_column, image_prompt_column, prompt_column_advanced, prompt_column_wizard_vars, prompt_column_wizard, lset_name, advanced_row] # show_advanced presets_column, From 8d998d645f0ca88dfb90faab3e03a71fbbceaa84 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Fri, 11 Apr 2025 00:10:03 +0200 Subject: [PATCH 56/69] Optimized Vace RAM usage --- wan/modules/model.py | 152 +++++++++++++++++++++++++++---------------- wan/text2video.py | 13 ++++ wgp.py | 23 ------- 3 files changed, 110 insertions(+), 78 deletions(-) diff --git a/wan/modules/model.py b/wan/modules/model.py index 5eba92b..d006f9e 100644 --- a/wan/modules/model.py +++ b/wan/modules/model.py @@ -447,6 +447,21 @@ class WanAttentionBlock(nn.Module): grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] """ + hint = None + if self.block_id is not None and hints is not None: + kwargs = { + "seq_lens" : seq_lens, + "grid_sizes" : grid_sizes, + "freqs" :freqs, + "context" : context, + "context_lens" : context_lens, + "e" : e, + } + if self.block_id == 0: + hint = self.vace(hints, x, **kwargs) + else: + hint = self.vace(hints, None, **kwargs) + e = (self.modulation + e).chunk(6, dim=1) # self-attention @@ -485,13 +500,16 @@ class WanAttentionBlock(nn.Module): x.addcmul_(y, e[5]) - - if self.block_id is not None and hints != None: + + + if hint is not None: if context_scale == 1: - x.add_(hints[self.block_id]) + x.add_(hint) else: - x.add_(hints[self.block_id], alpha =context_scale) - return x + x.add_(hint, alpha= context_scale) + return x + + class VaceWanAttentionBlock(WanAttentionBlock): def __init__( @@ -516,18 +534,29 @@ class VaceWanAttentionBlock(WanAttentionBlock): nn.init.zeros_(self.after_proj.weight) nn.init.zeros_(self.after_proj.bias) - def forward(self, c, x, **kwargs): + def forward(self, hints, x, **kwargs): # behold dbm magic ! + c = hints[0] + hints[0] = None 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 + hints[0] = c + return c_skip + + # 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): @@ -764,35 +793,37 @@ class WanModel(ModelMixin, ConfigMixin): 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] + # 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 + # ]) - return hints + # # 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, @@ -904,6 +935,34 @@ class WanModel(ModelMixin, ConfigMixin): x_list = [x] context_list = [context] del x + + # arguments + + kwargs = dict( + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=freqs, + context_lens=context_lens, + ) + + if vace_context == None: + hints_list = [None ] *len(x_list) + else: + # 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 + ]) + + kwargs['context_scale'] = vace_context_scale + hints_list = [ [c] if i==0 else [c.clone()] for i in range(len(x_list)) ] + del c + should_calc = True if self.enable_teacache: if is_uncond: @@ -935,23 +994,6 @@ class WanModel(ModelMixin, ConfigMixin): if joint_pass or not is_uncond: self.previous_residual_cond = None ori_hidden_states = x_list[0].clone() - # arguments - - kwargs = dict( - seq_lens=seq_lens, - grid_sizes=grid_sizes, - freqs=freqs, - 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 diff --git a/wan/text2video.py b/wan/text2video.py index f86284e..d77414d 100644 --- a/wan/text2video.py +++ b/wan/text2video.py @@ -143,6 +143,8 @@ class WanT2V: seq_len=32760, keep_last=True) + self.adapt_vace_model() + def vace_encode_frames(self, frames, ref_images, masks=None, tile_size = 0): if ref_images is None: ref_images = [None] * len(frames) @@ -505,3 +507,14 @@ class WanT2V: dist.barrier() return videos[0] if self.rank == 0 else None + + def adapt_vace_model(self): + model = self.model + modules_dict= { k: m for k, m in model.named_modules()} + for num in range(15): + module = modules_dict[f"vace_blocks.{num}"] + target = modules_dict[f"blocks.{2*num}"] + setattr(target, "vace", module ) + delattr(model, "vace_blocks") + + \ No newline at end of file diff --git a/wgp.py b/wgp.py index 53e6b5c..5a25676 100644 --- a/wgp.py +++ b/wgp.py @@ -910,14 +910,6 @@ def get_queue_table(queue): if len(queue) == 1: return data - # def td(l, content, width =None): - # if width !=None: - # l.append("" + content + "") - # else: - # l.append("" + content + "") - - # data.append("") - for i, item in enumerate(queue): if i==0: continue @@ -937,22 +929,7 @@ def get_queue_table(queue): start_img_md = f'Start' if end_img_uri: end_img_md = f'End' - # if i % 2 == 1: - # data.append("") - # else: - # data.append("") - # td(data,str(item.get('repeats', "1")) ) - # td(data, prompt_cell, "100%") - # td(data, num_steps, "100%") - # td(data, start_img_md) - # td(data, end_img_md) - # td(data, "↑") - # td(data, "↓") - # td(data, "✖") - # data.append("") - # data.append("
QtyPromptSteps
") - # return ''.join(data) data.append([item.get('repeats', "1"), prompt_cell, From bb73359772e322e5cffa81c6aa9aaf369541ac63 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Fri, 11 Apr 2025 08:10:01 +0200 Subject: [PATCH 57/69] Code cleaning --- wan/modules/model.py | 2 +- wan/text2video.py | 6 +++--- wgp.py | 19 +++++++++++++++---- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/wan/modules/model.py b/wan/modules/model.py index d006f9e..e7a76a9 100644 --- a/wan/modules/model.py +++ b/wan/modules/model.py @@ -960,7 +960,7 @@ class WanModel(ModelMixin, ConfigMixin): ]) kwargs['context_scale'] = vace_context_scale - hints_list = [ [c] if i==0 else [c.clone()] for i in range(len(x_list)) ] + hints_list = [ [c] for _ in range(len(x_list)) ] del c should_calc = True diff --git a/wan/text2video.py b/wan/text2video.py index d77414d..ba24e6d 100644 --- a/wan/text2video.py +++ b/wan/text2video.py @@ -511,9 +511,9 @@ class WanT2V: def adapt_vace_model(self): model = self.model modules_dict= { k: m for k, m in model.named_modules()} - for num in range(15): - module = modules_dict[f"vace_blocks.{num}"] - target = modules_dict[f"blocks.{2*num}"] + for model_layer, vace_layer in model.vace_layers_mapping.items(): + module = modules_dict[f"vace_blocks.{vace_layer}"] + target = modules_dict[f"blocks.{model_layer}"] setattr(target, "vace", module ) delattr(model, "vace_blocks") diff --git a/wgp.py b/wgp.py index eff8957..da7273a 100644 --- a/wgp.py +++ b/wgp.py @@ -421,7 +421,8 @@ def save_queue_action(state): if image_filenames_for_json: params_copy[key] = image_filenames_for_json if is_originally_list else image_filenames_for_json[0] else: - params_copy.pop(key, None) + pass + # params_copy.pop(key, None) #cant pop otherwise crash during reload for key in video_keys: video_path_orig = params_copy.get(key) @@ -885,6 +886,15 @@ def autoload_queue(state): update_global_queue_ref([]) dataframe_update = update_queue_data([]) + # need to remove queue otherwise every new tab will be processed it again + try: + if os.path.isfile(AUTOSAVE_FILENAME): + os.remove(AUTOSAVE_FILENAME) + print(f"Clear Queue: Deleted autosave file '{AUTOSAVE_FILENAME}'.") + except OSError as e: + print(f"Clear Queue: Error deleting autosave file '{AUTOSAVE_FILENAME}': {e}") + gr.Warning(f"Could not delete the autosave file '{AUTOSAVE_FILENAME}'. You may need to remove it manually.") + else: if original_queue: print("Autoload skipped: Queue is not empty.") @@ -895,6 +905,7 @@ def autoload_queue(state): update_global_queue_ref([]) dataframe_update = update_queue_data([]) + return dataframe_update, loaded_flag, state def run_autoload_and_prepare_ui(current_state): @@ -2313,10 +2324,10 @@ def generate_video( progress(*progress_args ) gen["progress_args"] = progress_args video_guide = preprocess_video(preprocess_type, width=width, height=height,video_in=video_guide, max_frames= video_length) - + image_refs = image_refs.copy() if image_refs != None else None # required since prepare_source do inplace modifications src_video, src_mask, src_ref_images = wan_model.prepare_source([video_guide], [video_mask], - [image_refs], + [image_refs], video_length, VACE_SIZE_CONFIGS[resolution_reformated], "cpu", original_video= "O" in video_prompt_type, trim_video=max_frames) @@ -3598,7 +3609,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non queue_df = gr.DataFrame( headers=["Qty","Prompt", "Length","Steps","", "", "", "", ""], datatype=[ "str","markdown","str", "markdown", "markdown", "markdown", "str", "str", "str"], - column_widths= ["5%", None, "7%", "7%", "10%", "10%", "3%", "3%", "3%"], + column_widths= ["5%", None, "7%", "7%", "10%", "10%", "3%", "3%", "34"], interactive=False, col_count=(9, "fixed"), wrap=True, From e934775eb11d0902b4e7617ec220c86f0d64a3ca Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Sat, 12 Apr 2025 02:54:54 +0200 Subject: [PATCH 58/69] Added Vace Inpainting Support and Create a Mask inside WanGP --- preprocessing/matanyone/__init__.py | 0 preprocessing/matanyone/app.py | 656 ++++++++++++++++++ .../matanyone/matanyone/config/__init__.py | 0 .../config/eval_matanyone_config.yaml | 47 ++ .../hydra/job_logging/custom-no-rank.yaml | 22 + .../config/hydra/job_logging/custom.yaml | 22 + .../matanyone/config/model/base.yaml | 58 ++ .../matanyone/matanyone/inference/__init__.py | 0 .../inference/image_feature_store.py | 56 ++ .../matanyone/inference/inference_core.py | 406 +++++++++++ .../matanyone/inference/kv_memory_store.py | 348 ++++++++++ .../matanyone/inference/memory_manager.py | 453 ++++++++++++ .../matanyone/inference/object_info.py | 24 + .../matanyone/inference/object_manager.py | 149 ++++ .../matanyone/inference/utils/__init__.py | 0 .../matanyone/inference/utils/args_utils.py | 30 + .../matanyone/matanyone/model/__init__.py | 0 .../matanyone/matanyone/model/aux_modules.py | 93 +++ .../matanyone/matanyone/model/big_modules.py | 365 ++++++++++ .../matanyone/matanyone/model/channel_attn.py | 39 ++ .../matanyone/model/group_modules.py | 126 ++++ .../matanyone/matanyone/model/matanyone.py | 333 +++++++++ .../matanyone/matanyone/model/modules.py | 149 ++++ .../matanyone/model/transformer/__init__.py | 0 .../model/transformer/object_summarizer.py | 89 +++ .../model/transformer/object_transformer.py | 206 ++++++ .../model/transformer/positional_encoding.py | 108 +++ .../model/transformer/transformer_layers.py | 161 +++++ .../matanyone/model/utils/__init__.py | 0 .../matanyone/model/utils/memory_utils.py | 107 +++ .../matanyone/model/utils/parameter_groups.py | 72 ++ .../matanyone/matanyone/model/utils/resnet.py | 179 +++++ preprocessing/matanyone/matanyone_wrapper.py | 73 ++ preprocessing/matanyone/tools/__init__.py | 0 .../matanyone/tools/base_segmenter.py | 141 ++++ .../matanyone/tools/download_util.py | 109 +++ .../matanyone/tools/interact_tools.py | 99 +++ preprocessing/matanyone/tools/mask_painter.py | 288 ++++++++ preprocessing/matanyone/tools/misc.py | 131 ++++ preprocessing/matanyone/tools/painter.py | 215 ++++++ preprocessing/matanyone/utils/__init__.py | 0 .../matanyone/utils/get_default_model.py | 27 + preprocessing/matanyone/utils/tensor_utils.py | 62 ++ requirements.txt | 2 + wan/modules/model.py | 49 +- wan/text2video.py | 7 +- wan/utils/utils.py | 5 +- wan/utils/vace_preprocessor.py | 2 +- wgp.py | 169 +++-- 49 files changed, 5578 insertions(+), 99 deletions(-) create mode 100644 preprocessing/matanyone/__init__.py create mode 100644 preprocessing/matanyone/app.py create mode 100644 preprocessing/matanyone/matanyone/config/__init__.py create mode 100644 preprocessing/matanyone/matanyone/config/eval_matanyone_config.yaml create mode 100644 preprocessing/matanyone/matanyone/config/hydra/job_logging/custom-no-rank.yaml create mode 100644 preprocessing/matanyone/matanyone/config/hydra/job_logging/custom.yaml create mode 100644 preprocessing/matanyone/matanyone/config/model/base.yaml create mode 100644 preprocessing/matanyone/matanyone/inference/__init__.py create mode 100644 preprocessing/matanyone/matanyone/inference/image_feature_store.py create mode 100644 preprocessing/matanyone/matanyone/inference/inference_core.py create mode 100644 preprocessing/matanyone/matanyone/inference/kv_memory_store.py create mode 100644 preprocessing/matanyone/matanyone/inference/memory_manager.py create mode 100644 preprocessing/matanyone/matanyone/inference/object_info.py create mode 100644 preprocessing/matanyone/matanyone/inference/object_manager.py create mode 100644 preprocessing/matanyone/matanyone/inference/utils/__init__.py create mode 100644 preprocessing/matanyone/matanyone/inference/utils/args_utils.py create mode 100644 preprocessing/matanyone/matanyone/model/__init__.py create mode 100644 preprocessing/matanyone/matanyone/model/aux_modules.py create mode 100644 preprocessing/matanyone/matanyone/model/big_modules.py create mode 100644 preprocessing/matanyone/matanyone/model/channel_attn.py create mode 100644 preprocessing/matanyone/matanyone/model/group_modules.py create mode 100644 preprocessing/matanyone/matanyone/model/matanyone.py create mode 100644 preprocessing/matanyone/matanyone/model/modules.py create mode 100644 preprocessing/matanyone/matanyone/model/transformer/__init__.py create mode 100644 preprocessing/matanyone/matanyone/model/transformer/object_summarizer.py create mode 100644 preprocessing/matanyone/matanyone/model/transformer/object_transformer.py create mode 100644 preprocessing/matanyone/matanyone/model/transformer/positional_encoding.py create mode 100644 preprocessing/matanyone/matanyone/model/transformer/transformer_layers.py create mode 100644 preprocessing/matanyone/matanyone/model/utils/__init__.py create mode 100644 preprocessing/matanyone/matanyone/model/utils/memory_utils.py create mode 100644 preprocessing/matanyone/matanyone/model/utils/parameter_groups.py create mode 100644 preprocessing/matanyone/matanyone/model/utils/resnet.py create mode 100644 preprocessing/matanyone/matanyone_wrapper.py create mode 100644 preprocessing/matanyone/tools/__init__.py create mode 100644 preprocessing/matanyone/tools/base_segmenter.py create mode 100644 preprocessing/matanyone/tools/download_util.py create mode 100644 preprocessing/matanyone/tools/interact_tools.py create mode 100644 preprocessing/matanyone/tools/mask_painter.py create mode 100644 preprocessing/matanyone/tools/misc.py create mode 100644 preprocessing/matanyone/tools/painter.py create mode 100644 preprocessing/matanyone/utils/__init__.py create mode 100644 preprocessing/matanyone/utils/get_default_model.py create mode 100644 preprocessing/matanyone/utils/tensor_utils.py diff --git a/preprocessing/matanyone/__init__.py b/preprocessing/matanyone/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/preprocessing/matanyone/app.py b/preprocessing/matanyone/app.py new file mode 100644 index 0000000..d302f03 --- /dev/null +++ b/preprocessing/matanyone/app.py @@ -0,0 +1,656 @@ +import sys + +import os +import json +import time +import psutil +import ffmpeg +import imageio +from PIL import Image + +import cv2 +import torch +import numpy as np +import gradio as gr +from .tools.painter import mask_painter +from .tools.interact_tools import SamControler +from .tools.misc import get_device +from .tools.download_util import load_file_from_url + +from .utils.get_default_model import get_matanyone_model +from .matanyone.inference.inference_core import InferenceCore +from .matanyone_wrapper import matanyone + +arg_device = "cuda" +arg_sam_model_type="vit_h" +arg_mask_save = False +model = None +matanyone_model = None + +# SAM generator +class MaskGenerator(): + def __init__(self, sam_checkpoint, device): + global args_device + args_device = device + self.samcontroler = SamControler(sam_checkpoint, arg_sam_model_type, arg_device) + + def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True): + mask, logit, painted_image = self.samcontroler.first_frame_click(image, points, labels, multimask) + return mask, logit, painted_image + +# convert points input to prompt state +def get_prompt(click_state, click_input): + inputs = json.loads(click_input) + points = click_state[0] + labels = click_state[1] + for input in inputs: + points.append(input[:2]) + labels.append(input[2]) + click_state[0] = points + click_state[1] = labels + prompt = { + "prompt_type":["click"], + "input_point":click_state[0], + "input_label":click_state[1], + "multimask_output":"True", + } + return prompt + +def get_frames_from_image(image_input, image_state): + """ + Args: + video_path:str + timestamp:float64 + Return + [[0:nearest_frame], [nearest_frame:], nearest_frame] + """ + + user_name = time.time() + frames = [image_input] * 2 # hardcode: mimic a video with 2 frames + image_size = (frames[0].shape[0],frames[0].shape[1]) + # initialize video_state + image_state = { + "user_name": user_name, + "image_name": "output.png", + "origin_images": frames, + "painted_images": frames.copy(), + "masks": [np.zeros((frames[0].shape[0],frames[0].shape[1]), np.uint8)]*len(frames), + "logits": [None]*len(frames), + "select_frame_number": 0, + "last_frame_numer": 0, + "fps": None + } + image_info = "Image Name: N/A,\nFPS: N/A,\nTotal Frames: {},\nImage Size:{}".format(len(frames), image_size) + model.samcontroler.sam_controler.reset_image() + model.samcontroler.sam_controler.set_image(image_state["origin_images"][0]) + return image_state, image_info, image_state["origin_images"][0], \ + gr.update(visible=True, maximum=10, value=10), gr.update(visible=True, maximum=len(frames), value=len(frames)), gr.update(visible=False, maximum=len(frames), value=len(frames)), \ + gr.update(visible=True), gr.update(visible=True), \ + gr.update(visible=True), gr.update(visible=True),\ + gr.update(visible=True), gr.update(visible=True), \ + gr.update(visible=True), gr.update(visible=False), \ + gr.update(visible=False), gr.update(visible=True), \ + gr.update(visible=True) + +# extract frames from upload video +def get_frames_from_video(video_input, video_state): + """ + Args: + video_path:str + timestamp:float64 + Return + [[0:nearest_frame], [nearest_frame:], nearest_frame] + """ + + while model == None: + time.sleep(1) + + video_path = video_input + frames = [] + user_name = time.time() + + # extract Audio + # try: + # audio_path = video_input.replace(".mp4", "_audio.wav") + # ffmpeg.input(video_path).output(audio_path, format='wav', acodec='pcm_s16le', ac=2, ar='44100').run(overwrite_output=True, quiet=True) + # except Exception as e: + # print(f"Audio extraction error: {str(e)}") + # audio_path = "" # Set to "" if extraction fails + # print(f'audio_path: {audio_path}') + audio_path = "" + # extract frames + try: + cap = cv2.VideoCapture(video_path) + fps = cap.get(cv2.CAP_PROP_FPS) + while cap.isOpened(): + ret, frame = cap.read() + if ret == True: + current_memory_usage = psutil.virtual_memory().percent + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + if current_memory_usage > 90: + break + else: + break + except (OSError, TypeError, ValueError, KeyError, SyntaxError) as e: + print("read_frame_source:{} error. {}\n".format(video_path, str(e))) + image_size = (frames[0].shape[0],frames[0].shape[1]) + + # resize if resolution too big + if image_size[0]>=1280 and image_size[0]>=1280: + scale = 1080 / min(image_size) + new_w = int(image_size[1] * scale) + new_h = int(image_size[0] * scale) + # update frames + frames = [cv2.resize(f, (new_w, new_h), interpolation=cv2.INTER_AREA) for f in frames] + # update image_size + image_size = (frames[0].shape[0],frames[0].shape[1]) + + # initialize video_state + video_state = { + "user_name": user_name, + "video_name": os.path.split(video_path)[-1], + "origin_images": frames, + "painted_images": frames.copy(), + "masks": [np.zeros((frames[0].shape[0],frames[0].shape[1]), np.uint8)]*len(frames), + "logits": [None]*len(frames), + "select_frame_number": 0, + "last_frame_number": 0, + "fps": fps, + "audio": audio_path + } + video_info = "Video Name: {},\nFPS: {},\nTotal Frames: {},\nImage Size:{}".format(video_state["video_name"], round(video_state["fps"], 0), len(frames), image_size) + model.samcontroler.sam_controler.reset_image() + model.samcontroler.sam_controler.set_image(video_state["origin_images"][0]) + return video_state, video_info, video_state["origin_images"][0], \ + gr.update(visible=True, maximum=len(frames), value=1), gr.update(visible=True, maximum=len(frames), value=len(frames)), gr.update(visible=False, maximum=len(frames), value=len(frames)), \ + gr.update(visible=True), gr.update(visible=True), \ + gr.update(visible=True), gr.update(visible=True),\ + gr.update(visible=True), gr.update(visible=True), \ + gr.update(visible=True), gr.update(visible=False), \ + gr.update(visible=False), gr.update(visible=True), \ + gr.update(visible=True) + +# get the select frame from gradio slider +def select_video_template(image_selection_slider, video_state, interactive_state): + + image_selection_slider -= 1 + video_state["select_frame_number"] = image_selection_slider + + # once select a new template frame, set the image in sam + model.samcontroler.sam_controler.reset_image() + model.samcontroler.sam_controler.set_image(video_state["origin_images"][image_selection_slider]) + + return video_state["painted_images"][image_selection_slider], video_state, interactive_state + +def select_image_template(image_selection_slider, video_state, interactive_state): + + image_selection_slider = 0 # fixed for image + video_state["select_frame_number"] = image_selection_slider + + # once select a new template frame, set the image in sam + model.samcontroler.sam_controler.reset_image() + model.samcontroler.sam_controler.set_image(video_state["origin_images"][image_selection_slider]) + + return video_state["painted_images"][image_selection_slider], video_state, interactive_state + +# set the tracking end frame +def get_end_number(track_pause_number_slider, video_state, interactive_state): + interactive_state["track_end_number"] = track_pause_number_slider + + return video_state["painted_images"][track_pause_number_slider],interactive_state + +# use sam to get the mask +def sam_refine(video_state, point_prompt, click_state, interactive_state, evt:gr.SelectData ): # + """ + Args: + template_frame: PIL.Image + point_prompt: flag for positive or negative button click + click_state: [[points], [labels]] + """ + if point_prompt == "Positive": + coordinate = "[[{},{},1]]".format(evt.index[0], evt.index[1]) + interactive_state["positive_click_times"] += 1 + else: + coordinate = "[[{},{},0]]".format(evt.index[0], evt.index[1]) + interactive_state["negative_click_times"] += 1 + + # prompt for sam model + model.samcontroler.sam_controler.reset_image() + model.samcontroler.sam_controler.set_image(video_state["origin_images"][video_state["select_frame_number"]]) + prompt = get_prompt(click_state=click_state, click_input=coordinate) + + mask, logit, painted_image = model.first_frame_click( + image=video_state["origin_images"][video_state["select_frame_number"]], + points=np.array(prompt["input_point"]), + labels=np.array(prompt["input_label"]), + multimask=prompt["multimask_output"], + ) + video_state["masks"][video_state["select_frame_number"]] = mask + video_state["logits"][video_state["select_frame_number"]] = logit + video_state["painted_images"][video_state["select_frame_number"]] = painted_image + + return painted_image, video_state, interactive_state + +def add_multi_mask(video_state, interactive_state, mask_dropdown): + mask = video_state["masks"][video_state["select_frame_number"]] + interactive_state["multi_mask"]["masks"].append(mask) + interactive_state["multi_mask"]["mask_names"].append("mask_{:03d}".format(len(interactive_state["multi_mask"]["masks"]))) + mask_dropdown.append("mask_{:03d}".format(len(interactive_state["multi_mask"]["masks"]))) + select_frame = show_mask(video_state, interactive_state, mask_dropdown) + + return interactive_state, gr.update(choices=interactive_state["multi_mask"]["mask_names"], value=mask_dropdown), select_frame, [[],[]] + +def clear_click(video_state, click_state): + click_state = [[],[]] + template_frame = video_state["origin_images"][video_state["select_frame_number"]] + return template_frame, click_state + +def remove_multi_mask(interactive_state, mask_dropdown): + interactive_state["multi_mask"]["mask_names"]= [] + interactive_state["multi_mask"]["masks"] = [] + + return interactive_state, gr.update(choices=[],value=[]) + +def show_mask(video_state, interactive_state, mask_dropdown): + mask_dropdown.sort() + if video_state["origin_images"]: + select_frame = video_state["origin_images"][video_state["select_frame_number"]] + for i in range(len(mask_dropdown)): + mask_number = int(mask_dropdown[i].split("_")[1]) - 1 + mask = interactive_state["multi_mask"]["masks"][mask_number] + select_frame = mask_painter(select_frame, mask.astype('uint8'), mask_color=mask_number+2) + + return select_frame + + +def save_video(frames, output_path, fps): + + writer = imageio.get_writer( output_path, fps=fps, codec='libx264', quality=8) + for frame in frames: + writer.append_data(frame) + writer.close() + + return output_path + +# video matting +def video_matting(video_state, end_slider, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size): + matanyone_processor = InferenceCore(matanyone_model, cfg=matanyone_model.cfg) + # if interactive_state["track_end_number"]: + # following_frames = video_state["origin_images"][video_state["select_frame_number"]:interactive_state["track_end_number"]] + # else: + end_slider = max(video_state["select_frame_number"] +1, end_slider) + following_frames = video_state["origin_images"][video_state["select_frame_number"]: end_slider] + + if interactive_state["multi_mask"]["masks"]: + if len(mask_dropdown) == 0: + mask_dropdown = ["mask_001"] + mask_dropdown.sort() + template_mask = interactive_state["multi_mask"]["masks"][int(mask_dropdown[0].split("_")[1]) - 1] * (int(mask_dropdown[0].split("_")[1])) + for i in range(1,len(mask_dropdown)): + mask_number = int(mask_dropdown[i].split("_")[1]) - 1 + template_mask = np.clip(template_mask+interactive_state["multi_mask"]["masks"][mask_number]*(mask_number+1), 0, mask_number+1) + video_state["masks"][video_state["select_frame_number"]]= template_mask + else: + template_mask = video_state["masks"][video_state["select_frame_number"]] + fps = video_state["fps"] + + audio_path = video_state["audio"] + + # operation error + if len(np.unique(template_mask))==1: + template_mask[0][0]=1 + foreground, alpha = matanyone(matanyone_processor, following_frames, template_mask*255, r_erode=erode_kernel_size, r_dilate=dilate_kernel_size) + output_frames = [] + for frame_origin, frame_alpha in zip(following_frames, alpha): + frame_alpha[frame_alpha > 127] = 255 + frame_alpha[frame_alpha <= 127] = 0 + output_frame = np.bitwise_and(frame_origin, 255-frame_alpha) + frame_grey = frame_alpha.copy() + frame_grey[frame_alpha == 255] = 127 + output_frame += frame_grey + output_frames.append(output_frame) + foreground = output_frames + + foreground_output = save_video(foreground, output_path="./results/{}_fg.mp4".format(video_state["video_name"]), fps=fps) + # foreground_output = generate_video_from_frames(foreground, output_path="./results/{}_fg.mp4".format(video_state["video_name"]), fps=fps, audio_path=audio_path) # import video_input to name the output video + alpha_output = save_video(alpha, output_path="./results/{}_alpha.mp4".format(video_state["video_name"]), fps=fps) + # alpha_output = generate_video_from_frames(alpha, output_path="./results/{}_alpha.mp4".format(video_state["video_name"]), fps=fps, gray2rgb=True, audio_path=audio_path) # import video_input to name the output video + + return foreground_output, alpha_output + + +def add_audio_to_video(video_path, audio_path, output_path): + try: + video_input = ffmpeg.input(video_path) + audio_input = ffmpeg.input(audio_path) + + _ = ( + ffmpeg + .output(video_input, audio_input, output_path, vcodec="copy", acodec="aac") + .run(overwrite_output=True, capture_stdout=True, capture_stderr=True) + ) + return output_path + except ffmpeg.Error as e: + print(f"FFmpeg error:\n{e.stderr.decode()}") + return None + + +def generate_video_from_frames(frames, output_path, fps=30, gray2rgb=False, audio_path=""): + """ + Generates a video from a list of frames. + + Args: + frames (list of numpy arrays): The frames to include in the video. + output_path (str): The path to save the generated video. + fps (int, optional): The frame rate of the output video. Defaults to 30. + """ + frames = torch.from_numpy(np.asarray(frames)) + _, h, w, _ = frames.shape + if gray2rgb: + frames = np.repeat(frames, 3, axis=3) + + if not os.path.exists(os.path.dirname(output_path)): + os.makedirs(os.path.dirname(output_path)) + video_temp_path = output_path.replace(".mp4", "_temp.mp4") + + # resize back to ensure input resolution + imageio.mimwrite(video_temp_path, frames, fps=fps, quality=7, + codec='libx264', ffmpeg_params=["-vf", f"scale={w}:{h}"]) + + # add audio to video if audio path exists + if audio_path != "" and os.path.exists(audio_path): + output_path = add_audio_to_video(video_temp_path, audio_path, output_path) + os.remove(video_temp_path) + return output_path + else: + return video_temp_path + +# reset all states for a new input +def restart(): + return { + "user_name": "", + "video_name": "", + "origin_images": None, + "painted_images": None, + "masks": None, + "inpaint_masks": None, + "logits": None, + "select_frame_number": 0, + "fps": 30 + }, { + "inference_times": 0, + "negative_click_times" : 0, + "positive_click_times": 0, + "mask_save": arg_mask_save, + "multi_mask": { + "mask_names": [], + "masks": [] + }, + "track_end_number": None, + }, [[],[]], None, None, \ + gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),\ + gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), \ + gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), \ + gr.update(visible=False), gr.update(visible=False, choices=[], value=[]), "", gr.update(visible=False) + +def load_unload_models(selected): + global model + global matanyone_model + if selected: + # args, defined in track_anything.py + sam_checkpoint_url_dict = { + 'vit_h': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth", + 'vit_l': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth", + 'vit_b': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth" + } + # os.path.join('.') + + from mmgp import offload + + # sam_checkpoint = load_file_from_url(sam_checkpoint_url_dict[arg_sam_model_type], ".") + sam_checkpoint = None + + transfer_stream = torch.cuda.Stream() + with torch.cuda.stream(transfer_stream): + # initialize sams + model = MaskGenerator(sam_checkpoint, "cuda") + from .matanyone.model.matanyone import MatAnyone + matanyone_model = MatAnyone.from_pretrained("PeiqingYang/MatAnyone") + # pipe ={"mat" : matanyone_model, "sam" :model.samcontroler.sam_controler.model } + # offload.profile(pipe) + matanyone_model = matanyone_model.to(arg_device).eval() + matanyone_processor = InferenceCore(matanyone_model, cfg=matanyone_model.cfg) + else: + import gc + model = None + matanyone_model = None + gc.collect() + torch.cuda.empty_cache() + + +def get_vmc_event_handler(): + return load_unload_models + +def export_to_vace_video_input(foreground_video_output): + gr.Info("Masked Video Input transferred to Vace For Inpainting") + return "V#" + str(time.time()), foreground_video_output + +def export_to_vace_video_mask(foreground_video_output, alpha_video_output): + gr.Info("Masked Video Input and Full Mask transferred to Vace For Stronger Inpainting") + return "MV#" + str(time.time()), foreground_video_output, alpha_video_output + +def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger): + # my_tab.select(fn=load_unload_models, inputs=[], outputs=[]) + + media_url = "https://github.com/pq-yang/MatAnyone/releases/download/media/" + + # download assets + + gr.Markdown("Mast Edition is provided by MatAnyone") + + with gr.Column( visible=True): + with gr.Row(): + with gr.Accordion("Video Tutorial (click to expand)", open=False, elem_classes="custom-bg"): + with gr.Row(): + with gr.Column(): + gr.Markdown("### Case 1: Single Target") + gr.Video(value="preprocessing/matanyone/tutorial_single_target.mp4", elem_classes="video") + + with gr.Column(): + gr.Markdown("### Case 2: Multiple Targets") + gr.Video(value="preprocessing/matanyone/tutorial_multi_targets.mp4", elem_classes="video") + + + click_state = gr.State([[],[]]) + + interactive_state = gr.State({ + "inference_times": 0, + "negative_click_times" : 0, + "positive_click_times": 0, + "mask_save": arg_mask_save, + "multi_mask": { + "mask_names": [], + "masks": [] + }, + "track_end_number": None, + } + ) + + video_state = gr.State( + { + "user_name": "", + "video_name": "", + "origin_images": None, + "painted_images": None, + "masks": None, + "inpaint_masks": None, + "logits": None, + "select_frame_number": 0, + "fps": 16, + "audio": "", + } + ) + + with gr.Column( visible=True): + with gr.Row(): + with gr.Accordion('MatAnyone Settings (click to expand)', open=False): + with gr.Row(): + erode_kernel_size = gr.Slider(label='Erode Kernel Size', + minimum=0, + maximum=30, + step=1, + value=10, + info="Erosion on the added mask", + interactive=True) + dilate_kernel_size = gr.Slider(label='Dilate Kernel Size', + minimum=0, + maximum=30, + step=1, + value=10, + info="Dilation on the added mask", + interactive=True) + + with gr.Row(): + image_selection_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Start Frame", info="Choose the start frame for target assignment and video matting", visible=False) + end_selection_slider = gr.Slider(minimum=1, maximum=300, step=1, value=81, label="Last Frame to Process", info="Last Frame to Process", visible=False) + + track_pause_number_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="End frame", visible=False) + with gr.Row(): + point_prompt = gr.Radio( + choices=["Positive", "Negative"], + value="Positive", + label="Point Prompt", + info="Click to add positive or negative point for target mask", + interactive=True, + visible=False, + min_width=100, + scale=1) + mask_dropdown = gr.Dropdown(multiselect=True, value=[], label="Mask Selection", info="Choose 1~all mask(s) added in Step 2", visible=False) + + gr.Markdown("---") + + with gr.Column(): + # input video + with gr.Row(equal_height=True): + with gr.Column(scale=2): + gr.Markdown("## Step1: Upload video") + with gr.Column(scale=2): + step2_title = gr.Markdown("## Step2: Add masks (Several clicks then **`Add Mask`** one by one)", visible=False) + with gr.Row(equal_height=True): + with gr.Column(scale=2): + video_input = gr.Video(label="Input Video", elem_classes="video") + extract_frames_button = gr.Button(value="Load Video", interactive=True, elem_classes="new_button") + with gr.Column(scale=2): + video_info = gr.Textbox(label="Video Info", visible=False) + template_frame = gr.Image(label="Start Frame", type="pil",interactive=True, elem_id="template_frame", visible=False, elem_classes="image") + with gr.Row(): + clear_button_click = gr.Button(value="Clear Clicks", interactive=True, visible=False, min_width=100) + add_mask_button = gr.Button(value="Add Mask", interactive=True, visible=False, min_width=100) + remove_mask_button = gr.Button(value="Remove Mask", interactive=True, visible=False, min_width=100) # no use + matting_button = gr.Button(value="Video Matting", interactive=True, visible=False, min_width=100) + with gr.Row(): + gr.Markdown("") + + # output video + with gr.Row(equal_height=True) as output_row: + with gr.Column(scale=2): + foreground_video_output = gr.Video(label="Masked Video Output", visible=False, elem_classes="video") + foreground_output_button = gr.Button(value="Black & White Video Output", visible=False, elem_classes="new_button") + export_to_vace_video_input_btn = gr.Button("Export to Vace Video Input Video For Inpainting") + with gr.Column(scale=2): + alpha_video_output = gr.Video(label="B & W Mask Video Output", visible=False, elem_classes="video") + alpha_output_button = gr.Button(value="Alpha Mask Output", visible=False, elem_classes="new_button") + export_to_vace_video_mask_btn = gr.Button("Export to Vace Video Input and Video Mask for stronger Inpainting") + + export_to_vace_video_input_btn.click(fn=export_to_vace_video_input, inputs= [foreground_video_output], outputs= [video_prompt_video_guide_trigger, vace_video_input]) + export_to_vace_video_mask_btn.click(fn=export_to_vace_video_mask, inputs= [foreground_video_output, alpha_video_output], outputs= [video_prompt_video_guide_trigger, vace_video_input, vace_video_mask]) + # first step: get the video information + extract_frames_button.click( + fn=get_frames_from_video, + inputs=[ + video_input, video_state + ], + outputs=[video_state, video_info, template_frame, + image_selection_slider, end_selection_slider, track_pause_number_slider, point_prompt, clear_button_click, add_mask_button, matting_button, template_frame, + foreground_video_output, alpha_video_output, foreground_output_button, alpha_output_button, mask_dropdown, step2_title] + ) + + # second step: select images from slider + image_selection_slider.release(fn=select_video_template, + inputs=[image_selection_slider, video_state, interactive_state], + outputs=[template_frame, video_state, interactive_state], api_name="select_image") + track_pause_number_slider.release(fn=get_end_number, + inputs=[track_pause_number_slider, video_state, interactive_state], + outputs=[template_frame, interactive_state], api_name="end_image") + + # click select image to get mask using sam + template_frame.select( + fn=sam_refine, + inputs=[video_state, point_prompt, click_state, interactive_state], + outputs=[template_frame, video_state, interactive_state] + ) + + # add different mask + add_mask_button.click( + fn=add_multi_mask, + inputs=[video_state, interactive_state, mask_dropdown], + outputs=[interactive_state, mask_dropdown, template_frame, click_state] + ) + + remove_mask_button.click( + fn=remove_multi_mask, + inputs=[interactive_state, mask_dropdown], + outputs=[interactive_state, mask_dropdown] + ) + + # video matting + matting_button.click( + fn=video_matting, + inputs=[video_state, end_selection_slider, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size], + outputs=[foreground_video_output, alpha_video_output] + ) + + # click to get mask + mask_dropdown.change( + fn=show_mask, + inputs=[video_state, interactive_state, mask_dropdown], + outputs=[template_frame] + ) + + # clear input + video_input.change( + fn=restart, + inputs=[], + outputs=[ + video_state, + interactive_state, + click_state, + foreground_video_output, alpha_video_output, + template_frame, + image_selection_slider , track_pause_number_slider,point_prompt, clear_button_click, + add_mask_button, matting_button, template_frame, foreground_video_output, alpha_video_output, remove_mask_button, foreground_output_button, alpha_output_button, mask_dropdown, video_info, step2_title + ], + queue=False, + show_progress=False) + + video_input.clear( + fn=restart, + inputs=[], + outputs=[ + video_state, + interactive_state, + click_state, + foreground_video_output, alpha_video_output, + template_frame, + image_selection_slider , track_pause_number_slider,point_prompt, clear_button_click, + add_mask_button, matting_button, template_frame, foreground_video_output, alpha_video_output, remove_mask_button, foreground_output_button, alpha_output_button, mask_dropdown, video_info, step2_title + ], + queue=False, + show_progress=False) + + # points clear + clear_button_click.click( + fn = clear_click, + inputs = [video_state, click_state,], + outputs = [template_frame,click_state], + ) diff --git a/preprocessing/matanyone/matanyone/config/__init__.py b/preprocessing/matanyone/matanyone/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/preprocessing/matanyone/matanyone/config/eval_matanyone_config.yaml b/preprocessing/matanyone/matanyone/config/eval_matanyone_config.yaml new file mode 100644 index 0000000..0c4d34f --- /dev/null +++ b/preprocessing/matanyone/matanyone/config/eval_matanyone_config.yaml @@ -0,0 +1,47 @@ +defaults: + - _self_ + - model: base + - override hydra/job_logging: custom-no-rank.yaml + +hydra: + run: + dir: ../output/${exp_id}/${dataset} + output_subdir: ${now:%Y-%m-%d_%H-%M-%S}-hydra + +amp: False +weights: pretrained_models/matanyone.pth # default (can be modified from outside) +output_dir: null # defaults to run_dir; specify this to override +flip_aug: False + + +# maximum shortest side of the input; -1 means no resizing +# With eval_vos.py, we usually just use the dataset's size (resizing done in dataloader) +# this parameter is added for the sole purpose for the GUI in the current codebase +# InferenceCore will downsize the input and restore the output to the original size if needed +# if you are using this code for some other project, you can also utilize this parameter +max_internal_size: -1 + +# these parameters, when set, override the dataset's default; useful for debugging +save_all: True +use_all_masks: False +use_long_term: False +mem_every: 5 + +# only relevant when long_term is not enabled +max_mem_frames: 5 + +# only relevant when long_term is enabled +long_term: + count_usage: True + max_mem_frames: 10 + min_mem_frames: 5 + num_prototypes: 128 + max_num_tokens: 10000 + buffer_tokens: 2000 + +top_k: 30 +stagger_updates: 5 +chunk_size: -1 # number of objects to process in parallel; -1 means unlimited +save_scores: False +save_aux: False +visualize: False diff --git a/preprocessing/matanyone/matanyone/config/hydra/job_logging/custom-no-rank.yaml b/preprocessing/matanyone/matanyone/config/hydra/job_logging/custom-no-rank.yaml new file mode 100644 index 0000000..0173c68 --- /dev/null +++ b/preprocessing/matanyone/matanyone/config/hydra/job_logging/custom-no-rank.yaml @@ -0,0 +1,22 @@ +# python logging configuration for tasks +version: 1 +formatters: + simple: + format: '[%(asctime)s][%(levelname)s] - %(message)s' + datefmt: '%Y-%m-%d %H:%M:%S' +handlers: + console: + class: logging.StreamHandler + formatter: simple + stream: ext://sys.stdout + file: + class: logging.FileHandler + formatter: simple + # absolute file path + filename: ${hydra.runtime.output_dir}/${now:%Y-%m-%d_%H-%M-%S}-eval.log + mode: w +root: + level: INFO + handlers: [console, file] + +disable_existing_loggers: false \ No newline at end of file diff --git a/preprocessing/matanyone/matanyone/config/hydra/job_logging/custom.yaml b/preprocessing/matanyone/matanyone/config/hydra/job_logging/custom.yaml new file mode 100644 index 0000000..16d4969 --- /dev/null +++ b/preprocessing/matanyone/matanyone/config/hydra/job_logging/custom.yaml @@ -0,0 +1,22 @@ +# python logging configuration for tasks +version: 1 +formatters: + simple: + format: '[%(asctime)s][%(levelname)s][r${oc.env:LOCAL_RANK}] - %(message)s' + datefmt: '%Y-%m-%d %H:%M:%S' +handlers: + console: + class: logging.StreamHandler + formatter: simple + stream: ext://sys.stdout + file: + class: logging.FileHandler + formatter: simple + # absolute file path + filename: ${hydra.runtime.output_dir}/${now:%Y-%m-%d_%H-%M-%S}-rank${oc.env:LOCAL_RANK}.log + mode: w +root: + level: INFO + handlers: [console, file] + +disable_existing_loggers: false \ No newline at end of file diff --git a/preprocessing/matanyone/matanyone/config/model/base.yaml b/preprocessing/matanyone/matanyone/config/model/base.yaml new file mode 100644 index 0000000..3d64dcc --- /dev/null +++ b/preprocessing/matanyone/matanyone/config/model/base.yaml @@ -0,0 +1,58 @@ +pixel_mean: [0.485, 0.456, 0.406] +pixel_std: [0.229, 0.224, 0.225] + +pixel_dim: 256 +key_dim: 64 +value_dim: 256 +sensory_dim: 256 +embed_dim: 256 + +pixel_encoder: + type: resnet50 + ms_dims: [1024, 512, 256, 64, 3] # f16, f8, f4, f2, f1 + +mask_encoder: + type: resnet18 + final_dim: 256 + +pixel_pe_scale: 32 +pixel_pe_temperature: 128 + +object_transformer: + embed_dim: ${model.embed_dim} + ff_dim: 2048 + num_heads: 8 + num_blocks: 3 + num_queries: 16 + read_from_pixel: + input_norm: False + input_add_pe: False + add_pe_to_qkv: [True, True, False] + read_from_past: + add_pe_to_qkv: [True, True, False] + read_from_memory: + add_pe_to_qkv: [True, True, False] + read_from_query: + add_pe_to_qkv: [True, True, False] + output_norm: False + query_self_attention: + add_pe_to_qkv: [True, True, False] + pixel_self_attention: + add_pe_to_qkv: [True, True, False] + +object_summarizer: + embed_dim: ${model.object_transformer.embed_dim} + num_summaries: ${model.object_transformer.num_queries} + add_pe: True + +aux_loss: + sensory: + enabled: True + weight: 0.01 + query: + enabled: True + weight: 0.01 + +mask_decoder: + # first value must equal embed_dim + up_dims: [256, 128, 128, 64, 16] diff --git a/preprocessing/matanyone/matanyone/inference/__init__.py b/preprocessing/matanyone/matanyone/inference/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/preprocessing/matanyone/matanyone/inference/image_feature_store.py b/preprocessing/matanyone/matanyone/inference/image_feature_store.py new file mode 100644 index 0000000..7195b05 --- /dev/null +++ b/preprocessing/matanyone/matanyone/inference/image_feature_store.py @@ -0,0 +1,56 @@ +import warnings +from typing import Iterable +import torch +from ..model.matanyone import MatAnyone + + +class ImageFeatureStore: + """ + A cache for image features. + These features might be reused at different parts of the inference pipeline. + This class provide an interface for reusing these features. + It is the user's responsibility to delete redundant features. + + Feature of a frame should be associated with a unique index -- typically the frame id. + """ + def __init__(self, network: MatAnyone, no_warning: bool = False): + self.network = network + self._store = {} + self.no_warning = no_warning + + def _encode_feature(self, index: int, image: torch.Tensor, last_feats=None) -> None: + ms_features, pix_feat = self.network.encode_image(image, last_feats=last_feats) + key, shrinkage, selection = self.network.transform_key(ms_features[0]) + self._store[index] = (ms_features, pix_feat, key, shrinkage, selection) + + def get_all_features(self, images: torch.Tensor) -> (Iterable[torch.Tensor], torch.Tensor): + seq_length = images.shape[0] + ms_features, pix_feat = self.network.encode_image(images, seq_length) + key, shrinkage, selection = self.network.transform_key(ms_features[0]) + for index in range(seq_length): + self._store[index] = ([f[index].unsqueeze(0) for f in ms_features], pix_feat[index].unsqueeze(0), key[index].unsqueeze(0), shrinkage[index].unsqueeze(0), selection[index].unsqueeze(0)) + + def get_features(self, index: int, + image: torch.Tensor, last_feats=None) -> (Iterable[torch.Tensor], torch.Tensor): + if index not in self._store: + self._encode_feature(index, image, last_feats) + + return self._store[index][:2] + + def get_key(self, index: int, + image: torch.Tensor, last_feats=None) -> (torch.Tensor, torch.Tensor, torch.Tensor): + if index not in self._store: + self._encode_feature(index, image, last_feats) + + return self._store[index][2:] + + def delete(self, index: int) -> None: + if index in self._store: + del self._store[index] + + def __len__(self): + return len(self._store) + + def __del__(self): + if len(self._store) > 0 and not self.no_warning: + warnings.warn(f'Leaking {self._store.keys()} in the image feature store') diff --git a/preprocessing/matanyone/matanyone/inference/inference_core.py b/preprocessing/matanyone/matanyone/inference/inference_core.py new file mode 100644 index 0000000..12a6365 --- /dev/null +++ b/preprocessing/matanyone/matanyone/inference/inference_core.py @@ -0,0 +1,406 @@ +from typing import List, Optional, Iterable +import logging +from omegaconf import DictConfig + +import numpy as np +import torch +import torch.nn.functional as F + +from .memory_manager import MemoryManager +from .object_manager import ObjectManager +from .image_feature_store import ImageFeatureStore +from ..model.matanyone import MatAnyone +from ...utils.tensor_utils import pad_divide_by, unpad, aggregate + +log = logging.getLogger() + + +class InferenceCore: + + def __init__(self, + network: MatAnyone, + cfg: DictConfig, + *, + image_feature_store: ImageFeatureStore = None): + self.network = network + self.cfg = cfg + self.mem_every = cfg.mem_every + stagger_updates = cfg.stagger_updates + self.chunk_size = cfg.chunk_size + self.save_aux = cfg.save_aux + self.max_internal_size = cfg.max_internal_size + self.flip_aug = cfg.flip_aug + + self.curr_ti = -1 + self.last_mem_ti = 0 + # at which time indices should we update the sensory memory + if stagger_updates >= self.mem_every: + self.stagger_ti = set(range(1, self.mem_every + 1)) + else: + self.stagger_ti = set( + np.round(np.linspace(1, self.mem_every, stagger_updates)).astype(int)) + self.object_manager = ObjectManager() + self.memory = MemoryManager(cfg=cfg, object_manager=self.object_manager) + + if image_feature_store is None: + self.image_feature_store = ImageFeatureStore(self.network) + else: + self.image_feature_store = image_feature_store + + self.last_mask = None + self.last_pix_feat = None + self.last_msk_value = None + + def clear_memory(self): + self.curr_ti = -1 + self.last_mem_ti = 0 + self.memory = MemoryManager(cfg=self.cfg, object_manager=self.object_manager) + + def clear_non_permanent_memory(self): + self.curr_ti = -1 + self.last_mem_ti = 0 + self.memory.clear_non_permanent_memory() + + def clear_sensory_memory(self): + self.curr_ti = -1 + self.last_mem_ti = 0 + self.memory.clear_sensory_memory() + + def update_config(self, cfg): + self.mem_every = cfg['mem_every'] + self.memory.update_config(cfg) + + def clear_temp_mem(self): + self.memory.clear_work_mem() + # self.object_manager = ObjectManager() + self.memory.clear_obj_mem() + # self.memory.clear_sensory_memory() + + def _add_memory(self, + image: torch.Tensor, + pix_feat: torch.Tensor, + prob: torch.Tensor, + key: torch.Tensor, + shrinkage: torch.Tensor, + selection: torch.Tensor, + *, + is_deep_update: bool = True, + force_permanent: bool = False) -> None: + """ + Memorize the given segmentation in all memory stores. + + The batch dimension is 1 if flip augmentation is not used. + image: RGB image, (1/2)*3*H*W + pix_feat: from the key encoder, (1/2)*_*H*W + prob: (1/2)*num_objects*H*W, in [0, 1] + key/shrinkage/selection: for anisotropic l2, (1/2)*_*H*W + selection can be None if not using long-term memory + is_deep_update: whether to use deep update (e.g. with the mask encoder) + force_permanent: whether to force the memory to be permanent + """ + if prob.shape[1] == 0: + # nothing to add + log.warn('Trying to add an empty object mask to memory!') + return + + if force_permanent: + as_permanent = 'all' + else: + as_permanent = 'first' + + self.memory.initialize_sensory_if_needed(key, self.object_manager.all_obj_ids) + msk_value, sensory, obj_value, _ = self.network.encode_mask( + image, + pix_feat, + self.memory.get_sensory(self.object_manager.all_obj_ids), + prob, + deep_update=is_deep_update, + chunk_size=self.chunk_size, + need_weights=self.save_aux) + self.memory.add_memory(key, + shrinkage, + msk_value, + obj_value, + self.object_manager.all_obj_ids, + selection=selection, + as_permanent=as_permanent) + self.last_mem_ti = self.curr_ti + if is_deep_update: + self.memory.update_sensory(sensory, self.object_manager.all_obj_ids) + self.last_msk_value = msk_value + + def _segment(self, + key: torch.Tensor, + selection: torch.Tensor, + pix_feat: torch.Tensor, + ms_features: Iterable[torch.Tensor], + update_sensory: bool = True) -> torch.Tensor: + """ + Produce a segmentation using the given features and the memory + + The batch dimension is 1 if flip augmentation is not used. + key/selection: for anisotropic l2: (1/2) * _ * H * W + pix_feat: from the key encoder, (1/2) * _ * H * W + ms_features: an iterable of multiscale features from the encoder, each is (1/2)*_*H*W + with strides 16, 8, and 4 respectively + update_sensory: whether to update the sensory memory + + Returns: (num_objects+1)*H*W normalized probability; the first channel is the background + """ + bs = key.shape[0] + if self.flip_aug: + assert bs == 2 + else: + assert bs == 1 + + if not self.memory.engaged: + log.warn('Trying to segment without any memory!') + return torch.zeros((1, key.shape[-2] * 16, key.shape[-1] * 16), + device=key.device, + dtype=key.dtype) + + uncert_output = None + + if self.curr_ti == 0: # ONLY for the first frame for prediction + memory_readout = self.memory.read_first_frame(self.last_msk_value, pix_feat, self.last_mask, self.network, uncert_output=uncert_output) + else: + memory_readout = self.memory.read(pix_feat, key, selection, self.last_mask, self.network, uncert_output=uncert_output, last_msk_value=self.last_msk_value, ti=self.curr_ti, + last_pix_feat=self.last_pix_feat, last_pred_mask=self.last_mask) + memory_readout = self.object_manager.realize_dict(memory_readout) + + sensory, _, pred_prob_with_bg = self.network.segment(ms_features, + memory_readout, + self.memory.get_sensory( + self.object_manager.all_obj_ids), + chunk_size=self.chunk_size, + update_sensory=update_sensory) + # remove batch dim + if self.flip_aug: + # average predictions of the non-flipped and flipped version + pred_prob_with_bg = (pred_prob_with_bg[0] + + torch.flip(pred_prob_with_bg[1], dims=[-1])) / 2 + else: + pred_prob_with_bg = pred_prob_with_bg[0] + if update_sensory: + self.memory.update_sensory(sensory, self.object_manager.all_obj_ids) + return pred_prob_with_bg + + def pred_all_flow(self, images): + self.total_len = images.shape[0] + images, self.pad = pad_divide_by(images, 16) + images = images.unsqueeze(0) # add the batch dimension: (1,t,c,h,w) + + self.flows_forward, self.flows_backward = self.network.pred_forward_backward_flow(images) + + def encode_all_images(self, images): + images, self.pad = pad_divide_by(images, 16) + self.image_feature_store.get_all_features(images) # t c h w + return images + + def step(self, + image: torch.Tensor, + mask: Optional[torch.Tensor] = None, + objects: Optional[List[int]] = None, + *, + idx_mask: bool = False, + end: bool = False, + delete_buffer: bool = True, + force_permanent: bool = False, + matting: bool = True, + first_frame_pred: bool = False) -> torch.Tensor: + """ + Take a step with a new incoming image. + If there is an incoming mask with new objects, we will memorize them. + If there is no incoming mask, we will segment the image using the memory. + In both cases, we will update the memory and return a segmentation. + + image: 3*H*W + mask: H*W (if idx mask) or len(objects)*H*W or None + objects: list of object ids that are valid in the mask Tensor. + The ids themselves do not need to be consecutive/in order, but they need to be + in the same position in the list as the corresponding mask + in the tensor in non-idx-mask mode. + objects is ignored if the mask is None. + If idx_mask is False and objects is None, we sequentially infer the object ids. + idx_mask: if True, mask is expected to contain an object id at every pixel. + If False, mask should have multiple channels with each channel representing one object. + end: if we are at the end of the sequence, we do not need to update memory + if unsure just set it to False + delete_buffer: whether to delete the image feature buffer after this step + force_permanent: the memory recorded this frame will be added to the permanent memory + """ + if objects is None and mask is not None: + assert not idx_mask + objects = list(range(1, mask.shape[0] + 1)) + + # resize input if needed -- currently only used for the GUI + resize_needed = False + if self.max_internal_size > 0: + h, w = image.shape[-2:] + min_side = min(h, w) + if min_side > self.max_internal_size: + resize_needed = True + new_h = int(h / min_side * self.max_internal_size) + new_w = int(w / min_side * self.max_internal_size) + image = F.interpolate(image.unsqueeze(0), + size=(new_h, new_w), + mode='bilinear', + align_corners=False)[0] + if mask is not None: + if idx_mask: + mask = F.interpolate(mask.unsqueeze(0).unsqueeze(0).float(), + size=(new_h, new_w), + mode='nearest-exact', + align_corners=False)[0, 0].round().long() + else: + mask = F.interpolate(mask.unsqueeze(0), + size=(new_h, new_w), + mode='bilinear', + align_corners=False)[0] + + self.curr_ti += 1 + + image, self.pad = pad_divide_by(image, 16) # DONE alreay for 3DCNN!! + image = image.unsqueeze(0) # add the batch dimension + if self.flip_aug: + image = torch.cat([image, torch.flip(image, dims=[-1])], dim=0) + + # whether to update the working memory + is_mem_frame = ((self.curr_ti - self.last_mem_ti >= self.mem_every) or + (mask is not None)) and (not end) + # segment when there is no input mask or when the input mask is incomplete + need_segment = (mask is None) or (self.object_manager.num_obj > 0 + and not self.object_manager.has_all(objects)) + update_sensory = ((self.curr_ti - self.last_mem_ti) in self.stagger_ti) and (not end) + + # reinit if it is the first frame for prediction + if first_frame_pred: + self.curr_ti = 0 + self.last_mem_ti = 0 + is_mem_frame = True + need_segment = True + update_sensory = True + + # encoding the image + ms_feat, pix_feat = self.image_feature_store.get_features(self.curr_ti, image) + key, shrinkage, selection = self.image_feature_store.get_key(self.curr_ti, image) + + # segmentation from memory if needed + if need_segment: + pred_prob_with_bg = self._segment(key, + selection, + pix_feat, + ms_feat, + update_sensory=update_sensory) + + # use the input mask if provided + if mask is not None: + # inform the manager of the new objects, and get a list of temporary id + # temporary ids -- indicates the position of objects in the tensor + # (starts with 1 due to the background channel) + corresponding_tmp_ids, _ = self.object_manager.add_new_objects(objects) + + mask, _ = pad_divide_by(mask, 16) + if need_segment: + # merge predicted mask with the incomplete input mask + pred_prob_no_bg = pred_prob_with_bg[1:] + # use the mutual exclusivity of segmentation + if idx_mask: + pred_prob_no_bg[:, mask > 0] = 0 + else: + pred_prob_no_bg[:, mask.max(0) > 0.5] = 0 + + new_masks = [] + for mask_id, tmp_id in enumerate(corresponding_tmp_ids): + if idx_mask: + this_mask = (mask == objects[mask_id]).type_as(pred_prob_no_bg) + else: + this_mask = mask[tmp_id] + if tmp_id > pred_prob_no_bg.shape[0]: + new_masks.append(this_mask.unsqueeze(0)) + else: + # +1 for padding the background channel + pred_prob_no_bg[tmp_id - 1] = this_mask + # new_masks are always in the order of tmp_id + mask = torch.cat([pred_prob_no_bg, *new_masks], dim=0) + elif idx_mask: + # simply convert cls to one-hot representation + if len(objects) == 0: + if delete_buffer: + self.image_feature_store.delete(self.curr_ti) + log.warn('Trying to insert an empty mask as memory!') + return torch.zeros((1, key.shape[-2] * 16, key.shape[-1] * 16), + device=key.device, + dtype=key.dtype) + mask = torch.stack( + [mask == objects[mask_id] for mask_id, _ in enumerate(corresponding_tmp_ids)], + dim=0) + if matting: + mask = mask.unsqueeze(0).float() / 255. + pred_prob_with_bg = torch.cat([1-mask, mask], 0) + else: + pred_prob_with_bg = aggregate(mask, dim=0) + pred_prob_with_bg = torch.softmax(pred_prob_with_bg, dim=0) + + self.last_mask = pred_prob_with_bg[1:].unsqueeze(0) + if self.flip_aug: + self.last_mask = torch.cat( + [self.last_mask, torch.flip(self.last_mask, dims=[-1])], dim=0) + self.last_pix_feat = pix_feat + + # save as memory if needed + if is_mem_frame or force_permanent: + # clear the memory for given mask and add the first predicted mask + if first_frame_pred: + self.clear_temp_mem() + self._add_memory(image, + pix_feat, + self.last_mask, + key, + shrinkage, + selection, + force_permanent=force_permanent, + is_deep_update=True) + else: # compute self.last_msk_value for non-memory frame + msk_value, _, _, _ = self.network.encode_mask( + image, + pix_feat, + self.memory.get_sensory(self.object_manager.all_obj_ids), + self.last_mask, + deep_update=False, + chunk_size=self.chunk_size, + need_weights=self.save_aux) + self.last_msk_value = msk_value + + if delete_buffer: + self.image_feature_store.delete(self.curr_ti) + + output_prob = unpad(pred_prob_with_bg, self.pad) + if resize_needed: + # restore output to the original size + output_prob = F.interpolate(output_prob.unsqueeze(0), + size=(h, w), + mode='bilinear', + align_corners=False)[0] + + return output_prob + + def delete_objects(self, objects: List[int]) -> None: + """ + Delete the given objects from the memory. + """ + self.object_manager.delete_objects(objects) + self.memory.purge_except(self.object_manager.all_obj_ids) + + def output_prob_to_mask(self, output_prob: torch.Tensor, matting: bool = True) -> torch.Tensor: + if matting: + new_mask = output_prob[1:].squeeze(0) + else: + mask = torch.argmax(output_prob, dim=0) + + # index in tensor != object id -- remap the ids here + new_mask = torch.zeros_like(mask) + for tmp_id, obj in self.object_manager.tmp_id_to_obj.items(): + new_mask[mask == tmp_id] = obj.id + + return new_mask diff --git a/preprocessing/matanyone/matanyone/inference/kv_memory_store.py b/preprocessing/matanyone/matanyone/inference/kv_memory_store.py new file mode 100644 index 0000000..e50b794 --- /dev/null +++ b/preprocessing/matanyone/matanyone/inference/kv_memory_store.py @@ -0,0 +1,348 @@ +from typing import Dict, List, Optional, Literal +from collections import defaultdict +import torch + + +def _add_last_dim(dictionary, key, new_value, prepend=False): + # append/prepend a new value to the last dimension of a tensor in a dictionary + # if the key does not exist, put the new value in + # append by default + if key in dictionary: + dictionary[key] = torch.cat([dictionary[key], new_value], -1) + else: + dictionary[key] = new_value + + +class KeyValueMemoryStore: + """ + Works for key/value pairs type storage + e.g., working and long-term memory + """ + def __init__(self, save_selection: bool = False, save_usage: bool = False): + """ + We store keys and values of objects that first appear in the same frame in a bucket. + Each bucket contains a set of object ids. + Each bucket is associated with a single key tensor + and a dictionary of value tensors indexed by object id. + + The keys and values are stored as the concatenation of a permanent part and a temporary part. + """ + self.save_selection = save_selection + self.save_usage = save_usage + + self.global_bucket_id = 0 # does not reduce even if buckets are removed + self.buckets: Dict[int, List[int]] = {} # indexed by bucket id + self.k: Dict[int, torch.Tensor] = {} # indexed by bucket id + self.v: Dict[int, torch.Tensor] = {} # indexed by object id + + # indexed by bucket id; the end point of permanent memory + self.perm_end_pt: Dict[int, int] = defaultdict(int) + + # shrinkage and selection are just like the keys + self.s = {} + if self.save_selection: + self.e = {} # does not contain the permanent memory part + + # usage + if self.save_usage: + self.use_cnt = {} # indexed by bucket id, does not contain the permanent memory part + self.life_cnt = {} # indexed by bucket id, does not contain the permanent memory part + + def add(self, + key: torch.Tensor, + values: Dict[int, torch.Tensor], + shrinkage: torch.Tensor, + selection: torch.Tensor, + supposed_bucket_id: int = -1, + as_permanent: Literal['no', 'first', 'all'] = 'no') -> None: + """ + key: (1/2)*C*N + values: dict of values ((1/2)*C*N), object ids are used as keys + shrinkage: (1/2)*1*N + selection: (1/2)*C*N + + supposed_bucket_id: used to sync the bucket id between working and long-term memory + if provided, the input should all be in a single bucket indexed by this id + as_permanent: whether to store the input as permanent memory + 'no': don't + 'first': only store it as permanent memory if the bucket is empty + 'all': always store it as permanent memory + """ + bs = key.shape[0] + ne = key.shape[-1] + assert len(key.shape) == 3 + assert len(shrinkage.shape) == 3 + assert not self.save_selection or len(selection.shape) == 3 + assert as_permanent in ['no', 'first', 'all'] + + # add the value and create new buckets if necessary + if supposed_bucket_id >= 0: + enabled_buckets = [supposed_bucket_id] + bucket_exist = supposed_bucket_id in self.buckets + for obj, value in values.items(): + if bucket_exist: + assert obj in self.v + assert obj in self.buckets[supposed_bucket_id] + _add_last_dim(self.v, obj, value, prepend=(as_permanent == 'all')) + else: + assert obj not in self.v + self.v[obj] = value + self.buckets[supposed_bucket_id] = list(values.keys()) + else: + new_bucket_id = None + enabled_buckets = set() + for obj, value in values.items(): + assert len(value.shape) == 3 + if obj in self.v: + _add_last_dim(self.v, obj, value, prepend=(as_permanent == 'all')) + bucket_used = [ + bucket_id for bucket_id, object_ids in self.buckets.items() + if obj in object_ids + ] + assert len(bucket_used) == 1 # each object should only be in one bucket + enabled_buckets.add(bucket_used[0]) + else: + self.v[obj] = value + if new_bucket_id is None: + # create new bucket + new_bucket_id = self.global_bucket_id + self.global_bucket_id += 1 + self.buckets[new_bucket_id] = [] + # put the new object into the corresponding bucket + self.buckets[new_bucket_id].append(obj) + enabled_buckets.add(new_bucket_id) + + # increment the permanent size if necessary + add_as_permanent = {} # indexed by bucket id + for bucket_id in enabled_buckets: + add_as_permanent[bucket_id] = False + if as_permanent == 'all': + self.perm_end_pt[bucket_id] += ne + add_as_permanent[bucket_id] = True + elif as_permanent == 'first': + if self.perm_end_pt[bucket_id] == 0: + self.perm_end_pt[bucket_id] = ne + add_as_permanent[bucket_id] = True + + # create new counters for usage if necessary + if self.save_usage and as_permanent != 'all': + new_count = torch.zeros((bs, ne), device=key.device, dtype=torch.float32) + new_life = torch.zeros((bs, ne), device=key.device, dtype=torch.float32) + 1e-7 + + # add the key to every bucket + for bucket_id in self.buckets: + if bucket_id not in enabled_buckets: + # if we are not adding new values to a bucket, we should skip it + continue + + _add_last_dim(self.k, bucket_id, key, prepend=add_as_permanent[bucket_id]) + _add_last_dim(self.s, bucket_id, shrinkage, prepend=add_as_permanent[bucket_id]) + if not add_as_permanent[bucket_id]: + if self.save_selection: + _add_last_dim(self.e, bucket_id, selection) + if self.save_usage: + _add_last_dim(self.use_cnt, bucket_id, new_count) + _add_last_dim(self.life_cnt, bucket_id, new_life) + + def update_bucket_usage(self, bucket_id: int, usage: torch.Tensor) -> None: + # increase all life count by 1 + # increase use of indexed elements + if not self.save_usage: + return + + usage = usage[:, self.perm_end_pt[bucket_id]:] + if usage.shape[-1] == 0: + # if there is no temporary memory, we don't need to update + return + self.use_cnt[bucket_id] += usage.view_as(self.use_cnt[bucket_id]) + self.life_cnt[bucket_id] += 1 + + def sieve_by_range(self, bucket_id: int, start: int, end: int, min_size: int) -> None: + # keep only the temporary elements *outside* of this range (with some boundary conditions) + # the permanent elements are ignored in this computation + # i.e., concat (a[:start], a[end:]) + # bucket with size <= min_size are not modified + + assert start >= 0 + assert end <= 0 + + object_ids = self.buckets[bucket_id] + bucket_num_elements = self.k[bucket_id].shape[-1] - self.perm_end_pt[bucket_id] + if bucket_num_elements <= min_size: + return + + if end == 0: + # negative 0 would not work as the end index! + # effectively make the second part an empty slice + end = self.k[bucket_id].shape[-1] + 1 + + p_size = self.perm_end_pt[bucket_id] + start = start + p_size + + k = self.k[bucket_id] + s = self.s[bucket_id] + if self.save_selection: + e = self.e[bucket_id] + if self.save_usage: + use_cnt = self.use_cnt[bucket_id] + life_cnt = self.life_cnt[bucket_id] + + self.k[bucket_id] = torch.cat([k[:, :, :start], k[:, :, end:]], -1) + self.s[bucket_id] = torch.cat([s[:, :, :start], s[:, :, end:]], -1) + if self.save_selection: + self.e[bucket_id] = torch.cat([e[:, :, :start - p_size], e[:, :, end:]], -1) + if self.save_usage: + self.use_cnt[bucket_id] = torch.cat([use_cnt[:, :start - p_size], use_cnt[:, end:]], -1) + self.life_cnt[bucket_id] = torch.cat([life_cnt[:, :start - p_size], life_cnt[:, end:]], + -1) + for obj_id in object_ids: + v = self.v[obj_id] + self.v[obj_id] = torch.cat([v[:, :, :start], v[:, :, end:]], -1) + + def remove_old_memory(self, bucket_id: int, max_len: int) -> None: + self.sieve_by_range(bucket_id, 0, -max_len, max_len) + + def remove_obsolete_features(self, bucket_id: int, max_size: int) -> None: + # for long-term memory only + object_ids = self.buckets[bucket_id] + + assert self.perm_end_pt[bucket_id] == 0 # permanent memory should be empty in LT memory + + # normalize with life duration + usage = self.get_usage(bucket_id) + bs = usage.shape[0] + + survivals = [] + + for bi in range(bs): + _, survived = torch.topk(usage[bi], k=max_size) + survivals.append(survived.flatten()) + assert survived.shape[-1] == survivals[0].shape[-1] + + self.k[bucket_id] = torch.stack( + [self.k[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0) + self.s[bucket_id] = torch.stack( + [self.s[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0) + + if self.save_selection: + # Long-term memory does not store selection so this should not be needed + self.e[bucket_id] = torch.stack( + [self.e[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0) + for obj_id in object_ids: + self.v[obj_id] = torch.stack( + [self.v[obj_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0) + + self.use_cnt[bucket_id] = torch.stack( + [self.use_cnt[bucket_id][bi, survived] for bi, survived in enumerate(survivals)], 0) + self.life_cnt[bucket_id] = torch.stack( + [self.life_cnt[bucket_id][bi, survived] for bi, survived in enumerate(survivals)], 0) + + def get_usage(self, bucket_id: int) -> torch.Tensor: + # return normalized usage + if not self.save_usage: + raise RuntimeError('I did not count usage!') + else: + usage = self.use_cnt[bucket_id] / self.life_cnt[bucket_id] + return usage + + def get_all_sliced( + self, bucket_id: int, start: int, end: int + ) -> (torch.Tensor, torch.Tensor, torch.Tensor, Dict[int, torch.Tensor], torch.Tensor): + # return k, sk, ek, value, normalized usage in order, sliced by start and end + # this only queries the temporary memory + + assert start >= 0 + assert end <= 0 + + p_size = self.perm_end_pt[bucket_id] + start = start + p_size + + if end == 0: + # negative 0 would not work as the end index! + k = self.k[bucket_id][:, :, start:] + sk = self.s[bucket_id][:, :, start:] + ek = self.e[bucket_id][:, :, start - p_size:] if self.save_selection else None + value = {obj_id: self.v[obj_id][:, :, start:] for obj_id in self.buckets[bucket_id]} + usage = self.get_usage(bucket_id)[:, start - p_size:] if self.save_usage else None + else: + k = self.k[bucket_id][:, :, start:end] + sk = self.s[bucket_id][:, :, start:end] + ek = self.e[bucket_id][:, :, start - p_size:end] if self.save_selection else None + value = {obj_id: self.v[obj_id][:, :, start:end] for obj_id in self.buckets[bucket_id]} + usage = self.get_usage(bucket_id)[:, start - p_size:end] if self.save_usage else None + + return k, sk, ek, value, usage + + def purge_except(self, obj_keep_idx: List[int]): + # purge certain objects from the memory except the one listed + obj_keep_idx = set(obj_keep_idx) + + # remove objects that are not in the keep list from the buckets + buckets_to_remove = [] + for bucket_id, object_ids in self.buckets.items(): + self.buckets[bucket_id] = [obj_id for obj_id in object_ids if obj_id in obj_keep_idx] + if len(self.buckets[bucket_id]) == 0: + buckets_to_remove.append(bucket_id) + + # remove object values that are not in the keep list + self.v = {k: v for k, v in self.v.items() if k in obj_keep_idx} + + # remove buckets that are empty + for bucket_id in buckets_to_remove: + del self.buckets[bucket_id] + del self.k[bucket_id] + del self.s[bucket_id] + if self.save_selection: + del self.e[bucket_id] + if self.save_usage: + del self.use_cnt[bucket_id] + del self.life_cnt[bucket_id] + + def clear_non_permanent_memory(self): + # clear all non-permanent memory + for bucket_id in self.buckets: + self.sieve_by_range(bucket_id, 0, 0, 0) + + def get_v_size(self, obj_id: int) -> int: + return self.v[obj_id].shape[-1] + + def size(self, bucket_id: int) -> int: + if bucket_id not in self.k: + return 0 + else: + return self.k[bucket_id].shape[-1] + + def perm_size(self, bucket_id: int) -> int: + return self.perm_end_pt[bucket_id] + + def non_perm_size(self, bucket_id: int) -> int: + return self.size(bucket_id) - self.perm_size(bucket_id) + + def engaged(self, bucket_id: Optional[int] = None) -> bool: + if bucket_id is None: + return len(self.buckets) > 0 + else: + return bucket_id in self.buckets + + @property + def num_objects(self) -> int: + return len(self.v) + + @property + def key(self) -> Dict[int, torch.Tensor]: + return self.k + + @property + def value(self) -> Dict[int, torch.Tensor]: + return self.v + + @property + def shrinkage(self) -> Dict[int, torch.Tensor]: + return self.s + + @property + def selection(self) -> Dict[int, torch.Tensor]: + return self.e + + def __contains__(self, key): + return key in self.v diff --git a/preprocessing/matanyone/matanyone/inference/memory_manager.py b/preprocessing/matanyone/matanyone/inference/memory_manager.py new file mode 100644 index 0000000..b70664c --- /dev/null +++ b/preprocessing/matanyone/matanyone/inference/memory_manager.py @@ -0,0 +1,453 @@ +import logging +from omegaconf import DictConfig +from typing import List, Dict +import torch + +from .object_manager import ObjectManager +from .kv_memory_store import KeyValueMemoryStore +from ..model.matanyone import MatAnyone +from ..model.utils.memory_utils import get_similarity, do_softmax + +log = logging.getLogger() + + +class MemoryManager: + """ + Manages all three memory stores and the transition between working/long-term memory + """ + def __init__(self, cfg: DictConfig, object_manager: ObjectManager): + self.object_manager = object_manager + self.sensory_dim = cfg.model.sensory_dim + self.top_k = cfg.top_k + self.chunk_size = cfg.chunk_size + + self.save_aux = cfg.save_aux + + self.use_long_term = cfg.use_long_term + self.count_long_term_usage = cfg.long_term.count_usage + # subtract 1 because the first-frame is now counted as "permanent memory" + # and is not counted towards max_mem_frames + # but we want to keep the hyperparameters consistent as before for the same behavior + if self.use_long_term: + self.max_mem_frames = cfg.long_term.max_mem_frames - 1 + self.min_mem_frames = cfg.long_term.min_mem_frames - 1 + self.num_prototypes = cfg.long_term.num_prototypes + self.max_long_tokens = cfg.long_term.max_num_tokens + self.buffer_tokens = cfg.long_term.buffer_tokens + else: + self.max_mem_frames = cfg.max_mem_frames - 1 + + # dimensions will be inferred from input later + self.CK = self.CV = None + self.H = self.W = None + + # The sensory memory is stored as a dictionary indexed by object ids + # each of shape bs * C^h * H * W + self.sensory = {} + + # a dictionary indexed by object ids, each of shape bs * T * Q * C + self.obj_v = {} + + self.work_mem = KeyValueMemoryStore(save_selection=self.use_long_term, + save_usage=self.use_long_term) + if self.use_long_term: + self.long_mem = KeyValueMemoryStore(save_usage=self.count_long_term_usage) + + self.config_stale = True + self.engaged = False + + def update_config(self, cfg: DictConfig) -> None: + self.config_stale = True + self.top_k = cfg['top_k'] + + assert self.use_long_term == cfg.use_long_term, 'cannot update this' + assert self.count_long_term_usage == cfg.long_term.count_usage, 'cannot update this' + + self.use_long_term = cfg.use_long_term + self.count_long_term_usage = cfg.long_term.count_usage + if self.use_long_term: + self.max_mem_frames = cfg.long_term.max_mem_frames - 1 + self.min_mem_frames = cfg.long_term.min_mem_frames - 1 + self.num_prototypes = cfg.long_term.num_prototypes + self.max_long_tokens = cfg.long_term.max_num_tokens + self.buffer_tokens = cfg.long_term.buffer_tokens + else: + self.max_mem_frames = cfg.max_mem_frames - 1 + + def _readout(self, affinity, v, uncert_mask=None) -> torch.Tensor: + # affinity: bs*N*HW + # v: bs*C*N or bs*num_objects*C*N + # returns bs*C*HW or bs*num_objects*C*HW + if len(v.shape) == 3: + # single object + if uncert_mask is not None: + return v @ affinity * uncert_mask + else: + return v @ affinity + else: + bs, num_objects, C, N = v.shape + v = v.view(bs, num_objects * C, N) + out = v @ affinity + if uncert_mask is not None: + uncert_mask = uncert_mask.flatten(start_dim=2).expand(-1, C, -1) + out = out * uncert_mask + return out.view(bs, num_objects, C, -1) + + def _get_mask_by_ids(self, mask: torch.Tensor, obj_ids: List[int]) -> torch.Tensor: + # -1 because the mask does not contain the background channel + return mask[:, [self.object_manager.find_tmp_by_id(obj) - 1 for obj in obj_ids]] + + def _get_sensory_by_ids(self, obj_ids: List[int]) -> torch.Tensor: + return torch.stack([self.sensory[obj] for obj in obj_ids], dim=1) + + def _get_object_mem_by_ids(self, obj_ids: List[int]) -> torch.Tensor: + return torch.stack([self.obj_v[obj] for obj in obj_ids], dim=1) + + def _get_visual_values_by_ids(self, obj_ids: List[int]) -> torch.Tensor: + # All the values that the object ids refer to should have the same shape + value = torch.stack([self.work_mem.value[obj] for obj in obj_ids], dim=1) + if self.use_long_term and obj_ids[0] in self.long_mem.value: + lt_value = torch.stack([self.long_mem.value[obj] for obj in obj_ids], dim=1) + value = torch.cat([lt_value, value], dim=-1) + + return value + + def read_first_frame(self, last_msk_value, pix_feat: torch.Tensor, + last_mask: torch.Tensor, network: MatAnyone, uncert_output=None) -> Dict[int, torch.Tensor]: + """ + Read from all memory stores and returns a single memory readout tensor for each object + + pix_feat: (1/2) x C x H x W + query_key: (1/2) x C^k x H x W + selection: (1/2) x C^k x H x W + last_mask: (1/2) x num_objects x H x W (at stride 16) + return a dict of memory readouts, indexed by object indices. Each readout is C*H*W + """ + h, w = pix_feat.shape[-2:] + bs = pix_feat.shape[0] + assert last_mask.shape[0] == bs + + """ + Compute affinity and perform readout + """ + all_readout_mem = {} + buckets = self.work_mem.buckets + for bucket_id, bucket in buckets.items(): + + if self.chunk_size < 1: + object_chunks = [bucket] + else: + object_chunks = [ + bucket[i:i + self.chunk_size] for i in range(0, len(bucket), self.chunk_size) + ] + + for objects in object_chunks: + this_sensory = self._get_sensory_by_ids(objects) + this_last_mask = self._get_mask_by_ids(last_mask, objects) + this_msk_value = self._get_visual_values_by_ids(objects) # (1/2)*num_objects*C*N + pixel_readout = network.pixel_fusion(pix_feat, last_msk_value, this_sensory, + this_last_mask) + this_obj_mem = self._get_object_mem_by_ids(objects).unsqueeze(2) + readout_memory, aux_features = network.readout_query(pixel_readout, this_obj_mem) + for i, obj in enumerate(objects): + all_readout_mem[obj] = readout_memory[:, i] + + if self.save_aux: + aux_output = { + # 'sensory': this_sensory, + # 'pixel_readout': pixel_readout, + 'q_logits': aux_features['logits'] if aux_features else None, + # 'q_weights': aux_features['q_weights'] if aux_features else None, + # 'p_weights': aux_features['p_weights'] if aux_features else None, + # 'attn_mask': aux_features['attn_mask'].float() if aux_features else None, + } + self.aux = aux_output + + return all_readout_mem + + def read(self, pix_feat: torch.Tensor, query_key: torch.Tensor, selection: torch.Tensor, + last_mask: torch.Tensor, network: MatAnyone, uncert_output=None, last_msk_value=None, ti=None, + last_pix_feat=None, last_pred_mask=None) -> Dict[int, torch.Tensor]: + """ + Read from all memory stores and returns a single memory readout tensor for each object + + pix_feat: (1/2) x C x H x W + query_key: (1/2) x C^k x H x W + selection: (1/2) x C^k x H x W + last_mask: (1/2) x num_objects x H x W (at stride 16) + return a dict of memory readouts, indexed by object indices. Each readout is C*H*W + """ + h, w = pix_feat.shape[-2:] + bs = pix_feat.shape[0] + assert query_key.shape[0] == bs + assert selection.shape[0] == bs + assert last_mask.shape[0] == bs + + uncert_mask = uncert_output["mask"] if uncert_output is not None else None + + query_key = query_key.flatten(start_dim=2) # bs*C^k*HW + selection = selection.flatten(start_dim=2) # bs*C^k*HW + """ + Compute affinity and perform readout + """ + all_readout_mem = {} + buckets = self.work_mem.buckets + for bucket_id, bucket in buckets.items(): + if self.use_long_term and self.long_mem.engaged(bucket_id): + # Use long-term memory + long_mem_size = self.long_mem.size(bucket_id) + memory_key = torch.cat([self.long_mem.key[bucket_id], self.work_mem.key[bucket_id]], + -1) + shrinkage = torch.cat( + [self.long_mem.shrinkage[bucket_id], self.work_mem.shrinkage[bucket_id]], -1) + + similarity = get_similarity(memory_key, shrinkage, query_key, selection) + affinity, usage = do_softmax(similarity, + top_k=self.top_k, + inplace=True, + return_usage=True) + """ + Record memory usage for working and long-term memory + """ + # ignore the index return for long-term memory + work_usage = usage[:, long_mem_size:] + self.work_mem.update_bucket_usage(bucket_id, work_usage) + + if self.count_long_term_usage: + # ignore the index return for working memory + long_usage = usage[:, :long_mem_size] + self.long_mem.update_bucket_usage(bucket_id, long_usage) + else: + # no long-term memory + memory_key = self.work_mem.key[bucket_id] + shrinkage = self.work_mem.shrinkage[bucket_id] + similarity = get_similarity(memory_key, shrinkage, query_key, selection, uncert_mask=uncert_mask) + + if self.use_long_term: + affinity, usage = do_softmax(similarity, + top_k=self.top_k, + inplace=True, + return_usage=True) + self.work_mem.update_bucket_usage(bucket_id, usage) + else: + affinity = do_softmax(similarity, top_k=self.top_k, inplace=True) + + if self.chunk_size < 1: + object_chunks = [bucket] + else: + object_chunks = [ + bucket[i:i + self.chunk_size] for i in range(0, len(bucket), self.chunk_size) + ] + + for objects in object_chunks: + this_sensory = self._get_sensory_by_ids(objects) + this_last_mask = self._get_mask_by_ids(last_mask, objects) + this_msk_value = self._get_visual_values_by_ids(objects) # (1/2)*num_objects*C*N + visual_readout = self._readout(affinity, + this_msk_value, uncert_mask).view(bs, len(objects), self.CV, h, w) + + uncert_output = network.pred_uncertainty(last_pix_feat, pix_feat, last_pred_mask, visual_readout[:,0]-last_msk_value[:,0]) + + if uncert_output is not None: + uncert_prob = uncert_output["prob"].unsqueeze(1) # b n 1 h w + visual_readout = visual_readout*uncert_prob + last_msk_value*(1-uncert_prob) + + pixel_readout = network.pixel_fusion(pix_feat, visual_readout, this_sensory, + this_last_mask) + this_obj_mem = self._get_object_mem_by_ids(objects).unsqueeze(2) + readout_memory, aux_features = network.readout_query(pixel_readout, this_obj_mem) + for i, obj in enumerate(objects): + all_readout_mem[obj] = readout_memory[:, i] + + if self.save_aux: + aux_output = { + # 'sensory': this_sensory, + # 'pixel_readout': pixel_readout, + 'q_logits': aux_features['logits'] if aux_features else None, + # 'q_weights': aux_features['q_weights'] if aux_features else None, + # 'p_weights': aux_features['p_weights'] if aux_features else None, + # 'attn_mask': aux_features['attn_mask'].float() if aux_features else None, + } + self.aux = aux_output + + return all_readout_mem + + def add_memory(self, + key: torch.Tensor, + shrinkage: torch.Tensor, + msk_value: torch.Tensor, + obj_value: torch.Tensor, + objects: List[int], + selection: torch.Tensor = None, + *, + as_permanent: bool = False) -> None: + # key: (1/2)*C*H*W + # msk_value: (1/2)*num_objects*C*H*W + # obj_value: (1/2)*num_objects*Q*C + # objects contains a list of object ids corresponding to the objects in msk_value/obj_value + bs = key.shape[0] + assert shrinkage.shape[0] == bs + assert msk_value.shape[0] == bs + assert obj_value.shape[0] == bs + + self.engaged = True + if self.H is None or self.config_stale: + self.config_stale = False + self.H, self.W = msk_value.shape[-2:] + self.HW = self.H * self.W + # convert from num. frames to num. tokens + self.max_work_tokens = self.max_mem_frames * self.HW + if self.use_long_term: + self.min_work_tokens = self.min_mem_frames * self.HW + + # key: bs*C*N + # value: bs*num_objects*C*N + key = key.flatten(start_dim=2) + shrinkage = shrinkage.flatten(start_dim=2) + self.CK = key.shape[1] + + msk_value = msk_value.flatten(start_dim=3) + self.CV = msk_value.shape[2] + + if selection is not None: + # not used in non-long-term mode + selection = selection.flatten(start_dim=2) + + # insert object values into object memory + for obj_id, obj in enumerate(objects): + if obj in self.obj_v: + """streaming average + each self.obj_v[obj] is (1/2)*num_summaries*(embed_dim+1) + first embed_dim keeps track of the sum of embeddings + the last dim keeps the total count + averaging in done inside the object transformer + + incoming obj_value is (1/2)*num_objects*num_summaries*(embed_dim+1) + self.obj_v[obj] = torch.cat([self.obj_v[obj], obj_value[:, obj_id]], dim=0) + """ + last_acc = self.obj_v[obj][:, :, -1] + new_acc = last_acc + obj_value[:, obj_id, :, -1] + + self.obj_v[obj][:, :, :-1] = (self.obj_v[obj][:, :, :-1] + + obj_value[:, obj_id, :, :-1]) + self.obj_v[obj][:, :, -1] = new_acc + else: + self.obj_v[obj] = obj_value[:, obj_id] + + # convert mask value tensor into a dict for insertion + msk_values = {obj: msk_value[:, obj_id] for obj_id, obj in enumerate(objects)} + self.work_mem.add(key, + msk_values, + shrinkage, + selection=selection, + as_permanent=as_permanent) + + for bucket_id in self.work_mem.buckets.keys(): + # long-term memory cleanup + if self.use_long_term: + # Do memory compressed if needed + if self.work_mem.non_perm_size(bucket_id) >= self.max_work_tokens: + # Remove obsolete features if needed + if self.long_mem.non_perm_size(bucket_id) >= (self.max_long_tokens - + self.num_prototypes): + self.long_mem.remove_obsolete_features( + bucket_id, + self.max_long_tokens - self.num_prototypes - self.buffer_tokens) + + self.compress_features(bucket_id) + else: + # FIFO + self.work_mem.remove_old_memory(bucket_id, self.max_work_tokens) + + def purge_except(self, obj_keep_idx: List[int]) -> None: + # purge certain objects from the memory except the one listed + self.work_mem.purge_except(obj_keep_idx) + if self.use_long_term and self.long_mem.engaged(): + self.long_mem.purge_except(obj_keep_idx) + self.sensory = {k: v for k, v in self.sensory.items() if k in obj_keep_idx} + + if not self.work_mem.engaged(): + # everything is removed! + self.engaged = False + + def compress_features(self, bucket_id: int) -> None: + + # perform memory consolidation + prototype_key, prototype_value, prototype_shrinkage = self.consolidation( + *self.work_mem.get_all_sliced(bucket_id, 0, -self.min_work_tokens)) + + # remove consolidated working memory + self.work_mem.sieve_by_range(bucket_id, + 0, + -self.min_work_tokens, + min_size=self.min_work_tokens) + + # add to long-term memory + self.long_mem.add(prototype_key, + prototype_value, + prototype_shrinkage, + selection=None, + supposed_bucket_id=bucket_id) + + def consolidation(self, candidate_key: torch.Tensor, candidate_shrinkage: torch.Tensor, + candidate_selection: torch.Tensor, candidate_value: Dict[int, torch.Tensor], + usage: torch.Tensor) -> (torch.Tensor, Dict[int, torch.Tensor], torch.Tensor): + # find the indices with max usage + bs = candidate_key.shape[0] + assert bs in [1, 2] + + prototype_key = [] + prototype_selection = [] + for bi in range(bs): + _, max_usage_indices = torch.topk(usage[bi], k=self.num_prototypes, dim=-1, sorted=True) + prototype_indices = max_usage_indices.flatten() + prototype_key.append(candidate_key[bi, :, prototype_indices]) + prototype_selection.append(candidate_selection[bi, :, prototype_indices]) + prototype_key = torch.stack(prototype_key, dim=0) + prototype_selection = torch.stack(prototype_selection, dim=0) + """ + Potentiation step + """ + similarity = get_similarity(candidate_key, candidate_shrinkage, prototype_key, + prototype_selection) + affinity = do_softmax(similarity) + + # readout the values + prototype_value = {k: self._readout(affinity, v) for k, v in candidate_value.items()} + + # readout the shrinkage term + prototype_shrinkage = self._readout(affinity, candidate_shrinkage) + + return prototype_key, prototype_value, prototype_shrinkage + + def initialize_sensory_if_needed(self, sample_key: torch.Tensor, ids: List[int]): + for obj in ids: + if obj not in self.sensory: + # also initializes the sensory memory + bs, _, h, w = sample_key.shape + self.sensory[obj] = torch.zeros((bs, self.sensory_dim, h, w), + device=sample_key.device) + + def update_sensory(self, sensory: torch.Tensor, ids: List[int]): + # sensory: 1*num_objects*C*H*W + for obj_id, obj in enumerate(ids): + self.sensory[obj] = sensory[:, obj_id] + + def get_sensory(self, ids: List[int]): + # returns (1/2)*num_objects*C*H*W + return self._get_sensory_by_ids(ids) + + def clear_non_permanent_memory(self): + self.work_mem.clear_non_permanent_memory() + if self.use_long_term: + self.long_mem.clear_non_permanent_memory() + + def clear_sensory_memory(self): + self.sensory = {} + + def clear_work_mem(self): + self.work_mem = KeyValueMemoryStore(save_selection=self.use_long_term, + save_usage=self.use_long_term) + + def clear_obj_mem(self): + self.obj_v = {} diff --git a/preprocessing/matanyone/matanyone/inference/object_info.py b/preprocessing/matanyone/matanyone/inference/object_info.py new file mode 100644 index 0000000..b0e0bd4 --- /dev/null +++ b/preprocessing/matanyone/matanyone/inference/object_info.py @@ -0,0 +1,24 @@ +class ObjectInfo: + """ + Store meta information for an object + """ + def __init__(self, id: int): + self.id = id + self.poke_count = 0 # count number of detections missed + + def poke(self) -> None: + self.poke_count += 1 + + def unpoke(self) -> None: + self.poke_count = 0 + + def __hash__(self): + return hash(self.id) + + def __eq__(self, other): + if type(other) == int: + return self.id == other + return self.id == other.id + + def __repr__(self): + return f'(ID: {self.id})' diff --git a/preprocessing/matanyone/matanyone/inference/object_manager.py b/preprocessing/matanyone/matanyone/inference/object_manager.py new file mode 100644 index 0000000..34a93a2 --- /dev/null +++ b/preprocessing/matanyone/matanyone/inference/object_manager.py @@ -0,0 +1,149 @@ +from typing import Union, List, Dict + +import torch +from .object_info import ObjectInfo + + +class ObjectManager: + """ + Object IDs are immutable. The same ID always represent the same object. + Temporary IDs are the positions of each object in the tensor. It changes as objects get removed. + Temporary IDs start from 1. + """ + + def __init__(self): + self.obj_to_tmp_id: Dict[ObjectInfo, int] = {} + self.tmp_id_to_obj: Dict[int, ObjectInfo] = {} + self.obj_id_to_obj: Dict[int, ObjectInfo] = {} + + self.all_historical_object_ids: List[int] = [] + + def _recompute_obj_id_to_obj_mapping(self) -> None: + self.obj_id_to_obj = {obj.id: obj for obj in self.obj_to_tmp_id} + + def add_new_objects( + self, objects: Union[List[ObjectInfo], ObjectInfo, + List[int]]) -> (List[int], List[int]): + if not isinstance(objects, list): + objects = [objects] + + corresponding_tmp_ids = [] + corresponding_obj_ids = [] + for obj in objects: + if isinstance(obj, int): + obj = ObjectInfo(id=obj) + + if obj in self.obj_to_tmp_id: + # old object + corresponding_tmp_ids.append(self.obj_to_tmp_id[obj]) + corresponding_obj_ids.append(obj.id) + else: + # new object + new_obj = ObjectInfo(id=obj.id) + + # new object + new_tmp_id = len(self.obj_to_tmp_id) + 1 + self.obj_to_tmp_id[new_obj] = new_tmp_id + self.tmp_id_to_obj[new_tmp_id] = new_obj + self.all_historical_object_ids.append(new_obj.id) + corresponding_tmp_ids.append(new_tmp_id) + corresponding_obj_ids.append(new_obj.id) + + self._recompute_obj_id_to_obj_mapping() + assert corresponding_tmp_ids == sorted(corresponding_tmp_ids) + return corresponding_tmp_ids, corresponding_obj_ids + + def delete_objects(self, obj_ids_to_remove: Union[int, List[int]]) -> None: + # delete an object or a list of objects + # re-sort the tmp ids + if isinstance(obj_ids_to_remove, int): + obj_ids_to_remove = [obj_ids_to_remove] + + new_tmp_id = 1 + total_num_id = len(self.obj_to_tmp_id) + + local_obj_to_tmp_id = {} + local_tmp_to_obj_id = {} + + for tmp_iter in range(1, total_num_id + 1): + obj = self.tmp_id_to_obj[tmp_iter] + if obj.id not in obj_ids_to_remove: + local_obj_to_tmp_id[obj] = new_tmp_id + local_tmp_to_obj_id[new_tmp_id] = obj + new_tmp_id += 1 + + self.obj_to_tmp_id = local_obj_to_tmp_id + self.tmp_id_to_obj = local_tmp_to_obj_id + self._recompute_obj_id_to_obj_mapping() + + def purge_inactive_objects(self, + max_missed_detection_count: int) -> (bool, List[int], List[int]): + # remove tmp ids of objects that are removed + obj_id_to_be_deleted = [] + tmp_id_to_be_deleted = [] + tmp_id_to_keep = [] + obj_id_to_keep = [] + + for obj in self.obj_to_tmp_id: + if obj.poke_count > max_missed_detection_count: + obj_id_to_be_deleted.append(obj.id) + tmp_id_to_be_deleted.append(self.obj_to_tmp_id[obj]) + else: + tmp_id_to_keep.append(self.obj_to_tmp_id[obj]) + obj_id_to_keep.append(obj.id) + + purge_activated = len(obj_id_to_be_deleted) > 0 + if purge_activated: + self.delete_objects(obj_id_to_be_deleted) + return purge_activated, tmp_id_to_keep, obj_id_to_keep + + def tmp_to_obj_cls(self, mask) -> torch.Tensor: + # remap tmp id cls representation to the true object id representation + new_mask = torch.zeros_like(mask) + for tmp_id, obj in self.tmp_id_to_obj.items(): + new_mask[mask == tmp_id] = obj.id + return new_mask + + def get_tmp_to_obj_mapping(self) -> Dict[int, ObjectInfo]: + # returns the mapping in a dict format for saving it with pickle + return {obj.id: tmp_id for obj, tmp_id in self.tmp_id_to_obj.items()} + + def realize_dict(self, obj_dict, dim=1) -> torch.Tensor: + # turns a dict indexed by obj id into a tensor, ordered by tmp IDs + output = [] + for _, obj in self.tmp_id_to_obj.items(): + if obj.id not in obj_dict: + raise NotImplementedError + output.append(obj_dict[obj.id]) + output = torch.stack(output, dim=dim) + return output + + def make_one_hot(self, cls_mask) -> torch.Tensor: + output = [] + for _, obj in self.tmp_id_to_obj.items(): + output.append(cls_mask == obj.id) + if len(output) == 0: + output = torch.zeros((0, *cls_mask.shape), dtype=torch.bool, device=cls_mask.device) + else: + output = torch.stack(output, dim=0) + return output + + @property + def all_obj_ids(self) -> List[int]: + return [k.id for k in self.obj_to_tmp_id] + + @property + def num_obj(self) -> int: + return len(self.obj_to_tmp_id) + + def has_all(self, objects: List[int]) -> bool: + for obj in objects: + if obj not in self.obj_to_tmp_id: + return False + return True + + def find_object_by_id(self, obj_id) -> ObjectInfo: + return self.obj_id_to_obj[obj_id] + + def find_tmp_by_id(self, obj_id) -> int: + return self.obj_to_tmp_id[self.obj_id_to_obj[obj_id]] diff --git a/preprocessing/matanyone/matanyone/inference/utils/__init__.py b/preprocessing/matanyone/matanyone/inference/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/preprocessing/matanyone/matanyone/inference/utils/args_utils.py b/preprocessing/matanyone/matanyone/inference/utils/args_utils.py new file mode 100644 index 0000000..a771cca --- /dev/null +++ b/preprocessing/matanyone/matanyone/inference/utils/args_utils.py @@ -0,0 +1,30 @@ +import logging +from omegaconf import DictConfig + +log = logging.getLogger() + + +def get_dataset_cfg(cfg: DictConfig): + dataset_name = cfg.dataset + data_cfg = cfg.datasets[dataset_name] + + potential_overrides = [ + 'image_directory', + 'mask_directory', + 'json_directory', + 'size', + 'save_all', + 'use_all_masks', + 'use_long_term', + 'mem_every', + ] + + for override in potential_overrides: + if cfg[override] is not None: + log.info(f'Overriding config {override} from {data_cfg[override]} to {cfg[override]}') + data_cfg[override] = cfg[override] + # escalte all potential overrides to the top-level config + if override in data_cfg: + cfg[override] = data_cfg[override] + + return data_cfg diff --git a/preprocessing/matanyone/matanyone/model/__init__.py b/preprocessing/matanyone/matanyone/model/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/preprocessing/matanyone/matanyone/model/aux_modules.py b/preprocessing/matanyone/matanyone/model/aux_modules.py new file mode 100644 index 0000000..efeb515 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/aux_modules.py @@ -0,0 +1,93 @@ +""" +For computing auxiliary outputs for auxiliary losses +""" +from typing import Dict +from omegaconf import DictConfig +import torch +import torch.nn as nn + +from .group_modules import GConv2d +from ...utils.tensor_utils import aggregate + + +class LinearPredictor(nn.Module): + def __init__(self, x_dim: int, pix_dim: int): + super().__init__() + self.projection = GConv2d(x_dim, pix_dim + 1, kernel_size=1) + + def forward(self, pix_feat: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + # pixel_feat: B*pix_dim*H*W + # x: B*num_objects*x_dim*H*W + num_objects = x.shape[1] + x = self.projection(x) + + pix_feat = pix_feat.unsqueeze(1).expand(-1, num_objects, -1, -1, -1) + logits = (pix_feat * x[:, :, :-1]).sum(dim=2) + x[:, :, -1] + return logits + + +class DirectPredictor(nn.Module): + def __init__(self, x_dim: int): + super().__init__() + self.projection = GConv2d(x_dim, 1, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: B*num_objects*x_dim*H*W + logits = self.projection(x).squeeze(2) + return logits + + +class AuxComputer(nn.Module): + def __init__(self, cfg: DictConfig): + super().__init__() + + use_sensory_aux = cfg.model.aux_loss.sensory.enabled + self.use_query_aux = cfg.model.aux_loss.query.enabled + self.use_sensory_aux = use_sensory_aux + + sensory_dim = cfg.model.sensory_dim + embed_dim = cfg.model.embed_dim + + if use_sensory_aux: + self.sensory_aux = LinearPredictor(sensory_dim, embed_dim) + + def _aggregate_with_selector(self, logits: torch.Tensor, selector: torch.Tensor) -> torch.Tensor: + prob = torch.sigmoid(logits) + if selector is not None: + prob = prob * selector + logits = aggregate(prob, dim=1) + return logits + + def forward(self, pix_feat: torch.Tensor, aux_input: Dict[str, torch.Tensor], + selector: torch.Tensor, seg_pass=False) -> Dict[str, torch.Tensor]: + sensory = aux_input['sensory'] + q_logits = aux_input['q_logits'] + + aux_output = {} + aux_output['attn_mask'] = aux_input['attn_mask'] + + if self.use_sensory_aux: + # B*num_objects*H*W + logits = self.sensory_aux(pix_feat, sensory) + aux_output['sensory_logits'] = self._aggregate_with_selector(logits, selector) + if self.use_query_aux: + # B*num_objects*num_levels*H*W + aux_output['q_logits'] = self._aggregate_with_selector( + torch.stack(q_logits, dim=2), + selector.unsqueeze(2) if selector is not None else None) + + return aux_output + + def compute_mask(self, aux_input: Dict[str, torch.Tensor], + selector: torch.Tensor) -> Dict[str, torch.Tensor]: + # sensory = aux_input['sensory'] + q_logits = aux_input['q_logits'] + + aux_output = {} + + # B*num_objects*num_levels*H*W + aux_output['q_logits'] = self._aggregate_with_selector( + torch.stack(q_logits, dim=2), + selector.unsqueeze(2) if selector is not None else None) + + return aux_output \ No newline at end of file diff --git a/preprocessing/matanyone/matanyone/model/big_modules.py b/preprocessing/matanyone/matanyone/model/big_modules.py new file mode 100644 index 0000000..4d09f53 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/big_modules.py @@ -0,0 +1,365 @@ +""" +big_modules.py - This file stores higher-level network blocks. + +x - usually denotes features that are shared between objects. +g - usually denotes features that are not shared between objects + with an extra "num_objects" dimension (batch_size * num_objects * num_channels * H * W). + +The trailing number of a variable usually denotes the stride +""" + +from typing import Iterable +from omegaconf import DictConfig +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .group_modules import MainToGroupDistributor, GroupFeatureFusionBlock, GConv2d +from .utils import resnet +from .modules import SensoryDeepUpdater, SensoryUpdater_fullscale, DecoderFeatureProcessor, MaskUpsampleBlock + +class UncertPred(nn.Module): + def __init__(self, model_cfg: DictConfig): + super().__init__() + self.conv1x1_v2 = nn.Conv2d(model_cfg.pixel_dim*2 + 1 + model_cfg.value_dim, 64, kernel_size=1, stride=1, bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.relu = nn.ReLU(inplace=True) + self.conv3x3 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1) + self.bn2 = nn.BatchNorm2d(32) + self.conv3x3_out = nn.Conv2d(32, 1, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1) + + def forward(self, last_frame_feat: torch.Tensor, cur_frame_feat: torch.Tensor, last_mask: torch.Tensor, mem_val_diff:torch.Tensor): + last_mask = F.interpolate(last_mask, size=last_frame_feat.shape[-2:], mode='area') + x = torch.cat([last_frame_feat, cur_frame_feat, last_mask, mem_val_diff], dim=1) + x = self.conv1x1_v2(x) + x = self.bn1(x) + x = self.relu(x) + x = self.conv3x3(x) + x = self.bn2(x) + x = self.relu(x) + x = self.conv3x3_out(x) + return x + + # override the default train() to freeze BN statistics + def train(self, mode=True): + self.training = False + for module in self.children(): + module.train(False) + return self + +class PixelEncoder(nn.Module): + def __init__(self, model_cfg: DictConfig): + super().__init__() + + self.is_resnet = 'resnet' in model_cfg.pixel_encoder.type + # if model_cfg.pretrained_resnet is set in the model_cfg we get the value + # else default to True + is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True) + if self.is_resnet: + if model_cfg.pixel_encoder.type == 'resnet18': + network = resnet.resnet18(pretrained=is_pretrained_resnet) + elif model_cfg.pixel_encoder.type == 'resnet50': + network = resnet.resnet50(pretrained=is_pretrained_resnet) + else: + raise NotImplementedError + self.conv1 = network.conv1 + self.bn1 = network.bn1 + self.relu = network.relu + self.maxpool = network.maxpool + + self.res2 = network.layer1 + self.layer2 = network.layer2 + self.layer3 = network.layer3 + else: + raise NotImplementedError + + def forward(self, x: torch.Tensor, seq_length=None) -> (torch.Tensor, torch.Tensor, torch.Tensor): + f1 = x + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + f2 = x + x = self.maxpool(x) + f4 = self.res2(x) + f8 = self.layer2(f4) + f16 = self.layer3(f8) + + return f16, f8, f4, f2, f1 + + # override the default train() to freeze BN statistics + def train(self, mode=True): + self.training = False + for module in self.children(): + module.train(False) + return self + + +class KeyProjection(nn.Module): + def __init__(self, model_cfg: DictConfig): + super().__init__() + in_dim = model_cfg.pixel_encoder.ms_dims[0] + mid_dim = model_cfg.pixel_dim + key_dim = model_cfg.key_dim + + self.pix_feat_proj = nn.Conv2d(in_dim, mid_dim, kernel_size=1) + self.key_proj = nn.Conv2d(mid_dim, key_dim, kernel_size=3, padding=1) + # shrinkage + self.d_proj = nn.Conv2d(mid_dim, 1, kernel_size=3, padding=1) + # selection + self.e_proj = nn.Conv2d(mid_dim, key_dim, kernel_size=3, padding=1) + + nn.init.orthogonal_(self.key_proj.weight.data) + nn.init.zeros_(self.key_proj.bias.data) + + def forward(self, x: torch.Tensor, *, need_s: bool, + need_e: bool) -> (torch.Tensor, torch.Tensor, torch.Tensor): + x = self.pix_feat_proj(x) + shrinkage = self.d_proj(x)**2 + 1 if (need_s) else None + selection = torch.sigmoid(self.e_proj(x)) if (need_e) else None + + return self.key_proj(x), shrinkage, selection + + +class MaskEncoder(nn.Module): + def __init__(self, model_cfg: DictConfig, single_object=False): + super().__init__() + pixel_dim = model_cfg.pixel_dim + value_dim = model_cfg.value_dim + sensory_dim = model_cfg.sensory_dim + final_dim = model_cfg.mask_encoder.final_dim + + self.single_object = single_object + extra_dim = 1 if single_object else 2 + + # if model_cfg.pretrained_resnet is set in the model_cfg we get the value + # else default to True + is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True) + if model_cfg.mask_encoder.type == 'resnet18': + network = resnet.resnet18(pretrained=is_pretrained_resnet, extra_dim=extra_dim) + elif model_cfg.mask_encoder.type == 'resnet50': + network = resnet.resnet50(pretrained=is_pretrained_resnet, extra_dim=extra_dim) + else: + raise NotImplementedError + self.conv1 = network.conv1 + self.bn1 = network.bn1 + self.relu = network.relu + self.maxpool = network.maxpool + + self.layer1 = network.layer1 + self.layer2 = network.layer2 + self.layer3 = network.layer3 + + self.distributor = MainToGroupDistributor() + self.fuser = GroupFeatureFusionBlock(pixel_dim, final_dim, value_dim) + + self.sensory_update = SensoryDeepUpdater(value_dim, sensory_dim) + + def forward(self, + image: torch.Tensor, + pix_feat: torch.Tensor, + sensory: torch.Tensor, + masks: torch.Tensor, + others: torch.Tensor, + *, + deep_update: bool = True, + chunk_size: int = -1) -> (torch.Tensor, torch.Tensor): + # ms_features are from the key encoder + # we only use the first one (lowest resolution), following XMem + if self.single_object: + g = masks.unsqueeze(2) + else: + g = torch.stack([masks, others], dim=2) + + g = self.distributor(image, g) + + batch_size, num_objects = g.shape[:2] + if chunk_size < 1 or chunk_size >= num_objects: + chunk_size = num_objects + fast_path = True + new_sensory = sensory + else: + if deep_update: + new_sensory = torch.empty_like(sensory) + else: + new_sensory = sensory + fast_path = False + + # chunk-by-chunk inference + all_g = [] + for i in range(0, num_objects, chunk_size): + if fast_path: + g_chunk = g + else: + g_chunk = g[:, i:i + chunk_size] + actual_chunk_size = g_chunk.shape[1] + g_chunk = g_chunk.flatten(start_dim=0, end_dim=1) + + g_chunk = self.conv1(g_chunk) + g_chunk = self.bn1(g_chunk) # 1/2, 64 + g_chunk = self.maxpool(g_chunk) # 1/4, 64 + g_chunk = self.relu(g_chunk) + + g_chunk = self.layer1(g_chunk) # 1/4 + g_chunk = self.layer2(g_chunk) # 1/8 + g_chunk = self.layer3(g_chunk) # 1/16 + + g_chunk = g_chunk.view(batch_size, actual_chunk_size, *g_chunk.shape[1:]) + g_chunk = self.fuser(pix_feat, g_chunk) + all_g.append(g_chunk) + if deep_update: + if fast_path: + new_sensory = self.sensory_update(g_chunk, sensory) + else: + new_sensory[:, i:i + chunk_size] = self.sensory_update( + g_chunk, sensory[:, i:i + chunk_size]) + g = torch.cat(all_g, dim=1) + + return g, new_sensory + + # override the default train() to freeze BN statistics + def train(self, mode=True): + self.training = False + for module in self.children(): + module.train(False) + return self + + +class PixelFeatureFuser(nn.Module): + def __init__(self, model_cfg: DictConfig, single_object=False): + super().__init__() + value_dim = model_cfg.value_dim + sensory_dim = model_cfg.sensory_dim + pixel_dim = model_cfg.pixel_dim + embed_dim = model_cfg.embed_dim + self.single_object = single_object + + self.fuser = GroupFeatureFusionBlock(pixel_dim, value_dim, embed_dim) + if self.single_object: + self.sensory_compress = GConv2d(sensory_dim + 1, value_dim, kernel_size=1) + else: + self.sensory_compress = GConv2d(sensory_dim + 2, value_dim, kernel_size=1) + + def forward(self, + pix_feat: torch.Tensor, + pixel_memory: torch.Tensor, + sensory_memory: torch.Tensor, + last_mask: torch.Tensor, + last_others: torch.Tensor, + *, + chunk_size: int = -1) -> torch.Tensor: + batch_size, num_objects = pixel_memory.shape[:2] + + if self.single_object: + last_mask = last_mask.unsqueeze(2) + else: + last_mask = torch.stack([last_mask, last_others], dim=2) + + if chunk_size < 1: + chunk_size = num_objects + + # chunk-by-chunk inference + all_p16 = [] + for i in range(0, num_objects, chunk_size): + sensory_readout = self.sensory_compress( + torch.cat([sensory_memory[:, i:i + chunk_size], last_mask[:, i:i + chunk_size]], 2)) + p16 = pixel_memory[:, i:i + chunk_size] + sensory_readout + p16 = self.fuser(pix_feat, p16) + all_p16.append(p16) + p16 = torch.cat(all_p16, dim=1) + + return p16 + + +class MaskDecoder(nn.Module): + def __init__(self, model_cfg: DictConfig): + super().__init__() + embed_dim = model_cfg.embed_dim + sensory_dim = model_cfg.sensory_dim + ms_image_dims = model_cfg.pixel_encoder.ms_dims + up_dims = model_cfg.mask_decoder.up_dims + + assert embed_dim == up_dims[0] + + self.sensory_update = SensoryUpdater_fullscale([up_dims[0], up_dims[1], up_dims[2], up_dims[3], up_dims[4] + 1], sensory_dim, + sensory_dim) + + self.decoder_feat_proc = DecoderFeatureProcessor(ms_image_dims[1:], up_dims[:-1]) + self.up_16_8 = MaskUpsampleBlock(up_dims[0], up_dims[1]) + self.up_8_4 = MaskUpsampleBlock(up_dims[1], up_dims[2]) + # newly add for alpha matte + self.up_4_2 = MaskUpsampleBlock(up_dims[2], up_dims[3]) + self.up_2_1 = MaskUpsampleBlock(up_dims[3], up_dims[4]) + + self.pred_seg = nn.Conv2d(up_dims[-1], 1, kernel_size=3, padding=1) + self.pred_mat = nn.Conv2d(up_dims[-1], 1, kernel_size=3, padding=1) + + def forward(self, + ms_image_feat: Iterable[torch.Tensor], + memory_readout: torch.Tensor, + sensory: torch.Tensor, + *, + chunk_size: int = -1, + update_sensory: bool = True, + seg_pass: bool = False, + last_mask=None, + sigmoid_residual=False) -> (torch.Tensor, torch.Tensor): + + batch_size, num_objects = memory_readout.shape[:2] + f8, f4, f2, f1 = self.decoder_feat_proc(ms_image_feat[1:]) + if chunk_size < 1 or chunk_size >= num_objects: + chunk_size = num_objects + fast_path = True + new_sensory = sensory + else: + if update_sensory: + new_sensory = torch.empty_like(sensory) + else: + new_sensory = sensory + fast_path = False + + # chunk-by-chunk inference + all_logits = [] + for i in range(0, num_objects, chunk_size): + if fast_path: + p16 = memory_readout + else: + p16 = memory_readout[:, i:i + chunk_size] + actual_chunk_size = p16.shape[1] + + p8 = self.up_16_8(p16, f8) + p4 = self.up_8_4(p8, f4) + p2 = self.up_4_2(p4, f2) + p1 = self.up_2_1(p2, f1) + with torch.amp.autocast("cuda"): + if seg_pass: + if last_mask is not None: + res = self.pred_seg(F.relu(p1.flatten(start_dim=0, end_dim=1).float())) + if sigmoid_residual: + res = (torch.sigmoid(res) - 0.5) * 2 # regularization: (-1, 1) change on last mask + logits = last_mask + res + else: + logits = self.pred_seg(F.relu(p1.flatten(start_dim=0, end_dim=1).float())) + else: + if last_mask is not None: + res = self.pred_mat(F.relu(p1.flatten(start_dim=0, end_dim=1).float())) + if sigmoid_residual: + res = (torch.sigmoid(res) - 0.5) * 2 # regularization: (-1, 1) change on last mask + logits = last_mask + res + else: + logits = self.pred_mat(F.relu(p1.flatten(start_dim=0, end_dim=1).float())) + ## SensoryUpdater_fullscale + if update_sensory: + p1 = torch.cat( + [p1, logits.view(batch_size, actual_chunk_size, 1, *logits.shape[-2:])], 2) + if fast_path: + new_sensory = self.sensory_update([p16, p8, p4, p2, p1], sensory) + else: + new_sensory[:, + i:i + chunk_size] = self.sensory_update([p16, p8, p4, p2, p1], + sensory[:, + i:i + chunk_size]) + all_logits.append(logits) + logits = torch.cat(all_logits, dim=0) + logits = logits.view(batch_size, num_objects, *logits.shape[-2:]) + + return new_sensory, logits diff --git a/preprocessing/matanyone/matanyone/model/channel_attn.py b/preprocessing/matanyone/matanyone/model/channel_attn.py new file mode 100644 index 0000000..a2096c1 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/channel_attn.py @@ -0,0 +1,39 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class CAResBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, residual: bool = True): + super().__init__() + self.residual = residual + self.conv1 = nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1) + self.conv2 = nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1) + + t = int((abs(math.log2(out_dim)) + 1) // 2) + k = t if t % 2 else t + 1 + self.pool = nn.AdaptiveAvgPool2d(1) + self.conv = nn.Conv1d(1, 1, kernel_size=k, padding=(k - 1) // 2, bias=False) + + if self.residual: + if in_dim == out_dim: + self.downsample = nn.Identity() + else: + self.downsample = nn.Conv2d(in_dim, out_dim, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + r = x + x = self.conv1(F.relu(x)) + x = self.conv2(F.relu(x)) + + b, c = x.shape[:2] + w = self.pool(x).view(b, 1, c) + w = self.conv(w).transpose(-1, -2).unsqueeze(-1).sigmoid() # B*C*1*1 + + if self.residual: + x = x * w + self.downsample(r) + else: + x = x * w + + return x diff --git a/preprocessing/matanyone/matanyone/model/group_modules.py b/preprocessing/matanyone/matanyone/model/group_modules.py new file mode 100644 index 0000000..f143f46 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/group_modules.py @@ -0,0 +1,126 @@ +from typing import Optional +import torch +import torch.nn as nn +import torch.nn.functional as F +from .channel_attn import CAResBlock + +def interpolate_groups(g: torch.Tensor, ratio: float, mode: str, + align_corners: bool) -> torch.Tensor: + batch_size, num_objects = g.shape[:2] + g = F.interpolate(g.flatten(start_dim=0, end_dim=1), + scale_factor=ratio, + mode=mode, + align_corners=align_corners) + g = g.view(batch_size, num_objects, *g.shape[1:]) + return g + + +def upsample_groups(g: torch.Tensor, + ratio: float = 2, + mode: str = 'bilinear', + align_corners: bool = False) -> torch.Tensor: + return interpolate_groups(g, ratio, mode, align_corners) + + +def downsample_groups(g: torch.Tensor, + ratio: float = 1 / 2, + mode: str = 'area', + align_corners: bool = None) -> torch.Tensor: + return interpolate_groups(g, ratio, mode, align_corners) + + +class GConv2d(nn.Conv2d): + def forward(self, g: torch.Tensor) -> torch.Tensor: + batch_size, num_objects = g.shape[:2] + g = super().forward(g.flatten(start_dim=0, end_dim=1)) + return g.view(batch_size, num_objects, *g.shape[1:]) + + +class GroupResBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int): + super().__init__() + + if in_dim == out_dim: + self.downsample = nn.Identity() + else: + self.downsample = GConv2d(in_dim, out_dim, kernel_size=1) + + self.conv1 = GConv2d(in_dim, out_dim, kernel_size=3, padding=1) + self.conv2 = GConv2d(out_dim, out_dim, kernel_size=3, padding=1) + + def forward(self, g: torch.Tensor) -> torch.Tensor: + out_g = self.conv1(F.relu(g)) + out_g = self.conv2(F.relu(out_g)) + + g = self.downsample(g) + + return out_g + g + + +class MainToGroupDistributor(nn.Module): + def __init__(self, + x_transform: Optional[nn.Module] = None, + g_transform: Optional[nn.Module] = None, + method: str = 'cat', + reverse_order: bool = False): + super().__init__() + + self.x_transform = x_transform + self.g_transform = g_transform + self.method = method + self.reverse_order = reverse_order + + def forward(self, x: torch.Tensor, g: torch.Tensor, skip_expand: bool = False) -> torch.Tensor: + num_objects = g.shape[1] + + if self.x_transform is not None: + x = self.x_transform(x) + + if self.g_transform is not None: + g = self.g_transform(g) + + if not skip_expand: + x = x.unsqueeze(1).expand(-1, num_objects, -1, -1, -1) + if self.method == 'cat': + if self.reverse_order: + g = torch.cat([g, x], 2) + else: + g = torch.cat([x, g], 2) + elif self.method == 'add': + g = x + g + elif self.method == 'mulcat': + g = torch.cat([x * g, g], dim=2) + elif self.method == 'muladd': + g = x * g + g + else: + raise NotImplementedError + + return g + + +class GroupFeatureFusionBlock(nn.Module): + def __init__(self, x_in_dim: int, g_in_dim: int, out_dim: int): + super().__init__() + + x_transform = nn.Conv2d(x_in_dim, out_dim, kernel_size=1) + g_transform = GConv2d(g_in_dim, out_dim, kernel_size=1) + + self.distributor = MainToGroupDistributor(x_transform=x_transform, + g_transform=g_transform, + method='add') + self.block1 = CAResBlock(out_dim, out_dim) + self.block2 = CAResBlock(out_dim, out_dim) + + def forward(self, x: torch.Tensor, g: torch.Tensor) -> torch.Tensor: + batch_size, num_objects = g.shape[:2] + + g = self.distributor(x, g) + + g = g.flatten(start_dim=0, end_dim=1) + + g = self.block1(g) + g = self.block2(g) + + g = g.view(batch_size, num_objects, *g.shape[1:]) + + return g \ No newline at end of file diff --git a/preprocessing/matanyone/matanyone/model/matanyone.py b/preprocessing/matanyone/matanyone/model/matanyone.py new file mode 100644 index 0000000..ec32c83 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/matanyone.py @@ -0,0 +1,333 @@ +from typing import List, Dict, Iterable +import logging +from omegaconf import DictConfig +import torch +import torch.nn as nn +import torch.nn.functional as F +from omegaconf import OmegaConf +from huggingface_hub import PyTorchModelHubMixin + +from .big_modules import PixelEncoder, UncertPred, KeyProjection, MaskEncoder, PixelFeatureFuser, MaskDecoder +from .aux_modules import AuxComputer +from .utils.memory_utils import get_affinity, readout +from .transformer.object_transformer import QueryTransformer +from .transformer.object_summarizer import ObjectSummarizer +from ...utils.tensor_utils import aggregate + +log = logging.getLogger() +class MatAnyone(nn.Module, + PyTorchModelHubMixin, + library_name="matanyone", + repo_url="https://github.com/pq-yang/MatAnyone", + coders={ + DictConfig: ( + lambda x: OmegaConf.to_container(x), + lambda data: OmegaConf.create(data), + ) + }, + ): + + def __init__(self, cfg: DictConfig, *, single_object=False): + super().__init__() + self.cfg = cfg + model_cfg = cfg.model + self.ms_dims = model_cfg.pixel_encoder.ms_dims + self.key_dim = model_cfg.key_dim + self.value_dim = model_cfg.value_dim + self.sensory_dim = model_cfg.sensory_dim + self.pixel_dim = model_cfg.pixel_dim + self.embed_dim = model_cfg.embed_dim + self.single_object = single_object + + log.info(f'Single object: {self.single_object}') + + self.pixel_encoder = PixelEncoder(model_cfg) + self.pix_feat_proj = nn.Conv2d(self.ms_dims[0], self.pixel_dim, kernel_size=1) + self.key_proj = KeyProjection(model_cfg) + self.mask_encoder = MaskEncoder(model_cfg, single_object=single_object) + self.mask_decoder = MaskDecoder(model_cfg) + self.pixel_fuser = PixelFeatureFuser(model_cfg, single_object=single_object) + self.object_transformer = QueryTransformer(model_cfg) + self.object_summarizer = ObjectSummarizer(model_cfg) + self.aux_computer = AuxComputer(cfg) + self.temp_sparity = UncertPred(model_cfg) + + self.register_buffer("pixel_mean", torch.Tensor(model_cfg.pixel_mean).view(-1, 1, 1), False) + self.register_buffer("pixel_std", torch.Tensor(model_cfg.pixel_std).view(-1, 1, 1), False) + + def _get_others(self, masks: torch.Tensor) -> torch.Tensor: + # for each object, return the sum of masks of all other objects + if self.single_object: + return None + + num_objects = masks.shape[1] + if num_objects >= 1: + others = (masks.sum(dim=1, keepdim=True) - masks).clamp(0, 1) + else: + others = torch.zeros_like(masks) + return others + + def pred_uncertainty(self, last_pix_feat: torch.Tensor, cur_pix_feat: torch.Tensor, last_mask: torch.Tensor, mem_val_diff:torch.Tensor): + logits = self.temp_sparity(last_frame_feat=last_pix_feat, + cur_frame_feat=cur_pix_feat, + last_mask=last_mask, + mem_val_diff=mem_val_diff) + + prob = torch.sigmoid(logits) + mask = (prob > 0) + 0 + + uncert_output = {"logits": logits, + "prob": prob, + "mask": mask} + + return uncert_output + + def encode_image(self, image: torch.Tensor, seq_length=None, last_feats=None) -> (Iterable[torch.Tensor], torch.Tensor): # type: ignore + image = (image - self.pixel_mean) / self.pixel_std + ms_image_feat = self.pixel_encoder(image, seq_length) # f16, f8, f4, f2, f1 + return ms_image_feat, self.pix_feat_proj(ms_image_feat[0]) + + def encode_mask( + self, + image: torch.Tensor, + ms_features: List[torch.Tensor], + sensory: torch.Tensor, + masks: torch.Tensor, + *, + deep_update: bool = True, + chunk_size: int = -1, + need_weights: bool = False) -> (torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor): + image = (image - self.pixel_mean) / self.pixel_std + others = self._get_others(masks) + mask_value, new_sensory = self.mask_encoder(image, + ms_features, + sensory, + masks, + others, + deep_update=deep_update, + chunk_size=chunk_size) + object_summaries, object_logits = self.object_summarizer(masks, mask_value, need_weights) + return mask_value, new_sensory, object_summaries, object_logits + + def transform_key(self, + final_pix_feat: torch.Tensor, + *, + need_sk: bool = True, + need_ek: bool = True) -> (torch.Tensor, torch.Tensor, torch.Tensor): + key, shrinkage, selection = self.key_proj(final_pix_feat, need_s=need_sk, need_e=need_ek) + return key, shrinkage, selection + + # Used in training only. + # This step is replaced by MemoryManager in test time + def read_memory(self, query_key: torch.Tensor, query_selection: torch.Tensor, + memory_key: torch.Tensor, memory_shrinkage: torch.Tensor, + msk_value: torch.Tensor, obj_memory: torch.Tensor, pix_feat: torch.Tensor, + sensory: torch.Tensor, last_mask: torch.Tensor, + selector: torch.Tensor, uncert_output=None, seg_pass=False, + last_pix_feat=None, last_pred_mask=None) -> (torch.Tensor, Dict[str, torch.Tensor]): + """ + query_key : B * CK * H * W + query_selection : B * CK * H * W + memory_key : B * CK * T * H * W + memory_shrinkage: B * 1 * T * H * W + msk_value : B * num_objects * CV * T * H * W + obj_memory : B * num_objects * T * num_summaries * C + pixel_feature : B * C * H * W + """ + batch_size, num_objects = msk_value.shape[:2] + + uncert_mask = uncert_output["mask"] if uncert_output is not None else None + + # read using visual attention + with torch.cuda.amp.autocast(enabled=False): + affinity = get_affinity(memory_key.float(), memory_shrinkage.float(), query_key.float(), + query_selection.float(), uncert_mask=uncert_mask) + + msk_value = msk_value.flatten(start_dim=1, end_dim=2).float() + + # B * (num_objects*CV) * H * W + pixel_readout = readout(affinity, msk_value, uncert_mask) + pixel_readout = pixel_readout.view(batch_size, num_objects, self.value_dim, + *pixel_readout.shape[-2:]) + + uncert_output = self.pred_uncertainty(last_pix_feat, pix_feat, last_pred_mask, pixel_readout[:,0]-msk_value[:,:,-1]) + uncert_prob = uncert_output["prob"].unsqueeze(1) # b n 1 h w + pixel_readout = pixel_readout*uncert_prob + msk_value[:,:,-1].unsqueeze(1)*(1-uncert_prob) + + pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask) + + + # read from query transformer + mem_readout, aux_features = self.readout_query(pixel_readout, obj_memory, selector=selector, seg_pass=seg_pass) + + aux_output = { + 'sensory': sensory, + 'q_logits': aux_features['logits'] if aux_features else None, + 'attn_mask': aux_features['attn_mask'] if aux_features else None, + } + + return mem_readout, aux_output, uncert_output + + def read_first_frame_memory(self, pixel_readout, + obj_memory: torch.Tensor, pix_feat: torch.Tensor, + sensory: torch.Tensor, last_mask: torch.Tensor, + selector: torch.Tensor, seg_pass=False) -> (torch.Tensor, Dict[str, torch.Tensor]): + """ + query_key : B * CK * H * W + query_selection : B * CK * H * W + memory_key : B * CK * T * H * W + memory_shrinkage: B * 1 * T * H * W + msk_value : B * num_objects * CV * T * H * W + obj_memory : B * num_objects * T * num_summaries * C + pixel_feature : B * C * H * W + """ + + pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask) + + # read from query transformer + mem_readout, aux_features = self.readout_query(pixel_readout, obj_memory, selector=selector, seg_pass=seg_pass) + + aux_output = { + 'sensory': sensory, + 'q_logits': aux_features['logits'] if aux_features else None, + 'attn_mask': aux_features['attn_mask'] if aux_features else None, + } + + return mem_readout, aux_output + + def pixel_fusion(self, + pix_feat: torch.Tensor, + pixel: torch.Tensor, + sensory: torch.Tensor, + last_mask: torch.Tensor, + *, + chunk_size: int = -1) -> torch.Tensor: + last_mask = F.interpolate(last_mask, size=sensory.shape[-2:], mode='area') + last_others = self._get_others(last_mask) + fused = self.pixel_fuser(pix_feat, + pixel, + sensory, + last_mask, + last_others, + chunk_size=chunk_size) + return fused + + def readout_query(self, + pixel_readout, + obj_memory, + *, + selector=None, + need_weights=False, + seg_pass=False) -> (torch.Tensor, Dict[str, torch.Tensor]): + return self.object_transformer(pixel_readout, + obj_memory, + selector=selector, + need_weights=need_weights, + seg_pass=seg_pass) + + def segment(self, + ms_image_feat: List[torch.Tensor], + memory_readout: torch.Tensor, + sensory: torch.Tensor, + *, + selector: bool = None, + chunk_size: int = -1, + update_sensory: bool = True, + seg_pass: bool = False, + clamp_mat: bool = True, + last_mask=None, + sigmoid_residual=False, + seg_mat=False) -> (torch.Tensor, torch.Tensor, torch.Tensor): + """ + multi_scale_features is from the key encoder for skip-connection + memory_readout is from working/long-term memory + sensory is the sensory memory + last_mask is the mask from the last frame, supplementing sensory memory + selector is 1 if an object exists, and 0 otherwise. We use it to filter padded objects + during training. + """ + #### use mat head for seg data + if seg_mat: + assert seg_pass + seg_pass = False + #### + sensory, logits = self.mask_decoder(ms_image_feat, + memory_readout, + sensory, + chunk_size=chunk_size, + update_sensory=update_sensory, + seg_pass = seg_pass, + last_mask=last_mask, + sigmoid_residual=sigmoid_residual) + if seg_pass: + prob = torch.sigmoid(logits) + if selector is not None: + prob = prob * selector + + # Softmax over all objects[] + logits = aggregate(prob, dim=1) + prob = F.softmax(logits, dim=1) + else: + if clamp_mat: + logits = logits.clamp(0.0, 1.0) + logits = torch.cat([torch.prod(1 - logits, dim=1, keepdim=True), logits], 1) + prob = logits + + return sensory, logits, prob + + def compute_aux(self, pix_feat: torch.Tensor, aux_inputs: Dict[str, torch.Tensor], + selector: torch.Tensor, seg_pass=False) -> Dict[str, torch.Tensor]: + return self.aux_computer(pix_feat, aux_inputs, selector, seg_pass=seg_pass) + + def forward(self, *args, **kwargs): + raise NotImplementedError + + def load_weights(self, src_dict, init_as_zero_if_needed=False) -> None: + if not self.single_object: + # Map single-object weight to multi-object weight (4->5 out channels in conv1) + for k in list(src_dict.keys()): + if k == 'mask_encoder.conv1.weight': + if src_dict[k].shape[1] == 4: + log.info(f'Converting {k} from single object to multiple objects.') + pads = torch.zeros((64, 1, 7, 7), device=src_dict[k].device) + if not init_as_zero_if_needed: + nn.init.orthogonal_(pads) + log.info(f'Randomly initialized padding for {k}.') + else: + log.info(f'Zero-initialized padding for {k}.') + src_dict[k] = torch.cat([src_dict[k], pads], 1) + elif k == 'pixel_fuser.sensory_compress.weight': + if src_dict[k].shape[1] == self.sensory_dim + 1: + log.info(f'Converting {k} from single object to multiple objects.') + pads = torch.zeros((self.value_dim, 1, 1, 1), device=src_dict[k].device) + if not init_as_zero_if_needed: + nn.init.orthogonal_(pads) + log.info(f'Randomly initialized padding for {k}.') + else: + log.info(f'Zero-initialized padding for {k}.') + src_dict[k] = torch.cat([src_dict[k], pads], 1) + elif self.single_object: + """ + If the model is multiple-object and we are training in single-object, + we strip the last channel of conv1. + This is not supposed to happen in standard training except when users are trying to + finetune a trained model with single object datasets. + """ + if src_dict['mask_encoder.conv1.weight'].shape[1] == 5: + log.warning('Converting mask_encoder.conv1.weight from multiple objects to single object.' + 'This is not supposed to happen in standard training.') + src_dict['mask_encoder.conv1.weight'] = src_dict['mask_encoder.conv1.weight'][:, :-1] + src_dict['pixel_fuser.sensory_compress.weight'] = src_dict['pixel_fuser.sensory_compress.weight'][:, :-1] + + for k in src_dict: + if k not in self.state_dict(): + log.info(f'Key {k} found in src_dict but not in self.state_dict()!!!') + for k in self.state_dict(): + if k not in src_dict: + log.info(f'Key {k} found in self.state_dict() but not in src_dict!!!') + + self.load_state_dict(src_dict, strict=False) + + @property + def device(self) -> torch.device: + return self.pixel_mean.device diff --git a/preprocessing/matanyone/matanyone/model/modules.py b/preprocessing/matanyone/matanyone/model/modules.py new file mode 100644 index 0000000..7350425 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/modules.py @@ -0,0 +1,149 @@ +from typing import List, Iterable +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .group_modules import MainToGroupDistributor, GroupResBlock, upsample_groups, GConv2d, downsample_groups + + +class UpsampleBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, scale_factor: int = 2): + super().__init__() + self.out_conv = ResBlock(in_dim, out_dim) + self.scale_factor = scale_factor + + def forward(self, in_g: torch.Tensor, skip_f: torch.Tensor) -> torch.Tensor: + g = F.interpolate(in_g, + scale_factor=self.scale_factor, + mode='bilinear') + g = self.out_conv(g) + g = g + skip_f + return g + +class MaskUpsampleBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, scale_factor: int = 2): + super().__init__() + self.distributor = MainToGroupDistributor(method='add') + self.out_conv = GroupResBlock(in_dim, out_dim) + self.scale_factor = scale_factor + + def forward(self, in_g: torch.Tensor, skip_f: torch.Tensor) -> torch.Tensor: + g = upsample_groups(in_g, ratio=self.scale_factor) + g = self.distributor(skip_f, g) + g = self.out_conv(g) + return g + + +class DecoderFeatureProcessor(nn.Module): + def __init__(self, decoder_dims: List[int], out_dims: List[int]): + super().__init__() + self.transforms = nn.ModuleList([ + nn.Conv2d(d_dim, p_dim, kernel_size=1) for d_dim, p_dim in zip(decoder_dims, out_dims) + ]) + + def forward(self, multi_scale_features: Iterable[torch.Tensor]) -> List[torch.Tensor]: + outputs = [func(x) for x, func in zip(multi_scale_features, self.transforms)] + return outputs + + +# @torch.jit.script +def _recurrent_update(h: torch.Tensor, values: torch.Tensor) -> torch.Tensor: + # h: batch_size * num_objects * hidden_dim * h * w + # values: batch_size * num_objects * (hidden_dim*3) * h * w + dim = values.shape[2] // 3 + forget_gate = torch.sigmoid(values[:, :, :dim]) + update_gate = torch.sigmoid(values[:, :, dim:dim * 2]) + new_value = torch.tanh(values[:, :, dim * 2:]) + new_h = forget_gate * h * (1 - update_gate) + update_gate * new_value + return new_h + + +class SensoryUpdater_fullscale(nn.Module): + # Used in the decoder, multi-scale feature + GRU + def __init__(self, g_dims: List[int], mid_dim: int, sensory_dim: int): + super().__init__() + self.g16_conv = GConv2d(g_dims[0], mid_dim, kernel_size=1) + self.g8_conv = GConv2d(g_dims[1], mid_dim, kernel_size=1) + self.g4_conv = GConv2d(g_dims[2], mid_dim, kernel_size=1) + self.g2_conv = GConv2d(g_dims[3], mid_dim, kernel_size=1) + self.g1_conv = GConv2d(g_dims[4], mid_dim, kernel_size=1) + + self.transform = GConv2d(mid_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1) + + nn.init.xavier_normal_(self.transform.weight) + + def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor: + g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \ + self.g4_conv(downsample_groups(g[2], ratio=1/4)) + \ + self.g2_conv(downsample_groups(g[3], ratio=1/8)) + \ + self.g1_conv(downsample_groups(g[4], ratio=1/16)) + + with torch.amp.autocast("cuda"): + g = g.float() + h = h.float() + values = self.transform(torch.cat([g, h], dim=2)) + new_h = _recurrent_update(h, values) + + return new_h + +class SensoryUpdater(nn.Module): + # Used in the decoder, multi-scale feature + GRU + def __init__(self, g_dims: List[int], mid_dim: int, sensory_dim: int): + super().__init__() + self.g16_conv = GConv2d(g_dims[0], mid_dim, kernel_size=1) + self.g8_conv = GConv2d(g_dims[1], mid_dim, kernel_size=1) + self.g4_conv = GConv2d(g_dims[2], mid_dim, kernel_size=1) + + self.transform = GConv2d(mid_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1) + + nn.init.xavier_normal_(self.transform.weight) + + def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor: + g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \ + self.g4_conv(downsample_groups(g[2], ratio=1/4)) + + with torch.amp.autocast("cuda"): + g = g.float() + h = h.float() + values = self.transform(torch.cat([g, h], dim=2)) + new_h = _recurrent_update(h, values) + + return new_h + + +class SensoryDeepUpdater(nn.Module): + def __init__(self, f_dim: int, sensory_dim: int): + super().__init__() + self.transform = GConv2d(f_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1) + + nn.init.xavier_normal_(self.transform.weight) + + def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor: + with torch.amp.autocast("cuda"): + g = g.float() + h = h.float() + values = self.transform(torch.cat([g, h], dim=2)) + new_h = _recurrent_update(h, values) + + return new_h + + +class ResBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int): + super().__init__() + + if in_dim == out_dim: + self.downsample = nn.Identity() + else: + self.downsample = nn.Conv2d(in_dim, out_dim, kernel_size=1) + + self.conv1 = nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1) + self.conv2 = nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1) + + def forward(self, g: torch.Tensor) -> torch.Tensor: + out_g = self.conv1(F.relu(g)) + out_g = self.conv2(F.relu(out_g)) + + g = self.downsample(g) + + return out_g + g \ No newline at end of file diff --git a/preprocessing/matanyone/matanyone/model/transformer/__init__.py b/preprocessing/matanyone/matanyone/model/transformer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/preprocessing/matanyone/matanyone/model/transformer/object_summarizer.py b/preprocessing/matanyone/matanyone/model/transformer/object_summarizer.py new file mode 100644 index 0000000..a2cf75a --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/transformer/object_summarizer.py @@ -0,0 +1,89 @@ +from typing import Optional +from omegaconf import DictConfig + +import torch +import torch.nn as nn +import torch.nn.functional as F +from .positional_encoding import PositionalEncoding + + +# @torch.jit.script +def _weighted_pooling(masks: torch.Tensor, value: torch.Tensor, + logits: torch.Tensor) -> (torch.Tensor, torch.Tensor): + # value: B*num_objects*H*W*value_dim + # logits: B*num_objects*H*W*num_summaries + # masks: B*num_objects*H*W*num_summaries: 1 if allowed + weights = logits.sigmoid() * masks + # B*num_objects*num_summaries*value_dim + sums = torch.einsum('bkhwq,bkhwc->bkqc', weights, value) + # B*num_objects*H*W*num_summaries -> B*num_objects*num_summaries*1 + area = weights.flatten(start_dim=2, end_dim=3).sum(2).unsqueeze(-1) + + # B*num_objects*num_summaries*value_dim + return sums, area + + +class ObjectSummarizer(nn.Module): + def __init__(self, model_cfg: DictConfig): + super().__init__() + + this_cfg = model_cfg.object_summarizer + self.value_dim = model_cfg.value_dim + self.embed_dim = this_cfg.embed_dim + self.num_summaries = this_cfg.num_summaries + self.add_pe = this_cfg.add_pe + self.pixel_pe_scale = model_cfg.pixel_pe_scale + self.pixel_pe_temperature = model_cfg.pixel_pe_temperature + + if self.add_pe: + self.pos_enc = PositionalEncoding(self.embed_dim, + scale=self.pixel_pe_scale, + temperature=self.pixel_pe_temperature) + + self.input_proj = nn.Linear(self.value_dim, self.embed_dim) + self.feature_pred = nn.Sequential( + nn.Linear(self.embed_dim, self.embed_dim), + nn.ReLU(inplace=True), + nn.Linear(self.embed_dim, self.embed_dim), + ) + self.weights_pred = nn.Sequential( + nn.Linear(self.embed_dim, self.embed_dim), + nn.ReLU(inplace=True), + nn.Linear(self.embed_dim, self.num_summaries), + ) + + def forward(self, + masks: torch.Tensor, + value: torch.Tensor, + need_weights: bool = False) -> (torch.Tensor, Optional[torch.Tensor]): + # masks: B*num_objects*(H0)*(W0) + # value: B*num_objects*value_dim*H*W + # -> B*num_objects*H*W*value_dim + h, w = value.shape[-2:] + masks = F.interpolate(masks, size=(h, w), mode='area') + masks = masks.unsqueeze(-1) + inv_masks = 1 - masks + repeated_masks = torch.cat([ + masks.expand(-1, -1, -1, -1, self.num_summaries // 2), + inv_masks.expand(-1, -1, -1, -1, self.num_summaries // 2), + ], + dim=-1) + + value = value.permute(0, 1, 3, 4, 2) + value = self.input_proj(value) + if self.add_pe: + pe = self.pos_enc(value) + value = value + pe + + with torch.amp.autocast("cuda"): + value = value.float() + feature = self.feature_pred(value) + logits = self.weights_pred(value) + sums, area = _weighted_pooling(repeated_masks, feature, logits) + + summaries = torch.cat([sums, area], dim=-1) + + if need_weights: + return summaries, logits + else: + return summaries, None \ No newline at end of file diff --git a/preprocessing/matanyone/matanyone/model/transformer/object_transformer.py b/preprocessing/matanyone/matanyone/model/transformer/object_transformer.py new file mode 100644 index 0000000..1aa6664 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/transformer/object_transformer.py @@ -0,0 +1,206 @@ +from typing import Dict, Optional +from omegaconf import DictConfig + +import torch +import torch.nn as nn +from ..group_modules import GConv2d +from ....utils.tensor_utils import aggregate +from .positional_encoding import PositionalEncoding +from .transformer_layers import CrossAttention, SelfAttention, FFN, PixelFFN + + +class QueryTransformerBlock(nn.Module): + def __init__(self, model_cfg: DictConfig): + super().__init__() + + this_cfg = model_cfg.object_transformer + self.embed_dim = this_cfg.embed_dim + self.num_heads = this_cfg.num_heads + self.num_queries = this_cfg.num_queries + self.ff_dim = this_cfg.ff_dim + + self.read_from_pixel = CrossAttention(self.embed_dim, + self.num_heads, + add_pe_to_qkv=this_cfg.read_from_pixel.add_pe_to_qkv) + self.self_attn = SelfAttention(self.embed_dim, + self.num_heads, + add_pe_to_qkv=this_cfg.query_self_attention.add_pe_to_qkv) + self.ffn = FFN(self.embed_dim, self.ff_dim) + self.read_from_query = CrossAttention(self.embed_dim, + self.num_heads, + add_pe_to_qkv=this_cfg.read_from_query.add_pe_to_qkv, + norm=this_cfg.read_from_query.output_norm) + self.pixel_ffn = PixelFFN(self.embed_dim) + + def forward( + self, + x: torch.Tensor, + pixel: torch.Tensor, + query_pe: torch.Tensor, + pixel_pe: torch.Tensor, + attn_mask: torch.Tensor, + need_weights: bool = False) -> (torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor): + # x: (bs*num_objects)*num_queries*embed_dim + # pixel: bs*num_objects*C*H*W + # query_pe: (bs*num_objects)*num_queries*embed_dim + # pixel_pe: (bs*num_objects)*(H*W)*C + # attn_mask: (bs*num_objects*num_heads)*num_queries*(H*W) + + # bs*num_objects*C*H*W -> (bs*num_objects)*(H*W)*C + pixel_flat = pixel.flatten(3, 4).flatten(0, 1).transpose(1, 2).contiguous() + x, q_weights = self.read_from_pixel(x, + pixel_flat, + query_pe, + pixel_pe, + attn_mask=attn_mask, + need_weights=need_weights) + x = self.self_attn(x, query_pe) + x = self.ffn(x) + + pixel_flat, p_weights = self.read_from_query(pixel_flat, + x, + pixel_pe, + query_pe, + need_weights=need_weights) + pixel = self.pixel_ffn(pixel, pixel_flat) + + if need_weights: + bs, num_objects, _, h, w = pixel.shape + q_weights = q_weights.view(bs, num_objects, self.num_heads, self.num_queries, h, w) + p_weights = p_weights.transpose(2, 3).view(bs, num_objects, self.num_heads, + self.num_queries, h, w) + + return x, pixel, q_weights, p_weights + + +class QueryTransformer(nn.Module): + def __init__(self, model_cfg: DictConfig): + super().__init__() + + this_cfg = model_cfg.object_transformer + self.value_dim = model_cfg.value_dim + self.embed_dim = this_cfg.embed_dim + self.num_heads = this_cfg.num_heads + self.num_queries = this_cfg.num_queries + + # query initialization and embedding + self.query_init = nn.Embedding(self.num_queries, self.embed_dim) + self.query_emb = nn.Embedding(self.num_queries, self.embed_dim) + + # projection from object summaries to query initialization and embedding + self.summary_to_query_init = nn.Linear(self.embed_dim, self.embed_dim) + self.summary_to_query_emb = nn.Linear(self.embed_dim, self.embed_dim) + + self.pixel_pe_scale = model_cfg.pixel_pe_scale + self.pixel_pe_temperature = model_cfg.pixel_pe_temperature + self.pixel_init_proj = GConv2d(self.embed_dim, self.embed_dim, kernel_size=1) + self.pixel_emb_proj = GConv2d(self.embed_dim, self.embed_dim, kernel_size=1) + self.spatial_pe = PositionalEncoding(self.embed_dim, + scale=self.pixel_pe_scale, + temperature=self.pixel_pe_temperature, + channel_last=False, + transpose_output=True) + + # transformer blocks + self.num_blocks = this_cfg.num_blocks + self.blocks = nn.ModuleList( + QueryTransformerBlock(model_cfg) for _ in range(self.num_blocks)) + self.mask_pred = nn.ModuleList( + nn.Sequential(nn.ReLU(), GConv2d(self.embed_dim, 1, kernel_size=1)) + for _ in range(self.num_blocks + 1)) + + self.act = nn.ReLU(inplace=True) + + def forward(self, + pixel: torch.Tensor, + obj_summaries: torch.Tensor, + selector: Optional[torch.Tensor] = None, + need_weights: bool = False, + seg_pass=False) -> (torch.Tensor, Dict[str, torch.Tensor]): + # pixel: B*num_objects*embed_dim*H*W + # obj_summaries: B*num_objects*T*num_queries*embed_dim + T = obj_summaries.shape[2] + bs, num_objects, _, H, W = pixel.shape + + # normalize object values + # the last channel is the cumulative area of the object + obj_summaries = obj_summaries.view(bs * num_objects, T, self.num_queries, + self.embed_dim + 1) + # sum over time + # during inference, T=1 as we already did streaming average in memory_manager + obj_sums = obj_summaries[:, :, :, :-1].sum(dim=1) + obj_area = obj_summaries[:, :, :, -1:].sum(dim=1) + obj_values = obj_sums / (obj_area + 1e-4) + obj_init = self.summary_to_query_init(obj_values) + obj_emb = self.summary_to_query_emb(obj_values) + + # positional embeddings for object queries + query = self.query_init.weight.unsqueeze(0).expand(bs * num_objects, -1, -1) + obj_init + query_emb = self.query_emb.weight.unsqueeze(0).expand(bs * num_objects, -1, -1) + obj_emb + + # positional embeddings for pixel features + pixel_init = self.pixel_init_proj(pixel) + pixel_emb = self.pixel_emb_proj(pixel) + pixel_pe = self.spatial_pe(pixel.flatten(0, 1)) + pixel_emb = pixel_emb.flatten(3, 4).flatten(0, 1).transpose(1, 2).contiguous() + pixel_pe = pixel_pe.flatten(1, 2) + pixel_emb + + pixel = pixel_init + + # run the transformer + aux_features = {'logits': []} + + # first aux output + aux_logits = self.mask_pred[0](pixel).squeeze(2) + attn_mask = self._get_aux_mask(aux_logits, selector, seg_pass=seg_pass) + aux_features['logits'].append(aux_logits) + for i in range(self.num_blocks): + query, pixel, q_weights, p_weights = self.blocks[i](query, + pixel, + query_emb, + pixel_pe, + attn_mask, + need_weights=need_weights) + + if self.training or i <= self.num_blocks - 1 or need_weights: + aux_logits = self.mask_pred[i + 1](pixel).squeeze(2) + attn_mask = self._get_aux_mask(aux_logits, selector, seg_pass=seg_pass) + aux_features['logits'].append(aux_logits) + + aux_features['q_weights'] = q_weights # last layer only + aux_features['p_weights'] = p_weights # last layer only + + if self.training: + # no need to save all heads + aux_features['attn_mask'] = attn_mask.view(bs, num_objects, self.num_heads, + self.num_queries, H, W)[:, :, 0] + + return pixel, aux_features + + def _get_aux_mask(self, logits: torch.Tensor, selector: torch.Tensor, seg_pass=False) -> torch.Tensor: + # logits: batch_size*num_objects*H*W + # selector: batch_size*num_objects*1*1 + # returns a mask of shape (batch_size*num_objects*num_heads)*num_queries*(H*W) + # where True means the attention is blocked + + if selector is None: + prob = logits.sigmoid() + else: + prob = logits.sigmoid() * selector + logits = aggregate(prob, dim=1) + + is_foreground = (logits[:, 1:] >= logits.max(dim=1, keepdim=True)[0]) + foreground_mask = is_foreground.bool().flatten(start_dim=2) + inv_foreground_mask = ~foreground_mask + inv_background_mask = foreground_mask + + aux_foreground_mask = inv_foreground_mask.unsqueeze(2).unsqueeze(2).repeat( + 1, 1, self.num_heads, self.num_queries // 2, 1).flatten(start_dim=0, end_dim=2) + aux_background_mask = inv_background_mask.unsqueeze(2).unsqueeze(2).repeat( + 1, 1, self.num_heads, self.num_queries // 2, 1).flatten(start_dim=0, end_dim=2) + + aux_mask = torch.cat([aux_foreground_mask, aux_background_mask], dim=1) + + aux_mask[torch.where(aux_mask.sum(-1) == aux_mask.shape[-1])] = False + + return aux_mask \ No newline at end of file diff --git a/preprocessing/matanyone/matanyone/model/transformer/positional_encoding.py b/preprocessing/matanyone/matanyone/model/transformer/positional_encoding.py new file mode 100644 index 0000000..6c15bb7 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/transformer/positional_encoding.py @@ -0,0 +1,108 @@ +# Reference: +# https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/transformer_decoder/position_encoding.py +# https://github.com/tatp22/multidim-positional-encoding/blob/master/positional_encodings/torch_encodings.py + +import math + +import numpy as np +import torch +from torch import nn + + +def get_emb(sin_inp: torch.Tensor) -> torch.Tensor: + """ + Gets a base embedding for one dimension with sin and cos intertwined + """ + emb = torch.stack((sin_inp.sin(), sin_inp.cos()), dim=-1) + return torch.flatten(emb, -2, -1) + + +class PositionalEncoding(nn.Module): + def __init__(self, + dim: int, + scale: float = math.pi * 2, + temperature: float = 10000, + normalize: bool = True, + channel_last: bool = True, + transpose_output: bool = False): + super().__init__() + dim = int(np.ceil(dim / 4) * 2) + self.dim = dim + inv_freq = 1.0 / (temperature**(torch.arange(0, dim, 2).float() / dim)) + self.register_buffer("inv_freq", inv_freq) + self.normalize = normalize + self.scale = scale + self.eps = 1e-6 + self.channel_last = channel_last + self.transpose_output = transpose_output + + self.cached_penc = None # the cache is irrespective of the number of objects + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + """ + :param tensor: A 4/5d tensor of size + channel_last=True: (batch_size, h, w, c) or (batch_size, k, h, w, c) + channel_last=False: (batch_size, c, h, w) or (batch_size, k, c, h, w) + :return: positional encoding tensor that has the same shape as the input if the input is 4d + if the input is 5d, the output is broadcastable along the k-dimension + """ + if len(tensor.shape) != 4 and len(tensor.shape) != 5: + raise RuntimeError(f'The input tensor has to be 4/5d, got {tensor.shape}!') + + if len(tensor.shape) == 5: + # take a sample from the k dimension + num_objects = tensor.shape[1] + tensor = tensor[:, 0] + else: + num_objects = None + + if self.channel_last: + batch_size, h, w, c = tensor.shape + else: + batch_size, c, h, w = tensor.shape + + if self.cached_penc is not None and self.cached_penc.shape == tensor.shape: + if num_objects is None: + return self.cached_penc + else: + return self.cached_penc.unsqueeze(1) + + self.cached_penc = None + + pos_y = torch.arange(h, device=tensor.device, dtype=self.inv_freq.dtype) + pos_x = torch.arange(w, device=tensor.device, dtype=self.inv_freq.dtype) + if self.normalize: + pos_y = pos_y / (pos_y[-1] + self.eps) * self.scale + pos_x = pos_x / (pos_x[-1] + self.eps) * self.scale + + sin_inp_y = torch.einsum("i,j->ij", pos_y, self.inv_freq) + sin_inp_x = torch.einsum("i,j->ij", pos_x, self.inv_freq) + emb_y = get_emb(sin_inp_y).unsqueeze(1) + emb_x = get_emb(sin_inp_x) + + emb = torch.zeros((h, w, self.dim * 2), device=tensor.device, dtype=tensor.dtype) + emb[:, :, :self.dim] = emb_x + emb[:, :, self.dim:] = emb_y + + if not self.channel_last and self.transpose_output: + # cancelled out + pass + elif (not self.channel_last) or (self.transpose_output): + emb = emb.permute(2, 0, 1) + + self.cached_penc = emb.unsqueeze(0).repeat(batch_size, 1, 1, 1) + if num_objects is None: + return self.cached_penc + else: + return self.cached_penc.unsqueeze(1) + + +if __name__ == '__main__': + pe = PositionalEncoding(8).cuda() + input = torch.ones((1, 8, 8, 8)).cuda() + output = pe(input) + # print(output) + print(output[0, :, 0, 0]) + print(output[0, :, 0, 5]) + print(output[0, 0, :, 0]) + print(output[0, 0, 0, :]) diff --git a/preprocessing/matanyone/matanyone/model/transformer/transformer_layers.py b/preprocessing/matanyone/matanyone/model/transformer/transformer_layers.py new file mode 100644 index 0000000..0b57bf2 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/transformer/transformer_layers.py @@ -0,0 +1,161 @@ +# Modified from PyTorch nn.Transformer + +from typing import List, Callable + +import torch +from torch import Tensor +import torch.nn as nn +import torch.nn.functional as F +from ...model.channel_attn import CAResBlock + + +class SelfAttention(nn.Module): + def __init__(self, + dim: int, + nhead: int, + dropout: float = 0.0, + batch_first: bool = True, + add_pe_to_qkv: List[bool] = [True, True, False]): + super().__init__() + self.self_attn = nn.MultiheadAttention(dim, nhead, dropout=dropout, batch_first=batch_first) + self.norm = nn.LayerNorm(dim) + self.dropout = nn.Dropout(dropout) + self.add_pe_to_qkv = add_pe_to_qkv + + def forward(self, + x: torch.Tensor, + pe: torch.Tensor, + attn_mask: bool = None, + key_padding_mask: bool = None) -> torch.Tensor: + x = self.norm(x) + if any(self.add_pe_to_qkv): + x_with_pe = x + pe + q = x_with_pe if self.add_pe_to_qkv[0] else x + k = x_with_pe if self.add_pe_to_qkv[1] else x + v = x_with_pe if self.add_pe_to_qkv[2] else x + else: + q = k = v = x + + r = x + x = self.self_attn(q, k, v, attn_mask=attn_mask, key_padding_mask=key_padding_mask)[0] + return r + self.dropout(x) + + +# https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html#torch.nn.functional.scaled_dot_product_attention +class CrossAttention(nn.Module): + def __init__(self, + dim: int, + nhead: int, + dropout: float = 0.0, + batch_first: bool = True, + add_pe_to_qkv: List[bool] = [True, True, False], + residual: bool = True, + norm: bool = True): + super().__init__() + self.cross_attn = nn.MultiheadAttention(dim, + nhead, + dropout=dropout, + batch_first=batch_first) + if norm: + self.norm = nn.LayerNorm(dim) + else: + self.norm = nn.Identity() + self.dropout = nn.Dropout(dropout) + self.add_pe_to_qkv = add_pe_to_qkv + self.residual = residual + + def forward(self, + x: torch.Tensor, + mem: torch.Tensor, + x_pe: torch.Tensor, + mem_pe: torch.Tensor, + attn_mask: bool = None, + *, + need_weights: bool = False) -> (torch.Tensor, torch.Tensor): + x = self.norm(x) + if self.add_pe_to_qkv[0]: + q = x + x_pe + else: + q = x + + if any(self.add_pe_to_qkv[1:]): + mem_with_pe = mem + mem_pe + k = mem_with_pe if self.add_pe_to_qkv[1] else mem + v = mem_with_pe if self.add_pe_to_qkv[2] else mem + else: + k = v = mem + r = x + x, weights = self.cross_attn(q, + k, + v, + attn_mask=attn_mask, + need_weights=need_weights, + average_attn_weights=False) + + if self.residual: + return r + self.dropout(x), weights + else: + return self.dropout(x), weights + + +class FFN(nn.Module): + def __init__(self, dim_in: int, dim_ff: int, activation=F.relu): + super().__init__() + self.linear1 = nn.Linear(dim_in, dim_ff) + self.linear2 = nn.Linear(dim_ff, dim_in) + self.norm = nn.LayerNorm(dim_in) + + if isinstance(activation, str): + self.activation = _get_activation_fn(activation) + else: + self.activation = activation + + def forward(self, x: torch.Tensor) -> torch.Tensor: + r = x + x = self.norm(x) + x = self.linear2(self.activation(self.linear1(x))) + x = r + x + return x + + +class PixelFFN(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.dim = dim + self.conv = CAResBlock(dim, dim) + + def forward(self, pixel: torch.Tensor, pixel_flat: torch.Tensor) -> torch.Tensor: + # pixel: batch_size * num_objects * dim * H * W + # pixel_flat: (batch_size*num_objects) * (H*W) * dim + bs, num_objects, _, h, w = pixel.shape + pixel_flat = pixel_flat.view(bs * num_objects, h, w, self.dim) + pixel_flat = pixel_flat.permute(0, 3, 1, 2).contiguous() + + x = self.conv(pixel_flat) + x = x.view(bs, num_objects, self.dim, h, w) + return x + + +class OutputFFN(nn.Module): + def __init__(self, dim_in: int, dim_out: int, activation=F.relu): + super().__init__() + self.linear1 = nn.Linear(dim_in, dim_out) + self.linear2 = nn.Linear(dim_out, dim_out) + + if isinstance(activation, str): + self.activation = _get_activation_fn(activation) + else: + self.activation = activation + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.linear2(self.activation(self.linear1(x))) + return x + + +def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]: + if activation == "relu": + return F.relu + elif activation == "gelu": + return F.gelu + + raise RuntimeError("activation should be relu/gelu, not {}".format(activation)) diff --git a/preprocessing/matanyone/matanyone/model/utils/__init__.py b/preprocessing/matanyone/matanyone/model/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/preprocessing/matanyone/matanyone/model/utils/memory_utils.py b/preprocessing/matanyone/matanyone/model/utils/memory_utils.py new file mode 100644 index 0000000..e7dd5e7 --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/utils/memory_utils.py @@ -0,0 +1,107 @@ +import math +import torch +from typing import Optional, Union, Tuple + + +# @torch.jit.script +def get_similarity(mk: torch.Tensor, + ms: torch.Tensor, + qk: torch.Tensor, + qe: torch.Tensor, + add_batch_dim: bool = False, + uncert_mask = None) -> torch.Tensor: + # used for training/inference and memory reading/memory potentiation + # mk: B x CK x [N] - Memory keys + # ms: B x 1 x [N] - Memory shrinkage + # qk: B x CK x [HW/P] - Query keys + # qe: B x CK x [HW/P] - Query selection + # Dimensions in [] are flattened + # Return: B*N*HW + if add_batch_dim: + mk, ms = mk.unsqueeze(0), ms.unsqueeze(0) + qk, qe = qk.unsqueeze(0), qe.unsqueeze(0) + + CK = mk.shape[1] + + mk = mk.flatten(start_dim=2) + ms = ms.flatten(start_dim=1).unsqueeze(2) if ms is not None else None + qk = qk.flatten(start_dim=2) + qe = qe.flatten(start_dim=2) if qe is not None else None + + # query token selection based on temporal sparsity + if uncert_mask is not None: + uncert_mask = uncert_mask.flatten(start_dim=2) + uncert_mask = uncert_mask.expand(-1, 64, -1) + qk = qk * uncert_mask + qe = qe * uncert_mask + + if qe is not None: + # See XMem's appendix for derivation + mk = mk.transpose(1, 2) + a_sq = (mk.pow(2) @ qe) + two_ab = 2 * (mk @ (qk * qe)) + b_sq = (qe * qk.pow(2)).sum(1, keepdim=True) + similarity = (-a_sq + two_ab - b_sq) + else: + # similar to STCN if we don't have the selection term + a_sq = mk.pow(2).sum(1).unsqueeze(2) + two_ab = 2 * (mk.transpose(1, 2) @ qk) + similarity = (-a_sq + two_ab) + + if ms is not None: + similarity = similarity * ms / math.sqrt(CK) # B*N*HW + else: + similarity = similarity / math.sqrt(CK) # B*N*HW + + return similarity + + +def do_softmax( + similarity: torch.Tensor, + top_k: Optional[int] = None, + inplace: bool = False, + return_usage: bool = False) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + # normalize similarity with top-k softmax + # similarity: B x N x [HW/P] + # use inplace with care + if top_k is not None: + values, indices = torch.topk(similarity, k=top_k, dim=1) + + x_exp = values.exp_() + x_exp /= torch.sum(x_exp, dim=1, keepdim=True) + if inplace: + similarity.zero_().scatter_(1, indices, x_exp) # B*N*HW + affinity = similarity + else: + affinity = torch.zeros_like(similarity).scatter_(1, indices, x_exp) # B*N*HW + else: + maxes = torch.max(similarity, dim=1, keepdim=True)[0] + x_exp = torch.exp(similarity - maxes) + x_exp_sum = torch.sum(x_exp, dim=1, keepdim=True) + affinity = x_exp / x_exp_sum + indices = None + + if return_usage: + return affinity, affinity.sum(dim=2) + + return affinity + + +def get_affinity(mk: torch.Tensor, ms: torch.Tensor, qk: torch.Tensor, + qe: torch.Tensor, uncert_mask = None) -> torch.Tensor: + # shorthand used in training with no top-k + similarity = get_similarity(mk, ms, qk, qe, uncert_mask=uncert_mask) + affinity = do_softmax(similarity) + return affinity + +def readout(affinity: torch.Tensor, mv: torch.Tensor, uncert_mask: torch.Tensor=None) -> torch.Tensor: + B, CV, T, H, W = mv.shape + + mo = mv.view(B, CV, T * H * W) + mem = torch.bmm(mo, affinity) + if uncert_mask is not None: + uncert_mask = uncert_mask.flatten(start_dim=2).expand(-1, CV, -1) + mem = mem * uncert_mask + mem = mem.view(B, CV, H, W) + + return mem diff --git a/preprocessing/matanyone/matanyone/model/utils/parameter_groups.py b/preprocessing/matanyone/matanyone/model/utils/parameter_groups.py new file mode 100644 index 0000000..177866a --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/utils/parameter_groups.py @@ -0,0 +1,72 @@ +import logging + +log = logging.getLogger() + + +def get_parameter_groups(model, stage_cfg, print_log=False): + """ + Assign different weight decays and learning rates to different parameters. + Returns a parameter group which can be passed to the optimizer. + """ + weight_decay = stage_cfg.weight_decay + embed_weight_decay = stage_cfg.embed_weight_decay + backbone_lr_ratio = stage_cfg.backbone_lr_ratio + base_lr = stage_cfg.learning_rate + + backbone_params = [] + embed_params = [] + other_params = [] + + embedding_names = ['summary_pos', 'query_init', 'query_emb', 'obj_pe'] + embedding_names = [e + '.weight' for e in embedding_names] + + # inspired by detectron2 + memo = set() + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + # Avoid duplicating parameters + if param in memo: + continue + memo.add(param) + + if name.startswith('module'): + name = name[7:] + + inserted = False + if name.startswith('pixel_encoder.'): + backbone_params.append(param) + inserted = True + if print_log: + log.info(f'{name} counted as a backbone parameter.') + else: + for e in embedding_names: + if name.endswith(e): + embed_params.append(param) + inserted = True + if print_log: + log.info(f'{name} counted as an embedding parameter.') + break + + if not inserted: + other_params.append(param) + + parameter_groups = [ + { + 'params': backbone_params, + 'lr': base_lr * backbone_lr_ratio, + 'weight_decay': weight_decay + }, + { + 'params': embed_params, + 'lr': base_lr, + 'weight_decay': embed_weight_decay + }, + { + 'params': other_params, + 'lr': base_lr, + 'weight_decay': weight_decay + }, + ] + + return parameter_groups \ No newline at end of file diff --git a/preprocessing/matanyone/matanyone/model/utils/resnet.py b/preprocessing/matanyone/matanyone/model/utils/resnet.py new file mode 100644 index 0000000..44886ee --- /dev/null +++ b/preprocessing/matanyone/matanyone/model/utils/resnet.py @@ -0,0 +1,179 @@ +""" +resnet.py - A modified ResNet structure +We append extra channels to the first conv by some network surgery +""" + +from collections import OrderedDict +import math + +import torch +import torch.nn as nn +from torch.utils import model_zoo + + +def load_weights_add_extra_dim(target, source_state, extra_dim=1): + new_dict = OrderedDict() + + for k1, v1 in target.state_dict().items(): + if 'num_batches_tracked' not in k1: + if k1 in source_state: + tar_v = source_state[k1] + + if v1.shape != tar_v.shape: + # Init the new segmentation channel with zeros + # print(v1.shape, tar_v.shape) + c, _, w, h = v1.shape + pads = torch.zeros((c, extra_dim, w, h), device=tar_v.device) + nn.init.orthogonal_(pads) + tar_v = torch.cat([tar_v, pads], 1) + + new_dict[k1] = tar_v + + target.load_state_dict(new_dict) + + +model_urls = { + 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', + 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', +} + + +def conv3x3(in_planes, out_planes, stride=1, dilation=1): + return nn.Conv2d(in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + dilation=dilation, + bias=False) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(inplanes, planes, stride=stride, dilation=dilation) + self.bn1 = nn.BatchNorm2d(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes, stride=1, dilation=dilation) + self.bn2 = nn.BatchNorm2d(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1): + super(Bottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, + planes, + kernel_size=3, + stride=stride, + dilation=dilation, + padding=dilation, + bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * 4) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + def __init__(self, block, layers=(3, 4, 23, 3), extra_dim=0): + self.inplanes = 64 + super(ResNet, self).__init__() + self.conv1 = nn.Conv2d(3 + extra_dim, 64, kernel_size=7, stride=2, padding=3, bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + def _make_layer(self, block, planes, blocks, stride=1, dilation=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [block(self.inplanes, planes, stride, downsample)] + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes, dilation=dilation)) + + return nn.Sequential(*layers) + + +def resnet18(pretrained=True, extra_dim=0): + model = ResNet(BasicBlock, [2, 2, 2, 2], extra_dim) + if pretrained: + load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet18']), extra_dim) + return model + + +def resnet50(pretrained=True, extra_dim=0): + model = ResNet(Bottleneck, [3, 4, 6, 3], extra_dim) + if pretrained: + load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet50']), extra_dim) + return model diff --git a/preprocessing/matanyone/matanyone_wrapper.py b/preprocessing/matanyone/matanyone_wrapper.py new file mode 100644 index 0000000..82fb773 --- /dev/null +++ b/preprocessing/matanyone/matanyone_wrapper.py @@ -0,0 +1,73 @@ +import tqdm +import torch +from torchvision.transforms.functional import to_tensor +import numpy as np +import random +import cv2 + +def gen_dilate(alpha, min_kernel_size, max_kernel_size): + kernel_size = random.randint(min_kernel_size, max_kernel_size) + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size,kernel_size)) + fg_and_unknown = np.array(np.not_equal(alpha, 0).astype(np.float32)) + dilate = cv2.dilate(fg_and_unknown, kernel, iterations=1)*255 + return dilate.astype(np.float32) + +def gen_erosion(alpha, min_kernel_size, max_kernel_size): + kernel_size = random.randint(min_kernel_size, max_kernel_size) + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size,kernel_size)) + fg = np.array(np.equal(alpha, 255).astype(np.float32)) + erode = cv2.erode(fg, kernel, iterations=1)*255 + return erode.astype(np.float32) + +@torch.inference_mode() +@torch.amp.autocast('cuda') +def matanyone(processor, frames_np, mask, r_erode=0, r_dilate=0, n_warmup=10): + """ + Args: + frames_np: [(H,W,C)]*n, uint8 + mask: (H,W), uint8 + Outputs: + com: [(H,W,C)]*n, uint8 + pha: [(H,W,C)]*n, uint8 + """ + + # print(f'===== [r_erode] {r_erode}; [r_dilate] {r_dilate} =====') + bgr = (np.array([120, 255, 155], dtype=np.float32)/255).reshape((1, 1, 3)) + objects = [1] + + # [optional] erode & dilate on given seg mask + if r_dilate > 0: + mask = gen_dilate(mask, r_dilate, r_dilate) + if r_erode > 0: + mask = gen_erosion(mask, r_erode, r_erode) + + mask = torch.from_numpy(mask).cuda() + + frames_np = [frames_np[0]]* n_warmup + frames_np + + frames = [] + phas = [] + for ti, frame_single in tqdm.tqdm(enumerate(frames_np)): + image = to_tensor(frame_single).cuda().float() + + if ti == 0: + output_prob = processor.step(image, mask, objects=objects) # encode given mask + output_prob = processor.step(image, first_frame_pred=True) # clear past memory for warmup frames + else: + if ti <= n_warmup: + output_prob = processor.step(image, first_frame_pred=True) # clear past memory for warmup frames + else: + output_prob = processor.step(image) + + # convert output probabilities to an object mask + mask = processor.output_prob_to_mask(output_prob) + + pha = mask.unsqueeze(2).cpu().numpy() + com_np = frame_single / 255. * pha + bgr * (1 - pha) + + # DONOT save the warmup frames + if ti > (n_warmup-1): + frames.append((com_np*255).astype(np.uint8)) + phas.append((pha*255).astype(np.uint8)) + + return frames, phas \ No newline at end of file diff --git a/preprocessing/matanyone/tools/__init__.py b/preprocessing/matanyone/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/preprocessing/matanyone/tools/base_segmenter.py b/preprocessing/matanyone/tools/base_segmenter.py new file mode 100644 index 0000000..096038e --- /dev/null +++ b/preprocessing/matanyone/tools/base_segmenter.py @@ -0,0 +1,141 @@ +import time +import torch +import cv2 +from PIL import Image, ImageDraw, ImageOps +import numpy as np +from typing import Union +from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator +import matplotlib.pyplot as plt +import PIL +from .mask_painter import mask_painter + + +class BaseSegmenter: + def __init__(self, SAM_checkpoint, model_type, device='cuda:0'): + """ + device: model device + SAM_checkpoint: path of SAM checkpoint + model_type: vit_b, vit_l, vit_h + """ + print(f"Initializing BaseSegmenter to {device}") + assert model_type in ['vit_b', 'vit_l', 'vit_h'], 'model_type must be vit_b, vit_l, or vit_h' + + self.device = device + # SAM_checkpoint = None + self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 + from accelerate import init_empty_weights + + # self.model = sam_model_registry[model_type](checkpoint=SAM_checkpoint) + with init_empty_weights(): + self.model = sam_model_registry[model_type](checkpoint=SAM_checkpoint) + from mmgp import offload + # self.model.to(torch.float16) + # offload.save_model(self.model, "ckpts/mask/sam_vit_h_4b8939_fp16.safetensors") + + offload.load_model_data(self.model, "ckpts/mask/sam_vit_h_4b8939_fp16.safetensors") + self.model.to(torch.float32) # need to be optimized, if not f32 crappy precision + self.model.to(device=self.device) + self.predictor = SamPredictor(self.model) + self.embedded = False + + @torch.no_grad() + def set_image(self, image: np.ndarray): + # PIL.open(image_path) 3channel: RGB + # image embedding: avoid encode the same image multiple times + self.orignal_image = image + if self.embedded: + print('repeat embedding, please reset_image.') + return + self.predictor.set_image(image) + self.embedded = True + return + + @torch.no_grad() + def reset_image(self): + # reset image embeding + self.predictor.reset_image() + self.embedded = False + + def predict(self, prompts, mode, multimask=True): + """ + image: numpy array, h, w, 3 + prompts: dictionary, 3 keys: 'point_coords', 'point_labels', 'mask_input' + prompts['point_coords']: numpy array [N,2] + prompts['point_labels']: numpy array [1,N] + prompts['mask_input']: numpy array [1,256,256] + mode: 'point' (points only), 'mask' (mask only), 'both' (consider both) + mask_outputs: True (return 3 masks), False (return 1 mask only) + whem mask_outputs=True, mask_input=logits[np.argmax(scores), :, :][None, :, :] + """ + assert self.embedded, 'prediction is called before set_image (feature embedding).' + assert mode in ['point', 'mask', 'both'], 'mode must be point, mask, or both' + + with torch.autocast(device_type='cuda', dtype=torch.float16): + if mode == 'point': + masks, scores, logits = self.predictor.predict(point_coords=prompts['point_coords'], + point_labels=prompts['point_labels'], + multimask_output=multimask) + elif mode == 'mask': + masks, scores, logits = self.predictor.predict(mask_input=prompts['mask_input'], + multimask_output=multimask) + elif mode == 'both': # both + masks, scores, logits = self.predictor.predict(point_coords=prompts['point_coords'], + point_labels=prompts['point_labels'], + mask_input=prompts['mask_input'], + multimask_output=multimask) + else: + raise("Not implement now!") + # masks (n, h, w), scores (n,), logits (n, 256, 256) + return masks, scores, logits + + +if __name__ == "__main__": + # load and show an image + image = cv2.imread('/hhd3/gaoshang/truck.jpg') + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # numpy array (h, w, 3) + + # initialise BaseSegmenter + SAM_checkpoint= '/ssd1/gaomingqi/checkpoints/sam_vit_h_4b8939.pth' + model_type = 'vit_h' + device = "cuda:4" + base_segmenter = BaseSegmenter(SAM_checkpoint=SAM_checkpoint, model_type=model_type, device=device) + + # image embedding (once embedded, multiple prompts can be applied) + base_segmenter.set_image(image) + + # examples + # point only ------------------------ + mode = 'point' + prompts = { + 'point_coords': np.array([[500, 375], [1125, 625]]), + 'point_labels': np.array([1, 1]), + } + masks, scores, logits = base_segmenter.predict(prompts, mode, multimask=False) # masks (n, h, w), scores (n,), logits (n, 256, 256) + painted_image = mask_painter(image, masks[np.argmax(scores)].astype('uint8'), background_alpha=0.8) + painted_image = cv2.cvtColor(painted_image, cv2.COLOR_RGB2BGR) # numpy array (h, w, 3) + cv2.imwrite('/hhd3/gaoshang/truck_point.jpg', painted_image) + + # both ------------------------ + mode = 'both' + mask_input = logits[np.argmax(scores), :, :] + prompts = {'mask_input': mask_input [None, :, :]} + prompts = { + 'point_coords': np.array([[500, 375], [1125, 625]]), + 'point_labels': np.array([1, 0]), + 'mask_input': mask_input[None, :, :] + } + masks, scores, logits = base_segmenter.predict(prompts, mode, multimask=True) # masks (n, h, w), scores (n,), logits (n, 256, 256) + painted_image = mask_painter(image, masks[np.argmax(scores)].astype('uint8'), background_alpha=0.8) + painted_image = cv2.cvtColor(painted_image, cv2.COLOR_RGB2BGR) # numpy array (h, w, 3) + cv2.imwrite('/hhd3/gaoshang/truck_both.jpg', painted_image) + + # mask only ------------------------ + mode = 'mask' + mask_input = logits[np.argmax(scores), :, :] + + prompts = {'mask_input': mask_input[None, :, :]} + + masks, scores, logits = base_segmenter.predict(prompts, mode, multimask=True) # masks (n, h, w), scores (n,), logits (n, 256, 256) + painted_image = mask_painter(image, masks[np.argmax(scores)].astype('uint8'), background_alpha=0.8) + painted_image = cv2.cvtColor(painted_image, cv2.COLOR_RGB2BGR) # numpy array (h, w, 3) + cv2.imwrite('/hhd3/gaoshang/truck_mask.jpg', painted_image) diff --git a/preprocessing/matanyone/tools/download_util.py b/preprocessing/matanyone/tools/download_util.py new file mode 100644 index 0000000..5e8fb1b --- /dev/null +++ b/preprocessing/matanyone/tools/download_util.py @@ -0,0 +1,109 @@ +import math +import os +import requests +from torch.hub import download_url_to_file, get_dir +from tqdm import tqdm +from urllib.parse import urlparse + +def sizeof_fmt(size, suffix='B'): + """Get human readable file size. + + Args: + size (int): File size. + suffix (str): Suffix. Default: 'B'. + + Return: + str: Formated file siz. + """ + for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']: + if abs(size) < 1024.0: + return f'{size:3.1f} {unit}{suffix}' + size /= 1024.0 + return f'{size:3.1f} Y{suffix}' + + +def download_file_from_google_drive(file_id, save_path): + """Download files from google drive. + Ref: + https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive # noqa E501 + Args: + file_id (str): File id. + save_path (str): Save path. + """ + + session = requests.Session() + URL = 'https://docs.google.com/uc?export=download' + params = {'id': file_id} + + response = session.get(URL, params=params, stream=True) + token = get_confirm_token(response) + if token: + params['confirm'] = token + response = session.get(URL, params=params, stream=True) + + # get file size + response_file_size = session.get(URL, params=params, stream=True, headers={'Range': 'bytes=0-2'}) + print(response_file_size) + if 'Content-Range' in response_file_size.headers: + file_size = int(response_file_size.headers['Content-Range'].split('/')[1]) + else: + file_size = None + + save_response_content(response, save_path, file_size) + + +def get_confirm_token(response): + for key, value in response.cookies.items(): + if key.startswith('download_warning'): + return value + return None + + +def save_response_content(response, destination, file_size=None, chunk_size=32768): + if file_size is not None: + pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk') + + readable_file_size = sizeof_fmt(file_size) + else: + pbar = None + + with open(destination, 'wb') as f: + downloaded_size = 0 + for chunk in response.iter_content(chunk_size): + downloaded_size += chunk_size + if pbar is not None: + pbar.update(1) + pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} / {readable_file_size}') + if chunk: # filter out keep-alive new chunks + f.write(chunk) + if pbar is not None: + pbar.close() + + +def load_file_from_url(url, model_dir=None, progress=True, file_name=None): + """Load file form http url, will download models if necessary. + Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py + Args: + url (str): URL to be downloaded. + model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir. + Default: None. + progress (bool): Whether to show the download progress. Default: True. + file_name (str): The downloaded file name. If None, use the file name in the url. Default: None. + Returns: + str: The path to the downloaded file. + """ + if model_dir is None: # use the pytorch hub_dir + hub_dir = get_dir() + model_dir = os.path.join(hub_dir, 'checkpoints') + + os.makedirs(model_dir, exist_ok=True) + + parts = urlparse(url) + filename = os.path.basename(parts.path) + if file_name is not None: + filename = file_name + cached_file = os.path.abspath(os.path.join(model_dir, filename)) + if not os.path.exists(cached_file): + print(f'Downloading: "{url}" to {cached_file}\n') + download_url_to_file(url, cached_file, hash_prefix=None, progress=progress) + return cached_file \ No newline at end of file diff --git a/preprocessing/matanyone/tools/interact_tools.py b/preprocessing/matanyone/tools/interact_tools.py new file mode 100644 index 0000000..c70b8c4 --- /dev/null +++ b/preprocessing/matanyone/tools/interact_tools.py @@ -0,0 +1,99 @@ +import time +import torch +import cv2 +from PIL import Image, ImageDraw, ImageOps +import numpy as np +from typing import Union +from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator +import matplotlib.pyplot as plt +import PIL +from .mask_painter import mask_painter as mask_painter2 +from .base_segmenter import BaseSegmenter +from .painter import mask_painter, point_painter +import os +import requests +import sys + + +mask_color = 3 +mask_alpha = 0.7 +contour_color = 1 +contour_width = 5 +point_color_ne = 8 +point_color_ps = 50 +point_alpha = 0.9 +point_radius = 15 +contour_color = 2 +contour_width = 5 + + +class SamControler(): + def __init__(self, SAM_checkpoint, model_type, device): + ''' + initialize sam controler + ''' + self.sam_controler = BaseSegmenter(SAM_checkpoint, model_type, device) + + + # def seg_again(self, image: np.ndarray): + # ''' + # it is used when interact in video + # ''' + # self.sam_controler.reset_image() + # self.sam_controler.set_image(image) + # return + + + def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True,mask_color=3): + ''' + it is used in first frame in video + return: mask, logit, painted image(mask+point) + ''' + # self.sam_controler.set_image(image) + origal_image = self.sam_controler.orignal_image + neg_flag = labels[-1] + if neg_flag==1: + #find neg + prompts = { + 'point_coords': points, + 'point_labels': labels, + } + masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask) + mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :] + prompts = { + 'point_coords': points, + 'point_labels': labels, + 'mask_input': logit[None, :, :] + } + masks, scores, logits = self.sam_controler.predict(prompts, 'both', multimask) + mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :] + else: + #find positive + prompts = { + 'point_coords': points, + 'point_labels': labels, + } + masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask) + mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :] + + + assert len(points)==len(labels) + + painted_image = mask_painter(image, mask.astype('uint8'), mask_color, mask_alpha, contour_color, contour_width) + painted_image = point_painter(painted_image, np.squeeze(points[np.argwhere(labels>0)],axis = 1), point_color_ne, point_alpha, point_radius, contour_color, contour_width) + painted_image = point_painter(painted_image, np.squeeze(points[np.argwhere(labels<1)],axis = 1), point_color_ps, point_alpha, point_radius, contour_color, contour_width) + painted_image = Image.fromarray(painted_image) + + return mask, logit, painted_image + + + + + + + + + + + + \ No newline at end of file diff --git a/preprocessing/matanyone/tools/mask_painter.py b/preprocessing/matanyone/tools/mask_painter.py new file mode 100644 index 0000000..f471ea0 --- /dev/null +++ b/preprocessing/matanyone/tools/mask_painter.py @@ -0,0 +1,288 @@ +import cv2 +import torch +import numpy as np +from PIL import Image +import copy +import time + + +def colormap(rgb=True): + color_list = np.array( + [ + 0.000, 0.000, 0.000, + 1.000, 1.000, 1.000, + 1.000, 0.498, 0.313, + 0.392, 0.581, 0.929, + 0.000, 0.447, 0.741, + 0.850, 0.325, 0.098, + 0.929, 0.694, 0.125, + 0.494, 0.184, 0.556, + 0.466, 0.674, 0.188, + 0.301, 0.745, 0.933, + 0.635, 0.078, 0.184, + 0.300, 0.300, 0.300, + 0.600, 0.600, 0.600, + 1.000, 0.000, 0.000, + 1.000, 0.500, 0.000, + 0.749, 0.749, 0.000, + 0.000, 1.000, 0.000, + 0.000, 0.000, 1.000, + 0.667, 0.000, 1.000, + 0.333, 0.333, 0.000, + 0.333, 0.667, 0.000, + 0.333, 1.000, 0.000, + 0.667, 0.333, 0.000, + 0.667, 0.667, 0.000, + 0.667, 1.000, 0.000, + 1.000, 0.333, 0.000, + 1.000, 0.667, 0.000, + 1.000, 1.000, 0.000, + 0.000, 0.333, 0.500, + 0.000, 0.667, 0.500, + 0.000, 1.000, 0.500, + 0.333, 0.000, 0.500, + 0.333, 0.333, 0.500, + 0.333, 0.667, 0.500, + 0.333, 1.000, 0.500, + 0.667, 0.000, 0.500, + 0.667, 0.333, 0.500, + 0.667, 0.667, 0.500, + 0.667, 1.000, 0.500, + 1.000, 0.000, 0.500, + 1.000, 0.333, 0.500, + 1.000, 0.667, 0.500, + 1.000, 1.000, 0.500, + 0.000, 0.333, 1.000, + 0.000, 0.667, 1.000, + 0.000, 1.000, 1.000, + 0.333, 0.000, 1.000, + 0.333, 0.333, 1.000, + 0.333, 0.667, 1.000, + 0.333, 1.000, 1.000, + 0.667, 0.000, 1.000, + 0.667, 0.333, 1.000, + 0.667, 0.667, 1.000, + 0.667, 1.000, 1.000, + 1.000, 0.000, 1.000, + 1.000, 0.333, 1.000, + 1.000, 0.667, 1.000, + 0.167, 0.000, 0.000, + 0.333, 0.000, 0.000, + 0.500, 0.000, 0.000, + 0.667, 0.000, 0.000, + 0.833, 0.000, 0.000, + 1.000, 0.000, 0.000, + 0.000, 0.167, 0.000, + 0.000, 0.333, 0.000, + 0.000, 0.500, 0.000, + 0.000, 0.667, 0.000, + 0.000, 0.833, 0.000, + 0.000, 1.000, 0.000, + 0.000, 0.000, 0.167, + 0.000, 0.000, 0.333, + 0.000, 0.000, 0.500, + 0.000, 0.000, 0.667, + 0.000, 0.000, 0.833, + 0.000, 0.000, 1.000, + 0.143, 0.143, 0.143, + 0.286, 0.286, 0.286, + 0.429, 0.429, 0.429, + 0.571, 0.571, 0.571, + 0.714, 0.714, 0.714, + 0.857, 0.857, 0.857 + ] + ).astype(np.float32) + color_list = color_list.reshape((-1, 3)) * 255 + if not rgb: + color_list = color_list[:, ::-1] + return color_list + + +color_list = colormap() +color_list = color_list.astype('uint8').tolist() + + +def vis_add_mask(image, background_mask, contour_mask, background_color, contour_color, background_alpha, contour_alpha): + background_color = np.array(background_color) + contour_color = np.array(contour_color) + + # background_mask = 1 - background_mask + # contour_mask = 1 - contour_mask + + for i in range(3): + image[:, :, i] = image[:, :, i] * (1-background_alpha+background_mask*background_alpha) \ + + background_color[i] * (background_alpha-background_mask*background_alpha) + + image[:, :, i] = image[:, :, i] * (1-contour_alpha+contour_mask*contour_alpha) \ + + contour_color[i] * (contour_alpha-contour_mask*contour_alpha) + + return image.astype('uint8') + + +def mask_generator_00(mask, background_radius, contour_radius): + # no background width when '00' + # distance map + dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3) + dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3) + dist_map = dist_transform_fore - dist_transform_back + # ...:::!!!:::... + contour_radius += 2 + contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius)) + contour_mask = contour_mask / np.max(contour_mask) + contour_mask[contour_mask>0.5] = 1. + + return mask, contour_mask + + +def mask_generator_01(mask, background_radius, contour_radius): + # no background width when '00' + # distance map + dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3) + dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3) + dist_map = dist_transform_fore - dist_transform_back + # ...:::!!!:::... + contour_radius += 2 + contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius)) + contour_mask = contour_mask / np.max(contour_mask) + return mask, contour_mask + + +def mask_generator_10(mask, background_radius, contour_radius): + # distance map + dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3) + dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3) + dist_map = dist_transform_fore - dist_transform_back + # .....:::::!!!!! + background_mask = np.clip(dist_map, -background_radius, background_radius) + background_mask = (background_mask - np.min(background_mask)) + background_mask = background_mask / np.max(background_mask) + # ...:::!!!:::... + contour_radius += 2 + contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius)) + contour_mask = contour_mask / np.max(contour_mask) + contour_mask[contour_mask>0.5] = 1. + return background_mask, contour_mask + + +def mask_generator_11(mask, background_radius, contour_radius): + # distance map + dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3) + dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3) + dist_map = dist_transform_fore - dist_transform_back + # .....:::::!!!!! + background_mask = np.clip(dist_map, -background_radius, background_radius) + background_mask = (background_mask - np.min(background_mask)) + background_mask = background_mask / np.max(background_mask) + # ...:::!!!:::... + contour_radius += 2 + contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius)) + contour_mask = contour_mask / np.max(contour_mask) + return background_mask, contour_mask + + +def mask_painter(input_image, input_mask, background_alpha=0.5, background_blur_radius=7, contour_width=3, contour_color=3, contour_alpha=1, mode='11'): + """ + Input: + input_image: numpy array + input_mask: numpy array + background_alpha: transparency of background, [0, 1], 1: all black, 0: do nothing + background_blur_radius: radius of background blur, must be odd number + contour_width: width of mask contour, must be odd number + contour_color: color index (in color map) of mask contour, 0: black, 1: white, >1: others + contour_alpha: transparency of mask contour, [0, 1], if 0: no contour highlighted + mode: painting mode, '00', no blur, '01' only blur contour, '10' only blur background, '11' blur both + + Output: + painted_image: numpy array + """ + assert input_image.shape[:2] == input_mask.shape, 'different shape' + assert background_blur_radius % 2 * contour_width % 2 > 0, 'background_blur_radius and contour_width must be ODD' + assert mode in ['00', '01', '10', '11'], 'mode should be 00, 01, 10, or 11' + + # downsample input image and mask + width, height = input_image.shape[0], input_image.shape[1] + res = 1024 + ratio = min(1.0 * res / max(width, height), 1.0) + input_image = cv2.resize(input_image, (int(height*ratio), int(width*ratio))) + input_mask = cv2.resize(input_mask, (int(height*ratio), int(width*ratio))) + + # 0: background, 1: foreground + msk = np.clip(input_mask, 0, 1) + + # generate masks for background and contour pixels + background_radius = (background_blur_radius - 1) // 2 + contour_radius = (contour_width - 1) // 2 + generator_dict = {'00':mask_generator_00, '01':mask_generator_01, '10':mask_generator_10, '11':mask_generator_11} + background_mask, contour_mask = generator_dict[mode](msk, background_radius, contour_radius) + + # paint + painted_image = vis_add_mask\ + (input_image, background_mask, contour_mask, color_list[0], color_list[contour_color], background_alpha, contour_alpha) # black for background + + return painted_image + + +if __name__ == '__main__': + + background_alpha = 0.7 # transparency of background 1: all black, 0: do nothing + background_blur_radius = 31 # radius of background blur, must be odd number + contour_width = 11 # contour width, must be odd number + contour_color = 3 # id in color map, 0: black, 1: white, >1: others + contour_alpha = 1 # transparency of background, 0: no contour highlighted + + # load input image and mask + input_image = np.array(Image.open('./test_img/painter_input_image.jpg').convert('RGB')) + input_mask = np.array(Image.open('./test_img/painter_input_mask.jpg').convert('P')) + + # paint + overall_time_1 = 0 + overall_time_2 = 0 + overall_time_3 = 0 + overall_time_4 = 0 + overall_time_5 = 0 + + for i in range(50): + t2 = time.time() + painted_image_00 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='00') + e2 = time.time() + + t3 = time.time() + painted_image_10 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='10') + e3 = time.time() + + t1 = time.time() + painted_image = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha) + e1 = time.time() + + t4 = time.time() + painted_image_01 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='01') + e4 = time.time() + + t5 = time.time() + painted_image_11 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='11') + e5 = time.time() + + overall_time_1 += (e1 - t1) + overall_time_2 += (e2 - t2) + overall_time_3 += (e3 - t3) + overall_time_4 += (e4 - t4) + overall_time_5 += (e5 - t5) + + print(f'average time w gaussian: {overall_time_1/50}') + print(f'average time w/o gaussian00: {overall_time_2/50}') + print(f'average time w/o gaussian10: {overall_time_3/50}') + print(f'average time w/o gaussian01: {overall_time_4/50}') + print(f'average time w/o gaussian11: {overall_time_5/50}') + + # save + painted_image_00 = Image.fromarray(painted_image_00) + painted_image_00.save('./test_img/painter_output_image_00.png') + + painted_image_10 = Image.fromarray(painted_image_10) + painted_image_10.save('./test_img/painter_output_image_10.png') + + painted_image_01 = Image.fromarray(painted_image_01) + painted_image_01.save('./test_img/painter_output_image_01.png') + + painted_image_11 = Image.fromarray(painted_image_11) + painted_image_11.save('./test_img/painter_output_image_11.png') diff --git a/preprocessing/matanyone/tools/misc.py b/preprocessing/matanyone/tools/misc.py new file mode 100644 index 0000000..43b8499 --- /dev/null +++ b/preprocessing/matanyone/tools/misc.py @@ -0,0 +1,131 @@ +import os +import re +import random +import time +import torch +import torch.nn as nn +import logging +import numpy as np +from os import path as osp + +def constant_init(module, val, bias=0): + if hasattr(module, 'weight') and module.weight is not None: + nn.init.constant_(module.weight, val) + if hasattr(module, 'bias') and module.bias is not None: + nn.init.constant_(module.bias, bias) + +initialized_logger = {} +def get_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=None): + """Get the root logger. + The logger will be initialized if it has not been initialized. By default a + StreamHandler will be added. If `log_file` is specified, a FileHandler will + also be added. + Args: + logger_name (str): root logger name. Default: 'basicsr'. + log_file (str | None): The log filename. If specified, a FileHandler + will be added to the root logger. + log_level (int): The root logger level. Note that only the process of + rank 0 is affected, while other processes will set the level to + "Error" and be silent most of the time. + Returns: + logging.Logger: The root logger. + """ + logger = logging.getLogger(logger_name) + # if the logger has been initialized, just return it + if logger_name in initialized_logger: + return logger + + format_str = '%(asctime)s %(levelname)s: %(message)s' + stream_handler = logging.StreamHandler() + stream_handler.setFormatter(logging.Formatter(format_str)) + logger.addHandler(stream_handler) + logger.propagate = False + + if log_file is not None: + logger.setLevel(log_level) + # add file handler + # file_handler = logging.FileHandler(log_file, 'w') + file_handler = logging.FileHandler(log_file, 'a') #Shangchen: keep the previous log + file_handler.setFormatter(logging.Formatter(format_str)) + file_handler.setLevel(log_level) + logger.addHandler(file_handler) + initialized_logger[logger_name] = True + return logger + + +IS_HIGH_VERSION = [int(m) for m in list(re.findall(r"^([0-9]+)\.([0-9]+)\.([0-9]+)([^0-9][a-zA-Z0-9]*)?(\+git.*)?$",\ + torch.__version__)[0][:3])] >= [1, 12, 0] + +def gpu_is_available(): + if IS_HIGH_VERSION: + if torch.backends.mps.is_available(): + return True + return True if torch.cuda.is_available() and torch.backends.cudnn.is_available() else False + +def get_device(gpu_id=None): + if gpu_id is None: + gpu_str = '' + elif isinstance(gpu_id, int): + gpu_str = f':{gpu_id}' + else: + raise TypeError('Input should be int value.') + + if IS_HIGH_VERSION: + if torch.backends.mps.is_available(): + return torch.device('mps'+gpu_str) + return torch.device('cuda'+gpu_str if torch.cuda.is_available() and torch.backends.cudnn.is_available() else 'cpu') + + +def set_random_seed(seed): + """Set random seeds.""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def get_time_str(): + return time.strftime('%Y%m%d_%H%M%S', time.localtime()) + + +def scandir(dir_path, suffix=None, recursive=False, full_path=False): + """Scan a directory to find the interested files. + + Args: + dir_path (str): Path of the directory. + suffix (str | tuple(str), optional): File suffix that we are + interested in. Default: None. + recursive (bool, optional): If set to True, recursively scan the + directory. Default: False. + full_path (bool, optional): If set to True, include the dir_path. + Default: False. + + Returns: + A generator for all the interested files with relative pathes. + """ + + if (suffix is not None) and not isinstance(suffix, (str, tuple)): + raise TypeError('"suffix" must be a string or tuple of strings') + + root = dir_path + + def _scandir(dir_path, suffix, recursive): + for entry in os.scandir(dir_path): + if not entry.name.startswith('.') and entry.is_file(): + if full_path: + return_path = entry.path + else: + return_path = osp.relpath(entry.path, root) + + if suffix is None: + yield return_path + elif return_path.endswith(suffix): + yield return_path + else: + if recursive: + yield from _scandir(entry.path, suffix=suffix, recursive=recursive) + else: + continue + + return _scandir(dir_path, suffix=suffix, recursive=recursive) \ No newline at end of file diff --git a/preprocessing/matanyone/tools/painter.py b/preprocessing/matanyone/tools/painter.py new file mode 100644 index 0000000..0e711d3 --- /dev/null +++ b/preprocessing/matanyone/tools/painter.py @@ -0,0 +1,215 @@ +# paint masks, contours, or points on images, with specified colors +import cv2 +import torch +import numpy as np +from PIL import Image +import copy +import time + + +def colormap(rgb=True): + color_list = np.array( + [ + 0.000, 0.000, 0.000, + 1.000, 1.000, 1.000, + 1.000, 0.498, 0.313, + 0.392, 0.581, 0.929, + 0.000, 0.447, 0.741, + 0.850, 0.325, 0.098, + 0.929, 0.694, 0.125, + 0.494, 0.184, 0.556, + 0.466, 0.674, 0.188, + 0.301, 0.745, 0.933, + 0.635, 0.078, 0.184, + 0.300, 0.300, 0.300, + 0.600, 0.600, 0.600, + 1.000, 0.000, 0.000, + 1.000, 0.500, 0.000, + 0.749, 0.749, 0.000, + 0.000, 1.000, 0.000, + 0.000, 0.000, 1.000, + 0.667, 0.000, 1.000, + 0.333, 0.333, 0.000, + 0.333, 0.667, 0.000, + 0.333, 1.000, 0.000, + 0.667, 0.333, 0.000, + 0.667, 0.667, 0.000, + 0.667, 1.000, 0.000, + 1.000, 0.333, 0.000, + 1.000, 0.667, 0.000, + 1.000, 1.000, 0.000, + 0.000, 0.333, 0.500, + 0.000, 0.667, 0.500, + 0.000, 1.000, 0.500, + 0.333, 0.000, 0.500, + 0.333, 0.333, 0.500, + 0.333, 0.667, 0.500, + 0.333, 1.000, 0.500, + 0.667, 0.000, 0.500, + 0.667, 0.333, 0.500, + 0.667, 0.667, 0.500, + 0.667, 1.000, 0.500, + 1.000, 0.000, 0.500, + 1.000, 0.333, 0.500, + 1.000, 0.667, 0.500, + 1.000, 1.000, 0.500, + 0.000, 0.333, 1.000, + 0.000, 0.667, 1.000, + 0.000, 1.000, 1.000, + 0.333, 0.000, 1.000, + 0.333, 0.333, 1.000, + 0.333, 0.667, 1.000, + 0.333, 1.000, 1.000, + 0.667, 0.000, 1.000, + 0.667, 0.333, 1.000, + 0.667, 0.667, 1.000, + 0.667, 1.000, 1.000, + 1.000, 0.000, 1.000, + 1.000, 0.333, 1.000, + 1.000, 0.667, 1.000, + 0.167, 0.000, 0.000, + 0.333, 0.000, 0.000, + 0.500, 0.000, 0.000, + 0.667, 0.000, 0.000, + 0.833, 0.000, 0.000, + 1.000, 0.000, 0.000, + 0.000, 0.167, 0.000, + 0.000, 0.333, 0.000, + 0.000, 0.500, 0.000, + 0.000, 0.667, 0.000, + 0.000, 0.833, 0.000, + 0.000, 1.000, 0.000, + 0.000, 0.000, 0.167, + 0.000, 0.000, 0.333, + 0.000, 0.000, 0.500, + 0.000, 0.000, 0.667, + 0.000, 0.000, 0.833, + 0.000, 0.000, 1.000, + 0.143, 0.143, 0.143, + 0.286, 0.286, 0.286, + 0.429, 0.429, 0.429, + 0.571, 0.571, 0.571, + 0.714, 0.714, 0.714, + 0.857, 0.857, 0.857 + ] + ).astype(np.float32) + color_list = color_list.reshape((-1, 3)) * 255 + if not rgb: + color_list = color_list[:, ::-1] + return color_list + + +color_list = colormap() +color_list = color_list.astype('uint8').tolist() + + +def vis_add_mask(image, mask, color, alpha): + color = np.array(color_list[color]) + mask = mask > 0.5 + image[mask] = image[mask] * (1-alpha) + color * alpha + return image.astype('uint8') + +def point_painter(input_image, input_points, point_color=5, point_alpha=0.9, point_radius=15, contour_color=2, contour_width=5): + h, w = input_image.shape[:2] + point_mask = np.zeros((h, w)).astype('uint8') + for point in input_points: + point_mask[point[1], point[0]] = 1 + + kernel = cv2.getStructuringElement(2, (point_radius, point_radius)) + point_mask = cv2.dilate(point_mask, kernel) + + contour_radius = (contour_width - 1) // 2 + dist_transform_fore = cv2.distanceTransform(point_mask, cv2.DIST_L2, 3) + dist_transform_back = cv2.distanceTransform(1-point_mask, cv2.DIST_L2, 3) + dist_map = dist_transform_fore - dist_transform_back + # ...:::!!!:::... + contour_radius += 2 + contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius)) + contour_mask = contour_mask / np.max(contour_mask) + contour_mask[contour_mask>0.5] = 1. + + # paint mask + painted_image = vis_add_mask(input_image.copy(), point_mask, point_color, point_alpha) + # paint contour + painted_image = vis_add_mask(painted_image.copy(), 1-contour_mask, contour_color, 1) + return painted_image + +def mask_painter(input_image, input_mask, mask_color=5, mask_alpha=0.7, contour_color=1, contour_width=3): + assert input_image.shape[:2] == input_mask.shape, 'different shape between image and mask' + # 0: background, 1: foreground + mask = np.clip(input_mask, 0, 1) + contour_radius = (contour_width - 1) // 2 + + dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3) + dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3) + dist_map = dist_transform_fore - dist_transform_back + # ...:::!!!:::... + contour_radius += 2 + contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius)) + contour_mask = contour_mask / np.max(contour_mask) + contour_mask[contour_mask>0.5] = 1. + + # paint mask + painted_image = vis_add_mask(input_image.copy(), mask.copy(), mask_color, mask_alpha) + # paint contour + painted_image = vis_add_mask(painted_image.copy(), 1-contour_mask, contour_color, 1) + + return painted_image + +def background_remover(input_image, input_mask): + """ + input_image: H, W, 3, np.array + input_mask: H, W, np.array + + image_wo_background: PIL.Image + """ + assert input_image.shape[:2] == input_mask.shape, 'different shape between image and mask' + # 0: background, 1: foreground + mask = np.expand_dims(np.clip(input_mask, 0, 1), axis=2)*255 + image_wo_background = np.concatenate([input_image, mask], axis=2) # H, W, 4 + image_wo_background = Image.fromarray(image_wo_background).convert('RGBA') + + return image_wo_background + +if __name__ == '__main__': + input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB')) + input_mask = np.array(Image.open('images/painter_input_mask.jpg').convert('P')) + + # example of mask painter + mask_color = 3 + mask_alpha = 0.7 + contour_color = 1 + contour_width = 5 + + # save + painted_image = Image.fromarray(input_image) + painted_image.save('images/original.png') + + painted_image = mask_painter(input_image, input_mask, mask_color, mask_alpha, contour_color, contour_width) + # save + painted_image = Image.fromarray(input_image) + painted_image.save('images/original1.png') + + # example of point painter + input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB')) + input_points = np.array([[500, 375], [70, 600]]) # x, y + point_color = 5 + point_alpha = 0.9 + point_radius = 15 + contour_color = 2 + contour_width = 5 + painted_image_1 = point_painter(input_image, input_points, point_color, point_alpha, point_radius, contour_color, contour_width) + # save + painted_image = Image.fromarray(painted_image_1) + painted_image.save('images/point_painter_1.png') + + input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB')) + painted_image_2 = point_painter(input_image, input_points, point_color=9, point_radius=20, contour_color=29) + # save + painted_image = Image.fromarray(painted_image_2) + painted_image.save('images/point_painter_2.png') + + # example of background remover + input_image = np.array(Image.open('images/original.png').convert('RGB')) + image_wo_background = background_remover(input_image, input_mask) # return PIL.Image + image_wo_background.save('images/image_wo_background.png') diff --git a/preprocessing/matanyone/utils/__init__.py b/preprocessing/matanyone/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/preprocessing/matanyone/utils/get_default_model.py b/preprocessing/matanyone/utils/get_default_model.py new file mode 100644 index 0000000..c51eae6 --- /dev/null +++ b/preprocessing/matanyone/utils/get_default_model.py @@ -0,0 +1,27 @@ +""" +A helper function to get a default model for quick testing +""" +from omegaconf import open_dict +from hydra import compose, initialize + +import torch +from ..matanyone.model.matanyone import MatAnyone + +def get_matanyone_model(ckpt_path, device=None) -> MatAnyone: + initialize(version_base='1.3.2', config_path="../config", job_name="eval_our_config") + cfg = compose(config_name="eval_matanyone_config") + + with open_dict(cfg): + cfg['weights'] = ckpt_path + + # Load the network weights + if device is not None: + matanyone = MatAnyone(cfg, single_object=True).to(device).eval() + model_weights = torch.load(cfg.weights, map_location=device) + else: # if device is not specified, `.cuda()` by default + matanyone = MatAnyone(cfg, single_object=True).cuda().eval() + model_weights = torch.load(cfg.weights) + + matanyone.load_weights(model_weights) + + return matanyone diff --git a/preprocessing/matanyone/utils/tensor_utils.py b/preprocessing/matanyone/utils/tensor_utils.py new file mode 100644 index 0000000..bb25a45 --- /dev/null +++ b/preprocessing/matanyone/utils/tensor_utils.py @@ -0,0 +1,62 @@ +from typing import List, Iterable +import torch +import torch.nn.functional as F + + +# STM +def pad_divide_by(in_img: torch.Tensor, d: int) -> (torch.Tensor, Iterable[int]): + h, w = in_img.shape[-2:] + + if h % d > 0: + new_h = h + d - h % d + else: + new_h = h + if w % d > 0: + new_w = w + d - w % d + else: + new_w = w + lh, uh = int((new_h - h) / 2), int(new_h - h) - int((new_h - h) / 2) + lw, uw = int((new_w - w) / 2), int(new_w - w) - int((new_w - w) / 2) + pad_array = (int(lw), int(uw), int(lh), int(uh)) + out = F.pad(in_img, pad_array) + return out, pad_array + + +def unpad(img: torch.Tensor, pad: Iterable[int]) -> torch.Tensor: + if len(img.shape) == 4: + if pad[2] + pad[3] > 0: + img = img[:, :, pad[2]:-pad[3], :] + if pad[0] + pad[1] > 0: + img = img[:, :, :, pad[0]:-pad[1]] + elif len(img.shape) == 3: + if pad[2] + pad[3] > 0: + img = img[:, pad[2]:-pad[3], :] + if pad[0] + pad[1] > 0: + img = img[:, :, pad[0]:-pad[1]] + elif len(img.shape) == 5: + if pad[2] + pad[3] > 0: + img = img[:, :, :, pad[2]:-pad[3], :] + if pad[0] + pad[1] > 0: + img = img[:, :, :, :, pad[0]:-pad[1]] + else: + raise NotImplementedError + return img + + +# @torch.jit.script +def aggregate(prob: torch.Tensor, dim: int) -> torch.Tensor: + with torch.amp.autocast("cuda"): + prob = prob.float() + new_prob = torch.cat([torch.prod(1 - prob, dim=dim, keepdim=True), prob], + dim).clamp(1e-7, 1 - 1e-7) + logits = torch.log((new_prob / (1 - new_prob))) # (0, 1) --> (-inf, inf) + + return logits + + +# @torch.jit.script +def cls_to_one_hot(cls_gt: torch.Tensor, num_objects: int) -> torch.Tensor: + # cls_gt: B*1*H*W + B, _, H, W = cls_gt.shape + one_hot = torch.zeros(B, num_objects + 1, H, W, device=cls_gt.device).scatter_(1, cls_gt, 1) + return one_hot \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 70126cd..9eebc8c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,4 +24,6 @@ onnxruntime-gpu rembg[gpu]==2.0.65 matplotlib timm +segment-anything +ffmpeg-python # rembg==2.0.65 \ No newline at end of file diff --git a/wan/modules/model.py b/wan/modules/model.py index e7a76a9..5af4ae8 100644 --- a/wan/modules/model.py +++ b/wan/modules/model.py @@ -482,7 +482,6 @@ class WanAttentionBlock(nn.Module): y *= 1 + e[4] y += e[3] - ffn = self.ffn[0] gelu = self.ffn[1] ffn2= self.ffn[2] @@ -500,8 +499,6 @@ class WanAttentionBlock(nn.Module): x.addcmul_(y, e[5]) - - if hint is not None: if context_scale == 1: x.add_(hint) @@ -539,24 +536,13 @@ class VaceWanAttentionBlock(WanAttentionBlock): c = hints[0] hints[0] = None if self.block_id == 0: - c = self.before_proj(c) + x + c = self.before_proj(c) + c += x c = super().forward(c, **kwargs) c_skip = self.after_proj(c) hints[0] = c return c_skip - # 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): @@ -793,37 +779,6 @@ class WanModel(ModelMixin, ConfigMixin): 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, diff --git a/wan/text2video.py b/wan/text2video.py index ba24e6d..39e7bb3 100644 --- a/wan/text2video.py +++ b/wan/text2video.py @@ -209,8 +209,9 @@ class WanT2V: 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, original_video = False, trim_video= 0): + def prepare_source(self, src_video, src_mask, src_ref_images, num_frames, image_size, device, original_video = False, keep_frames= []): image_sizes = [] + trim_video = len(keep_frames) 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) @@ -237,6 +238,10 @@ class WanT2V: 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) image_sizes.append(src_video[i].shape[2:]) + for k, keep in enumerate(keep_frames): + if not keep: + src_video[i][:, k:k+1] = 0 + src_mask[i][:, k:k+1] = 1 for i, ref_images in enumerate(src_ref_images): if ref_images is not None: diff --git a/wan/utils/utils.py b/wan/utils/utils.py index ce4ecd5..9935ca9 100644 --- a/wan/utils/utils.py +++ b/wan/utils/utils.py @@ -37,12 +37,11 @@ def resample(video_fps, video_frames_count, max_frames, target_fps): break add_frames_count = math.ceil( (target_time -cur_time) / video_frame_duration ) frame_no += add_frames_count + if frame_no >= video_frames_count: + break frame_ids.append(frame_no) cur_time += add_frames_count * video_frame_duration target_time += target_frame_duration - if frame_no >= video_frames_count -1: - break - frame_ids = frame_ids[:video_frames_count] return frame_ids def get_video_frame(file_name, frame_no): diff --git a/wan/utils/vace_preprocessor.py b/wan/utils/vace_preprocessor.py index 3bfe885..0c4edf1 100644 --- a/wan/utils/vace_preprocessor.py +++ b/wan/utils/vace_preprocessor.py @@ -254,7 +254,7 @@ class VaceVideoProcessor(object): if src_video != None: fps = 16 - length = src_video.shape[1] + length = src_video.shape[0] if len(readers) > 0: min_readers = min([len(r) for r in readers]) length = min(length, min_readers ) diff --git a/wgp.py b/wgp.py index da7273a..310b867 100644 --- a/wgp.py +++ b/wgp.py @@ -153,7 +153,7 @@ def process_prompt_and_add_tasks(state, model_choice): if "Vace" in model_filename and "1.3B" in model_filename : resolution_reformated = str(height) + "*" + str(width) if not resolution_reformated in VACE_SIZE_CONFIGS: - res = VACE_SIZE_CONFIGS.keys().join(" and ") + res = (" and ").join(VACE_SIZE_CONFIGS.keys()) gr.Info(f"Video Resolution for Vace model is not supported. Only {res} resolutions are allowed.") return if "I" in video_prompt_type: @@ -175,12 +175,19 @@ def process_prompt_and_add_tasks(state, model_choice): else: video_mask = None if "O" in video_prompt_type : - max_frames= inputs["max_frames"] + keep_frames= inputs["keep_frames"] video_length = inputs["video_length"] - if max_frames ==0: + if len(keep_frames) ==0: gr.Info(f"Warning : you have asked to reuse all the frames of the control Video in the Alternate Video Ending it. Please make sure the number of frames of the control Video is lower than the total number of frames to generate otherwise it won't make a difference.") - elif max_frames >= video_length: - gr.Info(f"The number of frames in the control Video to reuse ({max_frames}) in Alternate Video Ending can not be bigger than the total number of frames ({video_length}) to generate.") + # elif keep_frames >= video_length: + # gr.Info(f"The number of frames in the control Video to reuse ({keep_frames}) in Alternate Video Ending can not be bigger than the total number of frames ({video_length}) to generate.") + # return + elif "V" in video_prompt_type: + keep_frames= inputs["keep_frames"] + video_length = inputs["video_length"] + _, error = parse_keep_frames(keep_frames, video_length) + if len(error) > 0: + gr.Info(f"Invalid Keep Frames property: {error}") return if isinstance(image_refs, list): @@ -1540,8 +1547,8 @@ def download_models(transformer_filename, text_encoder_filename): from huggingface_hub import hf_hub_download, snapshot_download repoId = "DeepBeepMeep/Wan2.1" - sourceFolderList = ["xlm-roberta-large", "pose", "depth", "", ] - fileList = [ [], [],[], ["Wan2.1_VAE_bf16.safetensors", "models_clip_open-clip-xlm-roberta-large-vit-huge-14-bf16.safetensors", "flownet.pkl" ] + computeList(text_encoder_filename) + computeList(transformer_filename) ] + sourceFolderList = ["xlm-roberta-large", "pose", "depth", "mask", "", ] + fileList = [ [], [],[], ["sam_vit_h_4b8939_fp16.safetensors"], ["Wan2.1_VAE_bf16.safetensors", "models_clip_open-clip-xlm-roberta-large-vit-huge-14-bf16.safetensors", "flownet.pkl" ] + computeList(text_encoder_filename) + computeList(transformer_filename) ] targetRoot = "ckpts/" for sourceFolder, files in zip(sourceFolderList,fileList ): if len(files)==0: @@ -1782,25 +1789,6 @@ def get_model_name(model_filename): return model_name -# def generate_header(model_filename, compile, attention_mode): - -# header = "

" - -# model_name = get_model_name(model_filename) - -# header += model_name -# header += " (attention mode: " + (attention_mode if attention_mode!="auto" else "auto/" + get_auto_attention() ) -# if attention_mode not in attention_modes_installed: -# header += " -NOT INSTALLED-" -# elif attention_mode not in attention_modes_supported: -# header += " -NOT SUPPORTED-" - -# if compile: -# header += ", pytorch compilation ON" -# header += ")

" - - -# return header def generate_header(model_filename, compile, attention_mode): @@ -2122,6 +2110,57 @@ def preprocess_video(process_type, height, width, video_in, max_frames): return torch.stack(torch_frames) +def parse_keep_frames(keep_frames, video_length): + def is_integer(n): + try: + float(n) + except ValueError: + return False + else: + return float(n).is_integer() + + def absolute(n): + if n==0: + return 0 + elif n < 0: + return max(0, video_length + n) + else: + return min(n-1, video_length-1) + + if len(keep_frames) == 0: + return [True] *video_length, "" + frames =[False] *video_length + error = "" + sections = keep_frames.split(" ") + for section in sections: + section = section.strip() + if ":" in section: + parts = section.split(":") + if not is_integer(parts[0]): + error =f"Invalid integer {parts[0]}" + break + start_range = absolute(int(parts[0])) + if not is_integer(parts[1]): + error =f"Invalid integer {parts[1]}" + break + end_range = absolute(int(parts[1])) + for i in range(start_range, end_range + 1): + frames[i] = True + else: + if not is_integer(section): + error =f"Invalid integer {section}" + break + index = absolute(int(section)) + frames[index] = True + + if len(error ) > 0: + return [], error + for i in range(len(frames)-1, 0, -1): + if frames[i]: + break + frames= frames[0: i+1] + return frames, error + def generate_video( task_id, progress, @@ -2147,7 +2186,7 @@ def generate_video( image_refs, video_guide, video_mask, - max_frames, + keep_frames, remove_background_image_ref, temporal_upsampling, spatial_upsampling, @@ -2325,12 +2364,16 @@ def generate_video( gen["progress_args"] = progress_args video_guide = preprocess_video(preprocess_type, width=width, height=height,video_in=video_guide, max_frames= video_length) image_refs = image_refs.copy() if image_refs != None else None # required since prepare_source do inplace modifications + keep_frames_parsed, error = parse_keep_frames(keep_frames, video_length) + if len(error) > 0: + raise gr.Error(f"invalid keep frames {keep_frames}") + src_video, src_mask, src_ref_images = wan_model.prepare_source([video_guide], [video_mask], [image_refs], video_length, VACE_SIZE_CONFIGS[resolution_reformated], "cpu", original_video= "O" in video_prompt_type, - trim_video=max_frames) + keep_frames=keep_frames_parsed) else: src_video, src_mask, src_ref_images = None, None, None @@ -2635,6 +2678,8 @@ def process_tasks(state, progress=gr.Progress()): finally: if not ok: queue.clear() + gen["prompts_max"] = 0 + gen["prompt"] = "" yield status queue[:] = [item for item in queue if item['id'] != task['id']] @@ -3014,7 +3059,7 @@ def prepare_inputs_dict(target, inputs ): if not "Vace" in model_filename: - unsaved_params = ["video_prompt_type", "max_frames", "remove_background_image_ref"] + unsaved_params = ["video_prompt_type", "keep_frames", "remove_background_image_ref"] for k in unsaved_params: inputs.pop(k) @@ -3056,7 +3101,7 @@ def save_inputs( image_refs, video_guide, video_mask, - max_frames, + keep_frames, remove_background_image_ref, temporal_upsampling, spatial_upsampling, @@ -3246,6 +3291,13 @@ def refresh_video_prompt_type_video_guide(video_prompt_type, video_prompt_type_v visible = "V" in video_prompt_type return video_prompt_type, gr.update(visible = visible), gr.update(visible = visible), gr.update(visible= "M" in video_prompt_type ) +def refresh_video_prompt_video_guide_trigger(video_prompt_type, video_prompt_type_video_guide): + video_prompt_type_video_guide = video_prompt_type_video_guide.split("#")[0] + video_prompt_type = del_in_sequence(video_prompt_type, "ODPCMV") + video_prompt_type = add_to_sequence(video_prompt_type, video_prompt_type_video_guide) + + return video_prompt_type, video_prompt_type_video_guide, gr.update(visible= "V" in video_prompt_type ), gr.update(visible= "M" in video_prompt_type) , gr.update(visible= "V" in video_prompt_type ) + def generate_video_tab(update_form = False, state_dict = None, ui_defaults = None, model_choice = None, header = None): global inputs_names #, advanced @@ -3365,12 +3417,13 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ("Transfer Depth from the Control Video", "DV"), ("Recolorize the Control Video", "CV"), # ("Alternate Video Ending", "OV"), - ("(adv) Video contains Open Pose, Depth or Black & White ", "V"), - ("(adv) Inpainting of Control Video using Mask Video ", "MV"), + ("Video contains Open Pose, Depth, Black & White, Inpainting ", "V"), + ("Control Video and Mask video for stronger Inpainting ", "MV"), ], value=filter_letters(video_prompt_type_value, "ODPCMV"), label="Video to Video", scale = 3 ) + video_prompt_video_guide_trigger = gr.Text(visible=False, value="") video_prompt_type_image_refs = gr.Dropdown( choices=[ @@ -3384,8 +3437,8 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non # video_prompt_type_image_refs = gr.Checkbox(value="I" in video_prompt_type_value , label= "Use References Images (Faces, Objects) to customize New Video", scale =1 ) video_guide = gr.Video(label= "Control Video", visible= "V" in video_prompt_type_value, value= ui_defaults.get("video_guide", None),) - max_frames = gr.Slider(0, 100, value=ui_defaults.get("max_frames",0), step=1, label="Nb of frames in Control Video to use (0 = max)", visible= "V" in video_prompt_type_value, scale = 2 ) - + # keep_frames = gr.Slider(0, 100, value=ui_defaults.get("keep_frames",0), step=1, label="Nb of frames in Control Video to use (0 = max)", visible= "V" in video_prompt_type_value, scale = 2 ) + keep_frames = gr.Text(value=ui_defaults.get("keep_frames","") , visible= "V" in video_prompt_type_value, scale = 2, label= "Frames to keep in Control Video (empty=All, 1=first, a:b for a range, space to separate values)" ) #, -1=last image_refs = gr.Gallery( label ="Reference Images", type ="pil", show_label= True, columns=[3], rows=[1], object_fit="contain", height="auto", selected_index=0, interactive= True, visible= "I" in video_prompt_type_value, @@ -3798,9 +3851,9 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non target_settings = gr.Text(value = "settings", interactive= False, visible= False) image_prompt_type.change(fn=refresh_image_prompt_type, inputs=[state, image_prompt_type], outputs=[image_start, image_end]) - # video_prompt_type.change(fn=refresh_video_prompt_type, inputs=[state, video_prompt_type], outputs=[image_refs, video_guide, video_mask, max_frames, remove_background_image_ref]) + video_prompt_video_guide_trigger.change(fn=refresh_video_prompt_video_guide_trigger, inputs=[video_prompt_type, video_prompt_video_guide_trigger], outputs=[video_prompt_type, video_prompt_type_video_guide, video_guide, video_mask, keep_frames]) video_prompt_type_image_refs.input(fn=refresh_video_prompt_type_image_refs, inputs = [video_prompt_type, video_prompt_type_image_refs], outputs = [video_prompt_type, image_refs, remove_background_image_ref ]) - video_prompt_type_video_guide.input(fn=refresh_video_prompt_type_video_guide, inputs = [video_prompt_type, video_prompt_type_video_guide], outputs = [video_prompt_type, video_guide, max_frames, video_mask]) + video_prompt_type_video_guide.input(fn=refresh_video_prompt_type_video_guide, inputs = [video_prompt_type, video_prompt_type_video_guide], outputs = [video_prompt_type, video_guide, keep_frames, video_mask]) 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]) @@ -3903,12 +3956,8 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non return ( loras_choices, lset_name, state, queue_df, current_gen_column, gen_status, output, abort_btn, generate_btn, add_to_queue_btn, - gen_info, - prompt, wizard_prompt, wizard_prompt_activated_var, wizard_variables_var, - prompt_column_advanced, prompt_column_wizard, prompt_column_wizard_vars, - advanced_row, image_prompt_column, video_prompt_column, queue_accordion, - *prompt_vars - ) + gen_info, queue_accordion, video_guide, video_mask, video_prompt_video_guide_trigger + ) def generate_download_tab(lset_name,loras_choices, state): @@ -4132,8 +4181,30 @@ def generate_dropdown_model_list(): ) +def select_tab(tab_state, evt:gr.SelectData): + tab_video_mask_creator = 2 + + old_tab_no = tab_state.get("tab_no",0) + new_tab_no = evt.index + if old_tab_no == tab_video_mask_creator: + vmc_event_handler(False) + elif new_tab_no == tab_video_mask_creator: + if gen_in_progress: + gr.Info("Unable to access this Tab while a Generation is in Progress. Please come back later") + tab_state["tab_auto"]=old_tab_no + else: + vmc_event_handler(True) + tab_state["tab_no"] = new_tab_no +def select_tab_auto(tab_state): + old_tab_no = tab_state.pop("tab_auto", -1) + if old_tab_no>= 0: + tab_state["tab_auto"]=old_tab_no + return gr.Tabs(selected=old_tab_no) # !! doesnt work !! + return gr.Tab() + def create_demo(): + global vmc_event_handler css = """ #model_list{ background-color:black; @@ -4370,6 +4441,8 @@ def create_demo(): gr.Markdown("

WanGP v4.0 by DeepBeepMeep ") # (Updates)

") global model_list + tab_state = gr.State({ "tab_no":0 }) + with gr.Tabs(selected="video_gen", ) as main_tabs: with gr.Tab("Video Generator", id="video_gen") as t2v_tab: with gr.Row(): @@ -4386,14 +4459,15 @@ def create_demo(): ( loras_choices, lset_name, state, queue_df, current_gen_column, gen_status, output, abort_btn, generate_btn, add_to_queue_btn, - gen_info, - prompt, wizard_prompt, wizard_prompt_activated_var, wizard_variables_var, - prompt_column_advanced, prompt_column_wizard, prompt_column_wizard_vars, - advanced_row, image_prompt_column, video_prompt_column, queue_accordion, - *prompt_vars_outputs + gen_info, queue_accordion, video_guide, video_mask, video_prompt_type_video_trigger ) = generate_video_tab(model_choice=model_choice, header=header) with gr.Tab("Informations"): generate_info_tab() + with gr.Tab("Video Mask Creator", id="video_mask_creator") as video_mask_creator: + from preprocessing.matanyone import app as matanyone_app + vmc_event_handler = matanyone_app.get_vmc_event_handler() + + matanyone_app.display(video_guide, video_mask, video_prompt_type_video_trigger) if not args.lock_config: with gr.Tab("Downloads", id="downloads") as downloads_tab: generate_download_tab(lset_name, loras_choices, state) @@ -4420,6 +4494,7 @@ def create_demo(): trigger_mode="always_last" ) + main_tabs.select(fn=select_tab, inputs= [tab_state], outputs= None).then(fn=select_tab_auto, inputs= [tab_state], outputs=[main_tabs]) return demo if __name__ == "__main__": From 4c79c62419a445a192f4bee4285e1989c3030ab1 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Sat, 12 Apr 2025 04:10:21 +0200 Subject: [PATCH 59/69] small fixes --- preprocessing/matanyone/app.py | 8 ++++++-- requirements.txt | 2 ++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/preprocessing/matanyone/app.py b/preprocessing/matanyone/app.py index d302f03..18f177e 100644 --- a/preprocessing/matanyone/app.py +++ b/preprocessing/matanyone/app.py @@ -311,9 +311,13 @@ def video_matting(video_state, end_slider, interactive_state, mask_dropdown, ero output_frames.append(output_frame) foreground = output_frames - foreground_output = save_video(foreground, output_path="./results/{}_fg.mp4".format(video_state["video_name"]), fps=fps) + if not os.path.exists("mask_outputs"): + os.makedirs("mask_outputs") + + + foreground_output = save_video(foreground, output_path="./mask_outputs/{}_fg.mp4".format(video_state["video_name"]), fps=fps) # foreground_output = generate_video_from_frames(foreground, output_path="./results/{}_fg.mp4".format(video_state["video_name"]), fps=fps, audio_path=audio_path) # import video_input to name the output video - alpha_output = save_video(alpha, output_path="./results/{}_alpha.mp4".format(video_state["video_name"]), fps=fps) + alpha_output = save_video(alpha, output_path="./mask_outputs/{}_alpha.mp4".format(video_state["video_name"]), fps=fps) # alpha_output = generate_video_from_frames(alpha, output_path="./results/{}_alpha.mp4".format(video_state["video_name"]), fps=fps, gray2rgb=True, audio_path=audio_path) # import video_input to name the output video return foreground_output, alpha_output diff --git a/requirements.txt b/requirements.txt index 9eebc8c..0b90776 100644 --- a/requirements.txt +++ b/requirements.txt @@ -26,4 +26,6 @@ matplotlib timm segment-anything ffmpeg-python +omegaconf +hydra-core # rembg==2.0.65 \ No newline at end of file From 033546ca429061b4ee7f50318191ed3501f42ec6 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Sun, 13 Apr 2025 01:36:57 +0200 Subject: [PATCH 60/69] Added Vace Sliding Window --- README.md | 6 +- preprocessing/matanyone/app.py | 65 +++++-- wan/text2video.py | 37 ++-- wan/utils/utils.py | 11 +- wan/utils/vace_preprocessor.py | 16 +- wgp.py | 337 +++++++++++++++++++++------------ 6 files changed, 304 insertions(+), 168 deletions(-) diff --git a/README.md b/README.md index fa9d69f..9b7ccbf 100644 --- a/README.md +++ b/README.md @@ -14,10 +14,14 @@ ## 🔥 Latest News!! -* April 9 2025: 👋 Wan 2.1GP v4.0: lots of goodies for you ! +* April 13 2025: 👋 Wan 2.1GP v4.0: lots of goodies for you ! - A new queuing system that lets you stack in a queue as many text2video and imag2video tasks as you want. Each task can rely on complete different generation parameters (different number of frames, steps, loras, ...). - Temporal upsampling (Rife) and spatial upsampling (Lanczos) for a smoother video (32 fps or 64 fps) and to enlarge you video by x2 or x4. Check these new advanced options. - Wan Vace Control Net support : with Vace you can inject in the scene people or objects, animate a person, perform inpainting or outpainting, continue a video, ... I have provided an introduction guide below. + - Integrated *Matanyone* tool directly inside WanGP so that you can create easily inpainting masks + - Sliding Window generation for Vace, create windows that can last dozen of seconds + - A new UI, tabs were replaced by a Dropdown box to easily switch models + * Mar 27 2025: 👋 Added support for the new Wan Fun InP models (image2video). The 14B Fun InP has probably better end image support but unfortunately existing loras do not work so well with it. The great novelty is the Fun InP image2 1.3B model : Image 2 Video is now accessible to even lower hardware configuration. It is not as good as the 14B models but very impressive for its size. You can choose any of those models in the Configuration tab. Many thanks to the VideoX-Fun team (https://github.com/aigc-apps/VideoX-Fun) * Mar 26 2025: 👋 Good news ! Official support for RTX 50xx please check the installation instructions below. * Mar 24 2025: 👋 Wan2.1GP v3.2: diff --git a/preprocessing/matanyone/app.py b/preprocessing/matanyone/app.py index 18f177e..a58fa67 100644 --- a/preprocessing/matanyone/app.py +++ b/preprocessing/matanyone/app.py @@ -163,10 +163,10 @@ def get_frames_from_video(video_input, video_state): model.samcontroler.sam_controler.set_image(video_state["origin_images"][0]) return video_state, video_info, video_state["origin_images"][0], \ gr.update(visible=True, maximum=len(frames), value=1), gr.update(visible=True, maximum=len(frames), value=len(frames)), gr.update(visible=False, maximum=len(frames), value=len(frames)), \ - gr.update(visible=True), gr.update(visible=True), \ + gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), \ gr.update(visible=True), gr.update(visible=True),\ - gr.update(visible=True), gr.update(visible=True), \ gr.update(visible=True), gr.update(visible=False), \ + gr.update(visible=False), gr.update(visible=False), \ gr.update(visible=False), gr.update(visible=True), \ gr.update(visible=True) @@ -273,7 +273,7 @@ def save_video(frames, output_path, fps): return output_path # video matting -def video_matting(video_state, end_slider, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size): +def video_matting(video_state, end_slider, matting_type, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size): matanyone_processor = InferenceCore(matanyone_model, cfg=matanyone_model.cfg) # if interactive_state["track_end_number"]: # following_frames = video_state["origin_images"][video_state["select_frame_number"]:interactive_state["track_end_number"]] @@ -301,9 +301,16 @@ def video_matting(video_state, end_slider, interactive_state, mask_dropdown, ero template_mask[0][0]=1 foreground, alpha = matanyone(matanyone_processor, following_frames, template_mask*255, r_erode=erode_kernel_size, r_dilate=dilate_kernel_size) output_frames = [] + foreground_mat = matting_type == "Foreground" for frame_origin, frame_alpha in zip(following_frames, alpha): - frame_alpha[frame_alpha > 127] = 255 - frame_alpha[frame_alpha <= 127] = 0 + if foreground_mat: + frame_alpha[frame_alpha > 127] = 255 + frame_alpha[frame_alpha <= 127] = 0 + else: + frame_temp = frame_alpha.copy() + frame_alpha[frame_temp > 127] = 0 + frame_alpha[frame_temp <= 127] = 255 + output_frame = np.bitwise_and(frame_origin, 255-frame_alpha) frame_grey = frame_alpha.copy() frame_grey[frame_alpha == 255] = 127 @@ -314,15 +321,19 @@ def video_matting(video_state, end_slider, interactive_state, mask_dropdown, ero if not os.path.exists("mask_outputs"): os.makedirs("mask_outputs") - - foreground_output = save_video(foreground, output_path="./mask_outputs/{}_fg.mp4".format(video_state["video_name"]), fps=fps) + file_name= video_state["video_name"] + file_name = ".".join(file_name.split(".")[:-1]) + foreground_output = save_video(foreground, output_path="./mask_outputs/{}_fg.mp4".format(file_name), fps=fps) # foreground_output = generate_video_from_frames(foreground, output_path="./results/{}_fg.mp4".format(video_state["video_name"]), fps=fps, audio_path=audio_path) # import video_input to name the output video - alpha_output = save_video(alpha, output_path="./mask_outputs/{}_alpha.mp4".format(video_state["video_name"]), fps=fps) + alpha_output = save_video(alpha, output_path="./mask_outputs/{}_alpha.mp4".format(file_name), fps=fps) # alpha_output = generate_video_from_frames(alpha, output_path="./results/{}_alpha.mp4".format(video_state["video_name"]), fps=fps, gray2rgb=True, audio_path=audio_path) # import video_input to name the output video - return foreground_output, alpha_output + return foreground_output, alpha_output, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True) +def show_outputs(): + return gr.update(visible=True), gr.update(visible=True) + def add_audio_to_video(video_path, audio_path, output_path): try: video_input = ffmpeg.input(video_path) @@ -392,8 +403,8 @@ def restart(): }, "track_end_number": None, }, [[],[]], None, None, \ - gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),\ - gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), \ + gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),\ + gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), \ gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), \ gr.update(visible=False), gr.update(visible=False, choices=[], value=[]), "", gr.update(visible=False) @@ -529,7 +540,16 @@ def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger) visible=False, min_width=100, scale=1) - mask_dropdown = gr.Dropdown(multiselect=True, value=[], label="Mask Selection", info="Choose 1~all mask(s) added in Step 2", visible=False) + matting_type = gr.Radio( + choices=["Foreground", "Background"], + value="Foreground", + label="Matting Type", + info="Type of Video Matting to Generate", + interactive=True, + visible=False, + min_width=100, + scale=1) + mask_dropdown = gr.Dropdown(multiselect=True, value=[], label="Mask Selection", info="Choose 1~all mask(s) added in Step 2", visible=False, scale=2) gr.Markdown("---") @@ -549,9 +569,9 @@ def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger) template_frame = gr.Image(label="Start Frame", type="pil",interactive=True, elem_id="template_frame", visible=False, elem_classes="image") with gr.Row(): clear_button_click = gr.Button(value="Clear Clicks", interactive=True, visible=False, min_width=100) - add_mask_button = gr.Button(value="Add Mask", interactive=True, visible=False, min_width=100) + add_mask_button = gr.Button(value="Set Mask", interactive=True, visible=False, min_width=100) remove_mask_button = gr.Button(value="Remove Mask", interactive=True, visible=False, min_width=100) # no use - matting_button = gr.Button(value="Video Matting", interactive=True, visible=False, min_width=100) + matting_button = gr.Button(value="Generate Video Matting", interactive=True, visible=False, min_width=100) with gr.Row(): gr.Markdown("") @@ -560,11 +580,11 @@ def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger) with gr.Column(scale=2): foreground_video_output = gr.Video(label="Masked Video Output", visible=False, elem_classes="video") foreground_output_button = gr.Button(value="Black & White Video Output", visible=False, elem_classes="new_button") - export_to_vace_video_input_btn = gr.Button("Export to Vace Video Input Video For Inpainting") + export_to_vace_video_input_btn = gr.Button("Export to Vace Video Input Video For Inpainting", visible= False) with gr.Column(scale=2): alpha_video_output = gr.Video(label="B & W Mask Video Output", visible=False, elem_classes="video") alpha_output_button = gr.Button(value="Alpha Mask Output", visible=False, elem_classes="new_button") - export_to_vace_video_mask_btn = gr.Button("Export to Vace Video Input and Video Mask for stronger Inpainting") + export_to_vace_video_mask_btn = gr.Button("Export to Vace Video Input and Video Mask for stronger Inpainting", visible= False) export_to_vace_video_input_btn.click(fn=export_to_vace_video_input, inputs= [foreground_video_output], outputs= [video_prompt_video_guide_trigger, vace_video_input]) export_to_vace_video_mask_btn.click(fn=export_to_vace_video_mask, inputs= [foreground_video_output, alpha_video_output], outputs= [video_prompt_video_guide_trigger, vace_video_input, vace_video_mask]) @@ -575,7 +595,7 @@ def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger) video_input, video_state ], outputs=[video_state, video_info, template_frame, - image_selection_slider, end_selection_slider, track_pause_number_slider, point_prompt, clear_button_click, add_mask_button, matting_button, template_frame, + image_selection_slider, end_selection_slider, track_pause_number_slider, point_prompt, matting_type, clear_button_click, add_mask_button, matting_button, template_frame, foreground_video_output, alpha_video_output, foreground_output_button, alpha_output_button, mask_dropdown, step2_title] ) @@ -609,9 +629,12 @@ def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger) # video matting matting_button.click( + fn=show_outputs, + inputs=[], + outputs=[foreground_video_output, alpha_video_output]).then( fn=video_matting, - inputs=[video_state, end_selection_slider, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size], - outputs=[foreground_video_output, alpha_video_output] + inputs=[video_state, end_selection_slider, matting_type, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size], + outputs=[foreground_video_output, alpha_video_output,foreground_video_output, alpha_video_output, export_to_vace_video_input_btn, export_to_vace_video_mask_btn] ) # click to get mask @@ -631,7 +654,7 @@ def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger) click_state, foreground_video_output, alpha_video_output, template_frame, - image_selection_slider , track_pause_number_slider,point_prompt, clear_button_click, + image_selection_slider, end_selection_slider, track_pause_number_slider,point_prompt, export_to_vace_video_input_btn, export_to_vace_video_mask_btn, matting_type, clear_button_click, add_mask_button, matting_button, template_frame, foreground_video_output, alpha_video_output, remove_mask_button, foreground_output_button, alpha_output_button, mask_dropdown, video_info, step2_title ], queue=False, @@ -646,7 +669,7 @@ def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger) click_state, foreground_video_output, alpha_video_output, template_frame, - image_selection_slider , track_pause_number_slider,point_prompt, clear_button_click, + image_selection_slider , end_selection_slider, track_pause_number_slider,point_prompt, export_to_vace_video_input_btn, export_to_vace_video_mask_btn, matting_type, clear_button_click, add_mask_button, matting_button, template_frame, foreground_video_output, alpha_video_output, remove_mask_button, foreground_output_button, alpha_output_button, mask_dropdown, video_info, step2_title ], queue=False, diff --git a/wan/text2video.py b/wan/text2video.py index 39e7bb3..035700a 100644 --- a/wan/text2video.py +++ b/wan/text2video.py @@ -209,34 +209,47 @@ class WanT2V: 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, original_video = False, keep_frames= []): + def prepare_source(self, src_video, src_mask, src_ref_images, total_frames, image_size, device, original_video = False, keep_frames= [], start_frame = 0, pre_src_video = None): image_sizes = [] trim_video = len(keep_frames) - for i, (sub_src_video, sub_src_mask) in enumerate(zip(src_video, src_mask)): + + for i, (sub_src_video, sub_src_mask, sub_pre_src_video) in enumerate(zip(src_video, src_mask,pre_src_video)): + prepend_count = 0 if sub_pre_src_video == None else sub_pre_src_video.shape[1] + num_frames = total_frames - prepend_count if sub_src_mask is not None and sub_src_video is not None: - src_video[i], src_mask[i], _, _, _ = self.vid_proc.load_video_pair(sub_src_video, sub_src_mask, max_frames= num_frames, trim_video = trim_video) + src_video[i], src_mask[i], _, _, _ = self.vid_proc.load_video_pair(sub_src_video, sub_src_mask, max_frames= num_frames, trim_video = trim_video - prepend_count, start_frame = start_frame) # src_video is [-1, 1], 0 = inpainting area (in fact 127 in [0, 255]) # src_mask is [-1, 1], 0 = preserve original video (in fact 127 in [0, 255]) and 1 = Inpainting (in fact 255 in [0, 255]) src_video[i] = src_video[i].to(device) src_mask[i] = src_mask[i].to(device) + if prepend_count > 0: + src_video[i] = torch.cat( [sub_pre_src_video, src_video[i]], dim=1) + src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), src_mask[i]] ,1) src_video_shape = src_video[i].shape - if src_video_shape[1] != 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) + if src_video_shape[1] != total_frames: + src_video[i] = torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) + src_mask[i] = torch.cat( [src_mask[i], src_mask[i].new_ones(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) src_mask[i] = torch.clamp((src_mask[i][:1, :, :, :] + 1) / 2, min=0, max=1) image_sizes.append(src_video[i].shape[2:]) elif sub_src_video is None: - 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) + if prepend_count > 0: + src_video[i] = torch.cat( [sub_pre_src_video, torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device)], dim=1) + src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), torch.ones((3, num_frames, image_size[0], image_size[1]), device=device)] ,1) + else: + src_video[i] = torch.zeros((3, num_frames, image_size[0], image_size[1]), device=device) + src_mask[i] = torch.ones_like(src_video[i], device=device) image_sizes.append(image_size) else: - src_video[i], _, _, _ = self.vid_proc.load_video(sub_src_video, max_frames= num_frames, trim_video = trim_video) + src_video[i], _, _, _ = self.vid_proc.load_video(sub_src_video, max_frames= num_frames, trim_video = trim_video - prepend_count, start_frame = start_frame) src_video[i] = src_video[i].to(device) src_mask[i] = torch.zeros_like(src_video[i], device=device) if original_video else torch.ones_like(src_video[i], device=device) + if prepend_count > 0: + src_video[i] = torch.cat( [sub_pre_src_video, src_video[i]], dim=1) + src_mask[i] = torch.cat( [torch.zeros_like(sub_pre_src_video), src_mask[i]] ,1) src_video_shape = src_video[i].shape - if src_video_shape[1] != 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) + if src_video_shape[1] != total_frames: + src_video[i] = torch.cat( [src_video[i], src_video[i].new_zeros(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) + src_mask[i] = torch.cat( [src_mask[i], src_mask[i].new_ones(src_video_shape[0], total_frames -src_video_shape[1], *src_video_shape[-2:])], dim=1) image_sizes.append(src_video[i].shape[2:]) for k, keep in enumerate(keep_frames): if not keep: diff --git a/wan/utils/utils.py b/wan/utils/utils.py index 9935ca9..5149464 100644 --- a/wan/utils/utils.py +++ b/wan/utils/utils.py @@ -22,18 +22,18 @@ __all__ = ['cache_video', 'cache_image', 'str2bool'] from PIL import Image -def resample(video_fps, video_frames_count, max_frames, target_fps): +def resample(video_fps, video_frames_count, max_target_frames_count, target_fps, start_target_frame ): import math video_frame_duration = 1 /video_fps target_frame_duration = 1 / target_fps - cur_time = 0 - target_time = 0 - frame_no = 0 + target_time = start_target_frame * target_frame_duration + frame_no = math.ceil(target_time / video_frame_duration) + cur_time = frame_no * video_frame_duration frame_ids =[] while True: - if max_frames != 0 and len(frame_ids) >= max_frames: + if max_target_frames_count != 0 and len(frame_ids) >= max_target_frames_count : break add_frames_count = math.ceil( (target_time -cur_time) / video_frame_duration ) frame_no += add_frames_count @@ -42,6 +42,7 @@ def resample(video_fps, video_frames_count, max_frames, target_fps): frame_ids.append(frame_no) cur_time += add_frames_count * video_frame_duration target_time += target_frame_duration + frame_ids = frame_ids[:max_target_frames_count] return frame_ids def get_video_frame(file_name, frame_no): diff --git a/wan/utils/vace_preprocessor.py b/wan/utils/vace_preprocessor.py index 0c4edf1..c591cca 100644 --- a/wan/utils/vace_preprocessor.py +++ b/wan/utils/vace_preprocessor.py @@ -182,14 +182,14 @@ class VaceVideoProcessor(object): - def _get_frameid_bbox_adjust_last(self, fps, video_frames_count, h, w, crop_box, rng, max_frames= 0): + def _get_frameid_bbox_adjust_last(self, fps, video_frames_count, h, w, crop_box, rng, max_frames= 0, start_frame =0): from wan.utils.utils import resample target_fps = self.max_fps # video_frames_count = len(frame_timestamps) - frame_ids= resample(fps, video_frames_count, max_frames, target_fps) + frame_ids= resample(fps, video_frames_count, max_frames, target_fps, start_frame ) x1, x2, y1, y2 = [0, w, 0, h] if crop_box is None else crop_box h, w = y2 - y1, x2 - x1 @@ -206,7 +206,7 @@ class VaceVideoProcessor(object): np.log2(np.sqrt(max_area_z)) ))) - seq_len = max_area_z * ((max_frames- 1) // df +1) + seq_len = max_area_z * ((max_frames- start_frame - 1) // df +1) # of = min( # (len(frame_ids) - 1) // df + 1, @@ -226,9 +226,9 @@ class VaceVideoProcessor(object): return frame_ids, (x1, x2, y1, y2), (oh, ow), target_fps - def _get_frameid_bbox(self, fps, video_frames_count, h, w, crop_box, rng, max_frames= 0): + def _get_frameid_bbox(self, fps, video_frames_count, h, w, crop_box, rng, max_frames= 0, start_frame= 0): if self.keep_last: - return self._get_frameid_bbox_adjust_last(fps, video_frames_count, h, w, crop_box, rng, max_frames= max_frames) + return self._get_frameid_bbox_adjust_last(fps, video_frames_count, h, w, crop_box, rng, max_frames= max_frames, start_frame= start_frame) else: return self._get_frameid_bbox_default(fps, video_frames_count, h, w, crop_box, rng, max_frames= max_frames) @@ -238,7 +238,7 @@ class VaceVideoProcessor(object): 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): + def load_video_batch(self, *data_key_batch, crop_box=None, seed=2024, max_frames= 0, trim_video =0, start_frame = 0, **kwargs): rng = np.random.default_rng(seed + hash(data_key_batch[0]) % 10000) # read video import decord @@ -254,7 +254,7 @@ class VaceVideoProcessor(object): if src_video != None: fps = 16 - length = src_video.shape[0] + length = src_video.shape[0] + start_frame if len(readers) > 0: min_readers = min([len(r) for r in readers]) length = min(length, min_readers ) @@ -269,7 +269,7 @@ class VaceVideoProcessor(object): h, w = src_video.shape[1:3] else: h, w = readers[0].next().shape[:2] - frame_ids, (x1, x2, y1, y2), (oh, ow), fps = self._get_frameid_bbox(fps, length, h, w, crop_box, rng, max_frames=max_frames) + frame_ids, (x1, x2, y1, y2), (oh, ow), fps = self._get_frameid_bbox(fps, length, h, w, crop_box, rng, max_frames=max_frames, start_frame = start_frame ) # preprocess video videos = [reader.get_batch(frame_ids)[:, y1:y2, x1:x2, :] for reader in readers] diff --git a/wgp.py b/wgp.py index 310b867..283aeb8 100644 --- a/wgp.py +++ b/wgp.py @@ -144,13 +144,23 @@ def process_prompt_and_add_tasks(state, model_choice): gr.Info("You must use the 14B model to generate videos with a resolution equivalent to 720P") return - + sliding_window_repeat = inputs["sliding_window_repeat"] + sliding_window = sliding_window_repeat > 0 if "Vace" in model_filename: video_prompt_type = inputs["video_prompt_type"] image_refs = inputs["image_refs"] video_guide = inputs["video_guide"] video_mask = inputs["video_mask"] - if "Vace" in model_filename and "1.3B" in model_filename : + + if sliding_window: + if inputs["repeat_generation"]!=1: + gr.Info("Only one Video generated per Prompt is supported when Sliding windows is used") + return + if inputs["sliding_window_overlap"]>=inputs["video_length"] : + gr.Info("The number of frames of the Sliding Window Overlap must be less than the Number of Frames to Generate") + return + + if "1.3B" in model_filename : resolution_reformated = str(height) + "*" + str(width) if not resolution_reformated in VACE_SIZE_CONFIGS: res = (" and ").join(VACE_SIZE_CONFIGS.keys()) @@ -197,6 +207,9 @@ def process_prompt_and_add_tasks(state, model_choice): image_refs = resize_and_remove_background(image_refs, width, height, inputs["remove_background_image_ref"] ==1) + if sliding_window and len(prompts) > 0: + prompts = ["\n".join(prompts)] + for single_prompt in prompts: extra_inputs = { "prompt" : single_prompt, @@ -2053,7 +2066,7 @@ def convert_image(image): return cast(Image, ImageOps.exif_transpose(image)) -def preprocess_video(process_type, height, width, video_in, max_frames): +def preprocess_video(process_type, height, width, video_in, max_frames, start_frame=0): from wan.utils.utils import resample @@ -2063,8 +2076,10 @@ def preprocess_video(process_type, height, width, video_in, max_frames): fps = reader.get_avg_fps() - frame_nos = resample(fps, len(reader), max_frames= max_frames, target_fps=16) + frame_nos = resample(fps, len(reader), max_target_frames_count= max_frames, target_fps=16, start_target_frame= start_frame) frames_list = reader.get_batch(frame_nos) + if len(frames_list) == 0: + return None frame_height, frame_width, _ = frames_list[0].shape scale = ((height * width ) / (frame_height * frame_width))**(1/2) @@ -2187,6 +2202,9 @@ def generate_video( video_guide, video_mask, keep_frames, + sliding_window_repeat, + sliding_window_overlap, + sliding_window_discard_last_frames, remove_background_image_ref, temporal_upsampling, spatial_upsampling, @@ -2342,41 +2360,6 @@ def generate_video( else: raise gr.Error("Teacache not supported for this model") - if "Vace" in model_filename: - # video_prompt_type = video_prompt_type +"G" - if any(process in video_prompt_type for process in ("P", "D", "G")) : - prompts_max = gen["prompts_max"] - - status = get_generation_status(prompt_no, prompts_max, 1, 1) - preprocess_type = None - if "P" in video_prompt_type : - progress_args = [0, status + " - Extracting Open Pose Information"] - preprocess_type = "pose" - elif "D" in video_prompt_type : - progress_args = [0, status + " - Extracting Depth Information"] - preprocess_type = "depth" - elif "G" in video_prompt_type : - progress_args = [0, status + " - Extracting Gray Level Information"] - preprocess_type = "gray" - - if preprocess_type != None : - progress(*progress_args ) - gen["progress_args"] = progress_args - video_guide = preprocess_video(preprocess_type, width=width, height=height,video_in=video_guide, max_frames= video_length) - image_refs = image_refs.copy() if image_refs != None else None # required since prepare_source do inplace modifications - keep_frames_parsed, error = parse_keep_frames(keep_frames, video_length) - if len(error) > 0: - raise gr.Error(f"invalid keep frames {keep_frames}") - - src_video, src_mask, src_ref_images = wan_model.prepare_source([video_guide], - [video_mask], - [image_refs], - video_length, VACE_SIZE_CONFIGS[resolution_reformated], "cpu", - original_video= "O" in video_prompt_type, - keep_frames=keep_frames_parsed) - else: - src_video, src_mask, src_ref_images = None, None, None - import random if seed == None or seed <0: @@ -2393,6 +2376,21 @@ def generate_video( gen["prompt"] = prompt repeat_no = 0 extra_generation = 0 + sliding_window = sliding_window_repeat > 0 + if sliding_window: + start_frame = 0 + reuse_frames = sliding_window_overlap + discard_last_frames = sliding_window_discard_last_frames #4 + repeat_generation = sliding_window_repeat + prompts = prompt.split("\n") + prompts = [part for part in prompts if len(prompt)>0] + + + gen["sliding_window"] = sliding_window + + frames_already_processed = None + pre_video_guide = None + while True: extra_generation += gen.get("extra_orders",0) gen["extra_orders"] = 0 @@ -2400,10 +2398,59 @@ def generate_video( gen["total_generation"] = total_generation if abort or repeat_no >= total_generation: break + + if "Vace" in model_filename and (repeat_no == 0 or sliding_window): + if sliding_window: + prompt = prompts[repeat_no] if repeat_no < len(prompts) else prompts[-1] + + # video_prompt_type = video_prompt_type +"G" + image_refs_copy = image_refs.copy() if image_refs != None else None # required since prepare_source do inplace modifications + video_guide_copy = video_guide + video_mask_copy = video_mask + if any(process in video_prompt_type for process in ("P", "D", "G")) : + prompts_max = gen["prompts_max"] + + status = get_generation_status(prompt_no, prompts_max, 1, 1, sliding_window) + preprocess_type = None + if "P" in video_prompt_type : + progress_args = [0, status + " - Extracting Open Pose Information"] + preprocess_type = "pose" + elif "D" in video_prompt_type : + progress_args = [0, status + " - Extracting Depth Information"] + preprocess_type = "depth" + elif "G" in video_prompt_type : + progress_args = [0, status + " - Extracting Gray Level Information"] + preprocess_type = "gray" + + if preprocess_type != None : + progress(*progress_args ) + gen["progress_args"] = progress_args + video_guide_copy = preprocess_video(preprocess_type, width=width, height=height,video_in=video_guide, max_frames= video_length if repeat_no ==0 else video_length - reuse_frames, start_frame = start_frame) + keep_frames_parsed, error = parse_keep_frames(keep_frames, video_length) + if len(error) > 0: + raise gr.Error(f"invalid keep frames {keep_frames}") + if repeat_no == 0: + image_size = VACE_SIZE_CONFIGS[resolution_reformated] # default frame dimensions until it is set by video_src (if there is any) + src_video, src_mask, src_ref_images = wan_model.prepare_source([video_guide_copy], + [video_mask_copy ], + [image_refs_copy], + video_length, image_size = image_size, device ="cpu", + original_video= "O" in video_prompt_type, + keep_frames=keep_frames_parsed, + start_frame = start_frame, + pre_src_video = [pre_video_guide] + ) + if repeat_no == 0 and src_video != None and len(src_video) > 0: + image_size = src_video[0].shape[-2:] + + else: + src_video, src_mask, src_ref_images = None, None, None + + repeat_no +=1 gen["repeat_no"] = repeat_no prompts_max = gen["prompts_max"] - status = get_generation_status(prompt_no, prompts_max, repeat_no, total_generation) + status = get_generation_status(prompt_no, prompts_max, repeat_no, total_generation, sliding_window) yield status @@ -2539,6 +2586,15 @@ def generate_video( # yield f"Video generation was aborted. Total Generation Time: {end_time-start_time:.1f}s" else: sample = samples.cpu() + if sliding_window : + start_frame += video_length + if discard_last_frames > 0: + sample = sample[: , :-discard_last_frames] + start_frame -= discard_last_frames + pre_video_guide = sample[:, -reuse_frames:] + if repeat_no > 1: + sample = sample[: , reuse_frames:] + start_frame -= reuse_frames time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%Hh%Mm%Ss") if os.name == 'nt': @@ -2565,7 +2621,13 @@ def generate_video( if exp > 0: from rife.inference import temporal_interpolation - sample = temporal_interpolation( os.path.join("ckpts", "flownet.pkl"), sample, exp, device=processing_device) + if sliding_window and repeat_no > 1: + sample = torch.cat([frames_already_processed[:, -2:-1], sample], dim=1) + sample = temporal_interpolation( os.path.join("ckpts", "flownet.pkl"), sample, exp, device=processing_device) + sample = sample[:, 1:] + else: + sample = temporal_interpolation( os.path.join("ckpts", "flownet.pkl"), sample, exp, device=processing_device) + fps = fps * 2**exp if len(spatial_upsampling) > 0: @@ -2590,6 +2652,12 @@ def generate_video( new_frames = None sample = sample * 2 - 1 + if sliding_window : + if repeat_no == 1: + frames_already_processed = sample + else: + sample = torch.cat([frames_already_processed, sample], dim=1) + frames_already_processed = sample cache_video( tensor=sample[None], @@ -2616,7 +2684,8 @@ def generate_video( print(f"New video saved to Path: "+video_path) file_list.append(video_path) state['update_gallery'] = True - seed += 1 + if not sliding_window: + seed += 1 if temp_filename!= None and os.path.isfile(temp_filename): os.remove(temp_filename) @@ -2694,17 +2763,19 @@ def process_tasks(state, progress=gr.Progress()): yield f"Total Generation Time: {end_time-start_time:.1f}s" -def get_generation_status(prompt_no, prompts_max, repeat_no, repeat_max): - if prompts_max == 1: +def get_generation_status(prompt_no, prompts_max, repeat_no, repeat_max, sliding_window): + + item = "Sliding Window" if sliding_window else "Sample" + if prompts_max == 1: if repeat_max == 1: return "Video" else: - return f"Sample {repeat_no}/{repeat_max}" + return f"{item} {repeat_no}/{repeat_max}" else: if repeat_max == 1: return f"Prompt {prompt_no}/{prompts_max}" else: - return f"Prompt {prompt_no}/{prompts_max}, Sample {repeat_no}/{repeat_max}" + return f"Prompt {prompt_no}/{prompts_max}, {item} {repeat_no}/{repeat_max}" refresh_id = 0 @@ -2720,7 +2791,8 @@ def update_status(state): prompts_max = gen.get("prompts_max",0) total_generation = gen["total_generation"] repeat_no = gen["repeat_no"] - status = get_generation_status(prompt_no, prompts_max, repeat_no, total_generation) + sliding_window = gen["sliding_window"] + status = get_generation_status(prompt_no, prompts_max, repeat_no, total_generation, sliding_window) gen["progress_status"] = status gen["refresh"] = get_new_refresh_id() @@ -2737,7 +2809,7 @@ def one_more_sample(state): prompts_max = gen.get("prompts_max",0) total_generation = gen["total_generation"] + extra_orders repeat_no = gen["repeat_no"] - status = get_generation_status(prompt_no, prompts_max, repeat_no, total_generation) + status = get_generation_status(prompt_no, prompts_max, repeat_no, total_generation, gen.get("sliding_window",False)) gen["progress_status"] = status @@ -3059,7 +3131,7 @@ def prepare_inputs_dict(target, inputs ): if not "Vace" in model_filename: - unsaved_params = ["video_prompt_type", "keep_frames", "remove_background_image_ref"] + unsaved_params = ["video_prompt_type", "keep_frames", "remove_background_image_ref", "sliding_window_repeat", "sliding_window_overlap", "sliding_window_discard_last_frames"] for k in unsaved_params: inputs.pop(k) @@ -3102,6 +3174,9 @@ def save_inputs( video_guide, video_mask, keep_frames, + sliding_window_repeat, + sliding_window_overlap, + sliding_window_discard_last_frames, remove_background_image_ref, temporal_upsampling, spatial_upsampling, @@ -3437,7 +3512,6 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non # video_prompt_type_image_refs = gr.Checkbox(value="I" in video_prompt_type_value , label= "Use References Images (Faces, Objects) to customize New Video", scale =1 ) video_guide = gr.Video(label= "Control Video", visible= "V" in video_prompt_type_value, value= ui_defaults.get("video_guide", None),) - # keep_frames = gr.Slider(0, 100, value=ui_defaults.get("keep_frames",0), step=1, label="Nb of frames in Control Video to use (0 = max)", visible= "V" in video_prompt_type_value, scale = 2 ) keep_frames = gr.Text(value=ui_defaults.get("keep_frames","") , visible= "V" in video_prompt_type_value, scale = 2, label= "Frames to keep in Control Video (empty=All, 1=first, a:b for a range, space to separate values)" ) #, -1=last image_refs = gr.Gallery( label ="Reference Images", type ="pil", show_label= True, @@ -3513,28 +3587,32 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non label="Resolution" ) with gr.Row(): - with gr.Column(): - video_length = gr.Slider(5, 193, value=ui_defaults.get("video_length", 81), step=4, label="Number of frames (16 = 1s)") - with gr.Column(): - num_inference_steps = gr.Slider(1, 100, value=ui_defaults.get("num_inference_steps",30), step=1, label="Number of Inference Steps") + video_length = gr.Slider(5, 193, value=ui_defaults.get("video_length", 81), step=4, label="Number of frames (16 = 1s)") + num_inference_steps = gr.Slider(1, 100, value=ui_defaults.get("num_inference_steps",30), step=1, label="Number of Inference Steps") + + + show_advanced = gr.Checkbox(label="Advanced Mode", value=advanced_ui) - with gr.Row(visible=advanced_ui) as advanced_row: - with gr.Column(): - seed = gr.Slider(-1, 999999999, value=ui_defaults["seed"], step=1, label="Seed (-1 for random)") - with gr.Row(): - repeat_generation = gr.Slider(1, 25.0, value=ui_defaults.get("repeat_generation",1), step=1, label="Default Number of Generated Videos per Prompt") - multi_images_gen_type = gr.Dropdown( value=ui_defaults.get("multi_images_gen_type",0), - choices=[ - ("Generate every combination of images and texts", 0), - ("Match images and text prompts", 1), - ], visible= args.multiple_images, label= "Multiple Images as Texts Prompts" - ) - with gr.Row(): - guidance_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("guidance_scale",5), step=0.5, label="Guidance Scale", visible=True) - embedded_guidance_scale = gr.Slider(1.0, 20.0, value=6.0, step=0.5, label="Embedded Guidance Scale", visible=False) - flow_shift = gr.Slider(0.0, 25.0, value=ui_defaults.get("flow_shift",3), step=0.1, label="Shift Scale") - with gr.Row(): - negative_prompt = gr.Textbox(label="Negative Prompt", value=ui_defaults.get("negative_prompt", "") ) + with gr.Tabs(visible=advanced_ui) as advanced_row: + # with gr.Row(visible=advanced_ui) as advanced_row: + with gr.Tab("Generation"): + with gr.Column(): + seed = gr.Slider(-1, 999999999, value=ui_defaults["seed"], step=1, label="Seed (-1 for random)") + with gr.Row(): + repeat_generation = gr.Slider(1, 25.0, value=ui_defaults.get("repeat_generation",1), step=1, label="Default Number of Generated Videos per Prompt") + multi_images_gen_type = gr.Dropdown( value=ui_defaults.get("multi_images_gen_type",0), + choices=[ + ("Generate every combination of images and texts", 0), + ("Match images and text prompts", 1), + ], visible= args.multiple_images, label= "Multiple Images as Texts Prompts" + ) + with gr.Row(): + guidance_scale = gr.Slider(1.0, 20.0, value=ui_defaults.get("guidance_scale",5), step=0.5, label="Guidance Scale", visible=True) + embedded_guidance_scale = gr.Slider(1.0, 20.0, value=6.0, step=0.5, label="Embedded Guidance Scale", visible=False) + flow_shift = gr.Slider(0.0, 25.0, value=ui_defaults.get("flow_shift",3), step=0.1, label="Shift Scale") + with gr.Row(): + negative_prompt = gr.Textbox(label="Negative Prompt", value=ui_defaults.get("negative_prompt", "") ) + with gr.Tab("Loras"): with gr.Column(visible = True): #as loras_column: gr.Markdown("Loras can be used to create special effects on the video by mentioning a trigger word in the Prompt. You can save Loras combinations in presets.") loras_choices = gr.Dropdown( @@ -3548,7 +3626,10 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non loras_multipliers = gr.Textbox(label="Loras Multipliers (1.0 by default) separated by space characters or carriage returns, line that starts with # are ignored", value=launch_multis_str) with gr.Row(): gr.Markdown("Tea Cache accelerates by skipping intelligently some steps, the more steps are skipped the lower the quality of the video (Tea Cache consumes also VRAM)") - with gr.Row(): + with gr.Tab("Speed"): + with gr.Column(): + gr.Markdown("Tea Cache accelerates the Video generation by skipping denoising steps. This may impact the quality") + tea_cache_setting = gr.Dropdown( choices=[ ("Tea Cache Disabled", 0), @@ -3564,9 +3645,10 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ) tea_cache_start_step_perc = gr.Slider(0, 100, value=ui_defaults.get("tea_cache_start_step_perc",0), step=1, label="Tea Cache starting moment in % of generation") - with gr.Row(): + with gr.Tab("Upsampling"): + + with gr.Column(): gr.Markdown("Upsampling - postprocessing that may improve fluidity and the size of the video") - with gr.Row(): temporal_upsampling = gr.Dropdown( choices=[ ("Disabled", ""), @@ -3590,6 +3672,59 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non label="Spatial Upsampling" ) + with gr.Tab("Quality"): + with gr.Row(): + gr.Markdown("Experimental: Skip Layer Guidance, should improve video quality") + with gr.Row(): + slg_switch = gr.Dropdown( + choices=[ + ("OFF", 0), + ("ON", 1), + ], + value=ui_defaults.get("slg_switch",0), + visible=True, + scale = 1, + label="Skip Layer guidance" + ) + slg_layers = gr.Dropdown( + choices=[ + (str(i), i ) for i in range(40) + ], + value=ui_defaults.get("slg_layers", ["9"]), + multiselect= True, + label="Skip Layers", + scale= 3 + ) + with gr.Row(): + slg_start_perc = gr.Slider(0, 100, value=ui_defaults.get("slg_start_perc",10), step=1, label="Denoising Steps % start") + slg_end_perc = gr.Slider(0, 100, value=ui_defaults.get("slg_end_perc",90), step=1, label="Denoising Steps % end") + + with gr.Row(): + gr.Markdown("Experimental: Classifier-Free Guidance Zero Star, better adherence to Text Prompt") + with gr.Row(): + cfg_star_switch = gr.Dropdown( + choices=[ + ("OFF", 0), + ("ON", 1), + ], + value=ui_defaults.get("cfg_star_switch",0), + visible=True, + scale = 1, + label="CFG Star" + ) + with gr.Row(): + cfg_zero_step = gr.Slider(-1, 39, value=ui_defaults.get("cfg_zero_step",-1), step=1, label="CFG Zero below this Layer (Extra Process)") + + with gr.Tab("Sliding Window", visible= "Vace" in model_filename ) as sliding_window_tab: + + with gr.Column(visible= "Vace" in model_filename ) as sliding_window_row: + gr.Markdown("A Sliding Window allows you to generate video longer than those of the model limits") + + sliding_window_repeat = gr.Slider(0, 20, value=ui_defaults.get("sliding_window_repeat", 0), step=1, label="Sliding Window Iterations (O=Disabled)") + sliding_window_overlap = gr.Slider(1, 32, value=ui_defaults.get("sliding_window_overlap",16), step=1, label="Windows Frames Overlap (needed to maintain continuity between windows, a higher value will require more windows)") + sliding_window_discard_last_frames = gr.Slider(1, 10, value=ui_defaults.get("sliding_window_discard_last_frames", 4), step=1, label="Discard Last Frames of a Window (that may have bad quality)") + + with gr.Tab("Miscellaneous"): gr.Markdown("With Riflex you can generate videos longer than 5s which is the default duration of videos used to train the model") RIFLEx_setting = gr.Dropdown( choices=[ @@ -3600,50 +3735,9 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non value=ui_defaults.get("RIFLEx_setting",0), label="RIFLEx positional embedding to generate long video" ) - with gr.Row(): - gr.Markdown("Experimental: Skip Layer Guidance, should improve video quality") - with gr.Row(): - slg_switch = gr.Dropdown( - choices=[ - ("OFF", 0), - ("ON", 1), - ], - value=ui_defaults.get("slg_switch",0), - visible=True, - scale = 1, - label="Skip Layer guidance" - ) - slg_layers = gr.Dropdown( - choices=[ - (str(i), i ) for i in range(40) - ], - value=ui_defaults.get("slg_layers", ["9"]), - multiselect= True, - label="Skip Layers", - scale= 3 - ) - with gr.Row(): - slg_start_perc = gr.Slider(0, 100, value=ui_defaults.get("slg_start_perc",10), step=1, label="Denoising Steps % start") - slg_end_perc = gr.Slider(0, 100, value=ui_defaults.get("slg_end_perc",90), step=1, label="Denoising Steps % end") - with gr.Row(): - gr.Markdown("Experimental: Classifier-Free Guidance Zero Star, better adherence to Text Prompt") - with gr.Row(): - cfg_star_switch = gr.Dropdown( - choices=[ - ("OFF", 0), - ("ON", 1), - ], - value=ui_defaults.get("cfg_star_switch",0), - visible=True, - scale = 1, - label="CFG Star" - ) - with gr.Row(): - cfg_zero_step = gr.Slider(-1, 39, value=ui_defaults.get("cfg_zero_step",-1), step=1, label="CFG Zero below this Layer (Extra Process)") - - with gr.Row(): - save_settings_btn = gr.Button("Set Settings as Default", visible = not args.lock_config) + with gr.Row(): + save_settings_btn = gr.Button("Set Settings as Default", visible = not args.lock_config) if not update_form: with gr.Column(): @@ -3697,11 +3791,11 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non let countdown = 5; const label = document.getElementById('quit_timer_label'); if (label) { - label.innerText = `Quitting in ${countdown}...`; + label.innerText = `${countdown}...`; window.quitCountdownInterval = setInterval(() => { countdown--; if (countdown > 0) { - label.innerText = `Quitting in ${countdown}...`; + label.innerText = `${countdown}`; } else { clearInterval(window.quitCountdownInterval); findAndClickGradioButton('comfirm_quit_btn_hidden'); @@ -3841,7 +3935,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non ) extra_inputs = prompt_vars + [wizard_prompt, wizard_variables_var, wizard_prompt_activated_var, video_prompt_column, image_prompt_column, - prompt_column_advanced, prompt_column_wizard_vars, prompt_column_wizard, lset_name, advanced_row] # show_advanced presets_column, + prompt_column_advanced, prompt_column_wizard_vars, prompt_column_wizard, lset_name, advanced_row, sliding_window_tab] # show_advanced presets_column, if update_form: locals_dict = locals() gen_inputs = [state_dict if k=="state" else locals_dict[k] for k in inputs_names] + [state_dict] + extra_inputs @@ -4141,13 +4235,14 @@ def generate_about_tab(): gr.Markdown("- Alibaba Wan team for the best open source video generator") gr.Markdown("- Alibaba Vace and Fun Teams for their incredible control net models") gr.Markdown("- Cocktail Peanuts : QA and simple installation via Pinokio.computer") - gr.Markdown("- Tophness : created multi tabs and queuing frameworks") + gr.Markdown("- Tophness : created (former) multi tabs and queuing frameworks") gr.Markdown("- AmericanPresidentJimmyCarter : added original support for Skip Layer Guidance") gr.Markdown("- Remade_AI : for their awesome Loras collection") gr.Markdown("
Huge acknowlegments to these great open source projects used in WanGP:") gr.Markdown("- Rife: temporal upsampler (https://github.com/hzwer/ECCV2022-RIFE)") gr.Markdown("- DwPose: Open Pose extractor (https://github.com/IDEA-Research/DWPose)") gr.Markdown("- Midas: Depth extractor (https://github.com/isl-org/MiDaS") + gr.Markdown("- Matanyone and SAM2: Mask Generation (https://github.com/pq-yang/MatAnyone) and (https://github.com/facebookresearch/sam2)") def generate_info_tab(): From 6b3ca42acb89835144d530962e945ca6ff3b15d4 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Sun, 13 Apr 2025 02:21:08 +0200 Subject: [PATCH 61/69] Fixed Sliding Window tab sometime partially hidden --- wgp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wgp.py b/wgp.py index 283aeb8..09a09ae 100644 --- a/wgp.py +++ b/wgp.py @@ -3717,7 +3717,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non with gr.Tab("Sliding Window", visible= "Vace" in model_filename ) as sliding_window_tab: - with gr.Column(visible= "Vace" in model_filename ) as sliding_window_row: + with gr.Column(): gr.Markdown("A Sliding Window allows you to generate video longer than those of the model limits") sliding_window_repeat = gr.Slider(0, 20, value=ui_defaults.get("sliding_window_repeat", 0), step=1, label="Sliding Window Iterations (O=Disabled)") @@ -3791,7 +3791,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non let countdown = 5; const label = document.getElementById('quit_timer_label'); if (label) { - label.innerText = `${countdown}...`; + label.innerText = `M${countdown}...`; window.quitCountdownInterval = setInterval(() => { countdown--; if (countdown > 0) { From 5efddd626d038647f8d8c0169cab731ed10629ac Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Sun, 13 Apr 2025 11:02:20 +0200 Subject: [PATCH 62/69] Fixed Vace bug when sliding windows not enabled --- wgp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wgp.py b/wgp.py index 09a09ae..d0dfc6d 100644 --- a/wgp.py +++ b/wgp.py @@ -2376,9 +2376,9 @@ def generate_video( gen["prompt"] = prompt repeat_no = 0 extra_generation = 0 + start_frame = 0 sliding_window = sliding_window_repeat > 0 if sliding_window: - start_frame = 0 reuse_frames = sliding_window_overlap discard_last_frames = sliding_window_discard_last_frames #4 repeat_generation = sliding_window_repeat @@ -3718,7 +3718,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non with gr.Tab("Sliding Window", visible= "Vace" in model_filename ) as sliding_window_tab: with gr.Column(): - gr.Markdown("A Sliding Window allows you to generate video longer than those of the model limits") + gr.Markdown("A Sliding Window allows you to generate video with a duration not limited by the Model") sliding_window_repeat = gr.Slider(0, 20, value=ui_defaults.get("sliding_window_repeat", 0), step=1, label="Sliding Window Iterations (O=Disabled)") sliding_window_overlap = gr.Slider(1, 32, value=ui_defaults.get("sliding_window_overlap",16), step=1, label="Windows Frames Overlap (needed to maintain continuity between windows, a higher value will require more windows)") From 3044a0b4866ffc7315710b77f75a5eca137908d8 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 14 Apr 2025 04:03:45 +1000 Subject: [PATCH 63/69] gr.info based quit countdown timer --- wgp.py | 85 +++++++++++++++++++++++++++++++--------------------------- 1 file changed, 46 insertions(+), 39 deletions(-) diff --git a/wgp.py b/wgp.py index da7273a..79dc592 100644 --- a/wgp.py +++ b/wgp.py @@ -732,11 +732,17 @@ def quit_application(): import signal os.kill(os.getpid(), signal.SIGINT) -def request_quit_confirmation(): - return gr.update(visible=False), gr.update(visible=True) +def start_quit_process(): + return 5, gr.update(visible=False), gr.update(visible=True) -def cancel_quit_confirmation(): - return gr.update(visible=True), gr.update(visible=False) +def cancel_quit_process(): + return -1, gr.update(visible=True), gr.update(visible=False) + +def show_countdown_info_from_state(current_value: int): + if current_value > 0: + gr.Info(f"Quitting in {current_value}...") + return current_value - 1 + return current_value def autosave_queue(): global global_queue_ref @@ -994,7 +1000,6 @@ def create_html_progress_bar(percentage=0.0, text="Idle", is_idle=True): """ return html - def update_generation_status(html_content): if(html_content): return gr.update(value=html_content) @@ -3625,51 +3630,46 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non clear_queue_btn = gr.Button("Clear Queue", size="sm", variant="stop") quit_button = gr.Button("Save and Quit", size="sm", variant="secondary") with gr.Row(visible=False) as quit_confirmation_row: - gr.Markdown("Quitting in 5 seconds...", elem_id="quit_timer_label") - confirm_quit_button = gr.Button("Confirm Quit Now", elem_id="comfirm_quit_btn_hidden", size="sm", variant="stop") - cancel_quit_button = gr.Button("Cancel Quit", size="sm", variant="secondary") + confirm_quit_button = gr.Button("Confirm", elem_id="comfirm_quit_btn_hidden", size="sm", variant="stop") + cancel_quit_button = gr.Button("Cancel", size="sm", variant="secondary") hidden_force_quit_trigger = gr.Button("force_quit", visible=False, elem_id="force_quit_btn_hidden") + hidden_countdown_state = gr.Number(value=-1, visible=False, elem_id="hidden_countdown_state_num") + single_hidden_trigger_btn = gr.Button("trigger_countdown", visible=False, elem_id="trigger_info_single_btn") start_quit_timer_js = """ () => { function findAndClickGradioButton(elemId) { const gradioApp = document.querySelector('gradio-app') || document; const button = gradioApp.querySelector(`#${elemId}`); - if (button) { - button.click(); + if (button) { button.click(); } + } + + if (window.quitCountdownTimeoutId) clearTimeout(window.quitCountdownTimeoutId); + + let js_click_count = 0; + const max_clicks = 5; + + function countdownStep() { + if (js_click_count < max_clicks) { + findAndClickGradioButton('trigger_info_single_btn'); + js_click_count++; + window.quitCountdownTimeoutId = setTimeout(countdownStep, 1000); + } else { + findAndClickGradioButton('force_quit_btn_hidden'); } } - window.quitTimerId = setTimeout(() => { - }, 5000); - let countdown = 5; - const label = document.getElementById('quit_timer_label'); - if (label) { - label.innerText = `Quitting in ${countdown}...`; - window.quitCountdownInterval = setInterval(() => { - countdown--; - if (countdown > 0) { - label.innerText = `Quitting in ${countdown}...`; - } else { - clearInterval(window.quitCountdownInterval); - findAndClickGradioButton('comfirm_quit_btn_hidden'); - } - }, 1000); - } + + countdownStep(); } """ cancel_quit_timer_js = """ () => { - if (window.quitTimerId) { - clearTimeout(window.quitTimerId); - window.quitTimerId = null; + if (window.quitCountdownTimeoutId) { + clearTimeout(window.quitCountdownTimeoutId); + window.quitCountdownTimeoutId = null; + console.log("Quit countdown cancelled (single trigger)."); } - if(window.quitCountdownInterval) { - clearInterval(window.quitCountdownInterval); - window.quitCountdownInterval = null; - } - const label = document.getElementById('quit_timer_label'); - if(label) { label.innerText = 'Quit cancelled.'; } } """ @@ -3705,10 +3705,15 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non } """ + single_hidden_trigger_btn.click( + fn=show_countdown_info_from_state, + inputs=[hidden_countdown_state], + outputs=[hidden_countdown_state] + ) quit_button.click( - fn=request_quit_confirmation, + fn=start_quit_process, inputs=[], - outputs=[quit_button, quit_confirmation_row] + outputs=[hidden_countdown_state, quit_button, quit_confirmation_row] ).then( fn=None, inputs=None, outputs=None, js=start_quit_timer_js ) @@ -3717,12 +3722,14 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non fn=quit_application, inputs=[], outputs=[] + ).then( + fn=None, inputs=None, outputs=None, js=cancel_quit_timer_js ) cancel_quit_button.click( - fn=cancel_quit_confirmation, + fn=cancel_quit_process, inputs=[], - outputs=[quit_button, quit_confirmation_row] + outputs=[hidden_countdown_state, quit_button, quit_confirmation_row] ).then( fn=None, inputs=None, outputs=None, js=cancel_quit_timer_js ) From aef624bd5bbf8655c2fcb680100b607f83012d5d Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 14 Apr 2025 04:53:21 +1000 Subject: [PATCH 64/69] fix pytorch version extraction --- preprocessing/matanyone/tools/misc.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/preprocessing/matanyone/tools/misc.py b/preprocessing/matanyone/tools/misc.py index 43b8499..868639c 100644 --- a/preprocessing/matanyone/tools/misc.py +++ b/preprocessing/matanyone/tools/misc.py @@ -52,9 +52,14 @@ def get_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=None initialized_logger[logger_name] = True return logger - -IS_HIGH_VERSION = [int(m) for m in list(re.findall(r"^([0-9]+)\.([0-9]+)\.([0-9]+)([^0-9][a-zA-Z0-9]*)?(\+git.*)?$",\ - torch.__version__)[0][:3])] >= [1, 12, 0] +match = re.match(r"^([0-9]+)\.([0-9]+)\.([0-9]+)", torch.__version__) +if match: + version_tuple = match.groups() + IS_HIGH_VERSION = [int(v) for v in version_tuple] >= [1, 12, 0] +else: + logger = get_root_logger() + logger.warning(f"Could not parse torch version '{torch.__version__}'. Assuming it's not a high version >= 1.12.0.") + IS_HIGH_VERSION = False def gpu_is_available(): if IS_HIGH_VERSION: From 28ea76b0d14d7df7cfc02cae3c1dee77a2417e7b Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 14 Apr 2025 05:44:03 +1000 Subject: [PATCH 65/69] add autodownloader for missing ffmpeg binaries --- wgp.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/wgp.py b/wgp.py index c30c5ec..2c71961 100644 --- a/wgp.py +++ b/wgp.py @@ -32,6 +32,9 @@ import zipfile import tempfile import atexit import shutil +import urllib.request +from tqdm import tqdm +import requests global_queue_ref = [] AUTOSAVE_FILENAME = "queue.zip" PROMPT_VARS_MAX = 10 @@ -48,6 +51,30 @@ task_id = 0 # progress_tracker = {} # tracker_lock = threading.Lock() +def download_ffmpeg(): + if os.name != 'nt': return + exes = ['ffmpeg.exe', 'ffprobe.exe', 'ffplay.exe'] + if all(os.path.exists(e) for e in exes): return + api_url = 'https://api.github.com/repos/GyanD/codexffmpeg/releases/latest' + r = requests.get(api_url, headers={'Accept': 'application/vnd.github+json'}) + assets = r.json().get('assets', []) + zip_asset = next((a for a in assets if 'essentials_build.zip' in a['name']), None) + if not zip_asset: return + zip_url = zip_asset['browser_download_url'] + zip_name = zip_asset['name'] + with requests.get(zip_url, stream=True) as resp: + total = int(resp.headers.get('Content-Length', 0)) + with open(zip_name, 'wb') as f, tqdm(total=total, unit='B', unit_scale=True) as pbar: + for chunk in resp.iter_content(chunk_size=8192): + f.write(chunk) + pbar.update(len(chunk)) + with zipfile.ZipFile(zip_name) as z: + for f in z.namelist(): + if f.endswith(tuple(exes)) and '/bin/' in f: + z.extract(f) + os.rename(f, os.path.basename(f)) + os.remove(zip_name) + def format_time(seconds): if seconds < 60: return f"{seconds:.1f}s" @@ -4601,6 +4628,7 @@ def create_demo(): if __name__ == "__main__": atexit.register(autosave_queue) + download_ffmpeg() # threading.Thread(target=runner, daemon=True).start() os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" server_port = int(args.server_port) From a0e805013f2d6aaad460f25f4b9de0ec9ac1c8a4 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Mon, 14 Apr 2025 05:46:29 +1000 Subject: [PATCH 66/69] remove redundant import --- wgp.py | 1 - 1 file changed, 1 deletion(-) diff --git a/wgp.py b/wgp.py index 2c71961..390c651 100644 --- a/wgp.py +++ b/wgp.py @@ -32,7 +32,6 @@ import zipfile import tempfile import atexit import shutil -import urllib.request from tqdm import tqdm import requests global_queue_ref = [] From c62beb7d9dba87e79e3ac80477c3ee0135a2e0f4 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Tue, 15 Apr 2025 01:02:06 +0200 Subject: [PATCH 67/69] Added Low VRAM support for RTX 10XX and RTX 20XX GPUs --- README.md | 33 +++- gradio/i2v_14B_singleGPU.py | 306 --------------------------------- gradio/t2i_14B_singleGPU.py | 206 ---------------------- gradio/t2v_1.3B_singleGPU.py | 207 ---------------------- gradio/t2v_14B_singleGPU.py | 216 ----------------------- preprocessing/matanyone/app.py | 83 +++++---- requirements.txt | 3 +- wan/image2video.py | 61 +++---- wan/modules/attention.py | 19 +- wan/modules/model.py | 75 +------- wan/modules/vae.py | 11 +- wan/text2video.py | 47 ++--- wgp.py | 227 ++++++++++++++---------- 13 files changed, 279 insertions(+), 1215 deletions(-) delete mode 100644 gradio/i2v_14B_singleGPU.py delete mode 100644 gradio/t2i_14B_singleGPU.py delete mode 100644 gradio/t2v_1.3B_singleGPU.py delete mode 100644 gradio/t2v_14B_singleGPU.py diff --git a/README.md b/README.md index 9b7ccbf..2bbb1ef 100644 --- a/README.md +++ b/README.md @@ -15,12 +15,13 @@ ## 🔥 Latest News!! * April 13 2025: 👋 Wan 2.1GP v4.0: lots of goodies for you ! - - A new queuing system that lets you stack in a queue as many text2video and imag2video tasks as you want. Each task can rely on complete different generation parameters (different number of frames, steps, loras, ...). - - Temporal upsampling (Rife) and spatial upsampling (Lanczos) for a smoother video (32 fps or 64 fps) and to enlarge you video by x2 or x4. Check these new advanced options. - - Wan Vace Control Net support : with Vace you can inject in the scene people or objects, animate a person, perform inpainting or outpainting, continue a video, ... I have provided an introduction guide below. - - Integrated *Matanyone* tool directly inside WanGP so that you can create easily inpainting masks - - Sliding Window generation for Vace, create windows that can last dozen of seconds - A new UI, tabs were replaced by a Dropdown box to easily switch models + - A new queuing system that lets you stack in a queue as many text2video, imag2video tasks, ... as you want. Each task can rely on complete different generation parameters (different number of frames, steps, loras, ...). Many thanks to *Tophness** for being a big contributor on this new feature + - Temporal upsampling (Rife) and spatial upsampling (Lanczos) for a smoother video (32 fps or 64 fps) and to enlarge your video by x2 or x4. Check these new advanced options. + - Wan Vace Control Net support : with Vace you can inject in the scene people or objects, animate a person, perform inpainting or outpainting, continue a video, ... I have provided an introduction guide below. + - Integrated *Matanyone* tool directly inside WanGP so that you can create easily inpainting masks used in Vace + - Sliding Window generation for Vace, create windows that can last dozen of seconds + - New optimisations for old generation GPUs: Generate 5s (81 frames, 15 steps) of Vace 1.3B with only 5GB and in only 6 minutes on a RTX 2080Ti and 5s of t2v 14B in less than 10 minutes. * Mar 27 2025: 👋 Added support for the new Wan Fun InP models (image2video). The 14B Fun InP has probably better end image support but unfortunately existing loras do not work so well with it. The great novelty is the Fun InP image2 1.3B model : Image 2 Video is now accessible to even lower hardware configuration. It is not as good as the 14B models but very impressive for its size. You can choose any of those models in the Configuration tab. Many thanks to the VideoX-Fun team (https://github.com/aigc-apps/VideoX-Fun) * Mar 26 2025: 👋 Good news ! Official support for RTX 50xx please check the installation instructions below. @@ -303,6 +304,20 @@ Vace provides on its github (https://github.com/ali-vilab/VACE/tree/main/vace/gr There is also a guide that describes the various combination of hints (https://github.com/ali-vilab/VACE/blob/main/UserGuide.md).Good luck ! It seems you will get better results if you turn on "Skip Layer Guidance" with its default configuration + +### VACE Slidig Window +With this mode (that works for the moment only with Vace) you can merge mutiple Videos to form a very long video (up to 1 min). What is this very nice a about this feature is that the resulting video can be driven by the same control video. For instance the first 0-4s of the control video will be used to generate the first window then the next 4-8s of the control video will be used to generate the second window, and so on. So if your control video contains a person walking, your generate video could contain up to one minute of this person walking. + +To turn on sliding window, you need to go in the Advanced Settings Tab *Sliding Window* and set the iteration number to a number greater than 1. This number corresponds to the default number of windows. You can still increase the number during the genreation by clicking the "One More Sample, Please !" button. + +Each window duration will be set by the *Number of frames (16 = 1s)* form field. However the actual number of frames generated by each iteration will be less, because the *overlap frames* and *discard last frames*: +- *overlap frames* : the first frames ofa new window are filled with last frames of the previous window in order to ensure continuity between the two windows +- *discard last frames* : quite often the last frames of a window have a worse quality. You decide here how many ending frames of a new window should be dropped. + +Number of Generated = [Number of iterations] * ([Number of frames] - [Overlap Frames] - [Discard Last Frames]) + [Overlap Frames] + +Experimental: if your prompt is broken into multiple lines (each line separated by a carriage return), then each line of the prompt will be used for a new window. If there are more windows to generate than prompt lines, the last prompt line will be repeated. + ### Command line parameters for Gradio Server --i2v : launch the image to video generator\ --t2v : launch the text to video generator (default defined in the configuration)\ @@ -324,7 +339,7 @@ It seems you will get better results if you turn on "Skip Layer Guidance" with i --compile : turn on pytorch compilation\ --attention mode: force attention mode among, sdpa, flash, sage, sage2\ --profile no : default (4) : no of profile between 1 and 5\ ---preload no : number in Megabytes to preload partially the diffusion model in VRAM , may offer slight speed gains especially on older hardware. Works only with profile 2 and 4.\ +--preload no : number in Megabytes to preload partially the diffusion model in VRAM , may offer speed gains on older hardware, on recent hardware (RTX 30XX, RTX40XX and RTX50XX) speed gain is only 10% and not worth it. Works only with profile 2 and 4.\ --seed no : set default seed value\ --frames no : set the default number of frames to generate\ --steps no : set the default number of denoising steps\ @@ -333,7 +348,11 @@ It seems you will get better results if you turn on "Skip Layer Guidance" with i --check-loras : filter loras that are incompatible (will take a few seconds while refreshing the lora list or while starting the app)\ --advanced : turn on the advanced mode while launching the app\ --listen : make server accessible on network\ ---gpu device : run Wan on device for instance "cuda:1" +--gpu device : run Wan on device for instance "cuda:1"\ +--settings: path a folder that contains the default settings for all the models\ +--fp16: force to use fp16 versions of models instead of bf16 versions\ +--perc-reserved-mem-max float_less_than_1 : max percentage of RAM to allocate to reserved RAM, allow faster transfers RAM<->VRAM. Value should remain below 0.5 to keep the OS stable\ +--theme theme_name: load the UI with the specified Theme Name, so far only two are supported, "default" and "gradio". You may submit your own nice looking Gradio theme and I will add them ### Profiles (for power users only) You can choose between 5 profiles, but two are really relevant here : diff --git a/gradio/i2v_14B_singleGPU.py b/gradio/i2v_14B_singleGPU.py deleted file mode 100644 index 031479b..0000000 --- a/gradio/i2v_14B_singleGPU.py +++ /dev/null @@ -1,306 +0,0 @@ -# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. -import argparse -import gc -import os.path as osp -import os -import sys -import warnings - -import gradio as gr - -warnings.filterwarnings('ignore') - -# Model -sys.path.insert(0, os.path.sep.join(osp.realpath(__file__).split(os.path.sep)[:-2])) -import wan -from wan.configs import MAX_AREA_CONFIGS, WAN_CONFIGS -from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander -from wan.utils.utils import cache_video - -# Global Var -prompt_expander = None -wan_i2v_480P = None -wan_i2v_720P = None - - -# Button Func -def load_i2v_model(value): - global wan_i2v_480P, wan_i2v_720P - from mmgp import offload - - if value == '------': - print("No model loaded") - return '------' - - if value == '720P': - if args.ckpt_dir_720p is None: - print("Please specify the checkpoint directory for 720P model") - return '------' - if wan_i2v_720P is not None: - pass - else: - del wan_i2v_480P - gc.collect() - wan_i2v_480P = None - - print("load 14B-720P i2v model...", end='', flush=True) - cfg = WAN_CONFIGS['i2v-14B'] - wan_i2v_720P = wan.WanI2V( - config=cfg, - checkpoint_dir=args.ckpt_dir_720p, - device_id=0, - rank=0, - t5_fsdp=False, - dit_fsdp=False, - use_usp=False, - i2v720p= True - ) - print("done", flush=True) - pipe = {"transformer": wan_i2v_720P.model, "text_encoder" : wan_i2v_720P.text_encoder.model, "text_encoder_2": wan_i2v_720P.clip.model, "vae": wan_i2v_720P.vae.model } # - offload.profile(pipe, profile_no=4, budgets = {"transformer":100, "*":3000}, verboseLevel=2, compile="transformer", quantizeTransformer = False, pinnedMemory = False) - return '720P' - - if value == '480P': - if args.ckpt_dir_480p is None: - print("Please specify the checkpoint directory for 480P model") - return '------' - if wan_i2v_480P is not None: - pass - else: - del wan_i2v_720P - gc.collect() - wan_i2v_720P = None - - print("load 14B-480P i2v model...", end='', flush=True) - cfg = WAN_CONFIGS['i2v-14B'] - wan_i2v_480P = wan.WanI2V( - config=cfg, - checkpoint_dir=args.ckpt_dir_480p, - device_id=0, - rank=0, - t5_fsdp=False, - dit_fsdp=False, - use_usp=False, - i2v720p= False - ) - print("done", flush=True) - pipe = {"transformer": wan_i2v_480P.model, "text_encoder" : wan_i2v_480P.text_encoder.model, "text_encoder_2": wan_i2v_480P.clip.model, "vae": wan_i2v_480P.vae.model } # - offload.profile(pipe, profile_no=4, budgets = {"model":100, "*":3000}, verboseLevel=2, compile="transformer" ) - - return '480P' - - - -def prompt_enc(prompt, img, tar_lang): - print('prompt extend...') - if img is None: - print('Please upload an image') - return prompt - global prompt_expander - prompt_output = prompt_expander( - prompt, image=img, tar_lang=tar_lang.lower()) - if prompt_output.status == False: - return prompt - else: - return prompt_output.prompt - - -def i2v_generation(img2vid_prompt, img2vid_image, res, sd_steps, - guide_scale, shift_scale, seed, n_prompt): - # print(f"{img2vid_prompt},{resolution},{sd_steps},{guide_scale},{shift_scale},{seed},{n_prompt}") - global resolution - from PIL import Image - img2vid_image = Image.open("d:\mammoth2.jpg") - if resolution == '------': - print( - 'Please specify at least one resolution ckpt dir or specify the resolution' - ) - return None - - else: - if resolution == '720P': - global wan_i2v_720P - video = wan_i2v_720P.generate( - img2vid_prompt, - img2vid_image, - max_area=MAX_AREA_CONFIGS['720*1280'], - shift=shift_scale, - sampling_steps=sd_steps, - guide_scale=guide_scale, - n_prompt=n_prompt, - seed=seed, - offload_model=False) - else: - global wan_i2v_480P - video = wan_i2v_480P.generate( - img2vid_prompt, - img2vid_image, - max_area=MAX_AREA_CONFIGS['480*832'], - shift=3.0, #shift_scale - sampling_steps=sd_steps, - guide_scale=guide_scale, - n_prompt=n_prompt, - seed=seed, - offload_model=False) - - cache_video( - tensor=video[None], - save_file="example.mp4", - fps=16, - nrow=1, - normalize=True, - value_range=(-1, 1)) - - return "example.mp4" - - -# Interface -def gradio_interface(): - with gr.Blocks() as demo: - gr.Markdown(""" -
- Wan2.1 (I2V-14B) -
-
- Wan: Open and Advanced Large-Scale Video Generative Models. -
- """) - - with gr.Row(): - with gr.Column(): - resolution = gr.Dropdown( - label='Resolution', - choices=['------', '720P', '480P'], - value='------') - - img2vid_image = gr.Image( - type="pil", - label="Upload Input Image", - elem_id="image_upload", - ) - img2vid_prompt = gr.Textbox( - label="Prompt", - value="Several giant wooly mammoths approach treading through a snowy meadow, their long wooly fur lightly blows in the wind as they walk, snow covered trees and dramatic snow capped mountains in the distance, mid afternoon light with wispy clouds and a sun high in the distance creates a warm glow, the low camera view is stunning capturing the large furry mammal with beautiful photography, depth of field.", - placeholder="Describe the video you want to generate", - ) - tar_lang = gr.Radio( - choices=["CH", "EN"], - label="Target language of prompt enhance", - value="CH") - run_p_button = gr.Button(value="Prompt Enhance") - - with gr.Accordion("Advanced Options", open=True): - with gr.Row(): - sd_steps = gr.Slider( - label="Diffusion steps", - minimum=1, - maximum=1000, - value=50, - step=1) - guide_scale = gr.Slider( - label="Guide scale", - minimum=0, - maximum=20, - value=5.0, - step=1) - with gr.Row(): - shift_scale = gr.Slider( - label="Shift scale", - minimum=0, - maximum=10, - value=5.0, - step=1) - seed = gr.Slider( - label="Seed", - minimum=-1, - maximum=2147483647, - step=1, - value=-1) - n_prompt = gr.Textbox( - label="Negative Prompt", - placeholder="Describe the negative prompt you want to add" - ) - - run_i2v_button = gr.Button("Generate Video") - - with gr.Column(): - result_gallery = gr.Video( - label='Generated Video', interactive=False, height=600) - - resolution.input( - fn=load_model, inputs=[resolution], outputs=[resolution]) - - run_p_button.click( - fn=prompt_enc, - inputs=[img2vid_prompt, img2vid_image, tar_lang], - outputs=[img2vid_prompt]) - - run_i2v_button.click( - fn=i2v_generation, - inputs=[ - img2vid_prompt, img2vid_image, resolution, sd_steps, - guide_scale, shift_scale, seed, n_prompt - ], - outputs=[result_gallery], - ) - - return demo - - -# Main -def _parse_args(): - parser = argparse.ArgumentParser( - description="Generate a video from a text prompt or image using Gradio") - parser.add_argument( - "--ckpt_dir_720p", - type=str, - default=None, - help="The path to the checkpoint directory.") - parser.add_argument( - "--ckpt_dir_480p", - type=str, - default=None, - help="The path to the checkpoint directory.") - parser.add_argument( - "--prompt_extend_method", - type=str, - default="local_qwen", - choices=["dashscope", "local_qwen"], - help="The prompt extend method to use.") - parser.add_argument( - "--prompt_extend_model", - type=str, - default=None, - help="The prompt extend model to use.") - - args = parser.parse_args() - args.ckpt_dir_720p = "../ckpts" # os.path.join("ckpt") - args.ckpt_dir_480p = "../ckpts" # os.path.join("ckpt") - assert args.ckpt_dir_720p is not None or args.ckpt_dir_480p is not None, "Please specify at least one checkpoint directory." - - return args - - -if __name__ == '__main__': - args = _parse_args() - global resolution - # load_model('720P') - # resolution = '720P' - resolution = '480P' - - load_i2v_model(resolution) - - print("Step1: Init prompt_expander...", end='', flush=True) - if args.prompt_extend_method == "dashscope": - prompt_expander = DashScopePromptExpander( - model_name=args.prompt_extend_model, is_vl=True) - elif args.prompt_extend_method == "local_qwen": - prompt_expander = QwenPromptExpander( - model_name=args.prompt_extend_model, is_vl=True, device=0) - else: - raise NotImplementedError( - f"Unsupport prompt_extend_method: {args.prompt_extend_method}") - print("done", flush=True) - - demo = gradio_interface() - demo.launch(server_name="0.0.0.0", share=False, server_port=7860) diff --git a/gradio/t2i_14B_singleGPU.py b/gradio/t2i_14B_singleGPU.py deleted file mode 100644 index 172a1b0..0000000 --- a/gradio/t2i_14B_singleGPU.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. -import argparse -import os.path as osp -import os -import sys -import warnings - -import gradio as gr - -warnings.filterwarnings('ignore') - -# Model -sys.path.insert(0, os.path.sep.join(osp.realpath(__file__).split(os.path.sep)[:-2])) -import wan -from wan.configs import WAN_CONFIGS -from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander -from wan.utils.utils import cache_image - -# Global Var -prompt_expander = None -wan_t2i = None - - -# Button Func -def prompt_enc(prompt, tar_lang): - global prompt_expander - prompt_output = prompt_expander(prompt, tar_lang=tar_lang.lower()) - if prompt_output.status == False: - return prompt - else: - return prompt_output.prompt - - -def t2i_generation(txt2img_prompt, resolution, sd_steps, guide_scale, - shift_scale, seed, n_prompt): - global wan_t2i - # print(f"{txt2img_prompt},{resolution},{sd_steps},{guide_scale},{shift_scale},{seed},{n_prompt}") - - W = int(resolution.split("*")[0]) - H = int(resolution.split("*")[1]) - video = wan_t2i.generate( - txt2img_prompt, - size=(W, H), - frame_num=1, - shift=shift_scale, - sampling_steps=sd_steps, - guide_scale=guide_scale, - n_prompt=n_prompt, - seed=seed, - offload_model=True) - - cache_image( - tensor=video.squeeze(1)[None], - save_file="example.png", - nrow=1, - normalize=True, - value_range=(-1, 1)) - - return "example.png" - - -# Interface -def gradio_interface(): - with gr.Blocks() as demo: - gr.Markdown(""" -
- Wan2.1 (T2I-14B) -
-
- Wan: Open and Advanced Large-Scale Video Generative Models. -
- """) - - with gr.Row(): - with gr.Column(): - txt2img_prompt = gr.Textbox( - label="Prompt", - placeholder="Describe the image you want to generate", - ) - tar_lang = gr.Radio( - choices=["CH", "EN"], - label="Target language of prompt enhance", - value="CH") - run_p_button = gr.Button(value="Prompt Enhance") - - with gr.Accordion("Advanced Options", open=True): - resolution = gr.Dropdown( - label='Resolution(Width*Height)', - choices=[ - '720*1280', '1280*720', '960*960', '1088*832', - '832*1088', '480*832', '832*480', '624*624', - '704*544', '544*704' - ], - value='720*1280') - - with gr.Row(): - sd_steps = gr.Slider( - label="Diffusion steps", - minimum=1, - maximum=1000, - value=50, - step=1) - guide_scale = gr.Slider( - label="Guide scale", - minimum=0, - maximum=20, - value=5.0, - step=1) - with gr.Row(): - shift_scale = gr.Slider( - label="Shift scale", - minimum=0, - maximum=10, - value=5.0, - step=1) - seed = gr.Slider( - label="Seed", - minimum=-1, - maximum=2147483647, - step=1, - value=-1) - n_prompt = gr.Textbox( - label="Negative Prompt", - placeholder="Describe the negative prompt you want to add" - ) - - run_t2i_button = gr.Button("Generate Image") - - with gr.Column(): - result_gallery = gr.Image( - label='Generated Image', interactive=False, height=600) - - run_p_button.click( - fn=prompt_enc, - inputs=[txt2img_prompt, tar_lang], - outputs=[txt2img_prompt]) - - run_t2i_button.click( - fn=t2i_generation, - inputs=[ - txt2img_prompt, resolution, sd_steps, guide_scale, shift_scale, - seed, n_prompt - ], - outputs=[result_gallery], - ) - - return demo - - -# Main -def _parse_args(): - parser = argparse.ArgumentParser( - description="Generate a image from a text prompt or image using Gradio") - parser.add_argument( - "--ckpt_dir", - type=str, - default="cache", - help="The path to the checkpoint directory.") - parser.add_argument( - "--prompt_extend_method", - type=str, - default="local_qwen", - choices=["dashscope", "local_qwen"], - help="The prompt extend method to use.") - parser.add_argument( - "--prompt_extend_model", - type=str, - default=None, - help="The prompt extend model to use.") - - args = parser.parse_args() - - return args - - -if __name__ == '__main__': - args = _parse_args() - - print("Step1: Init prompt_expander...", end='', flush=True) - if args.prompt_extend_method == "dashscope": - prompt_expander = DashScopePromptExpander( - model_name=args.prompt_extend_model, is_vl=False) - elif args.prompt_extend_method == "local_qwen": - prompt_expander = QwenPromptExpander( - model_name=args.prompt_extend_model, is_vl=False, device=0) - else: - raise NotImplementedError( - f"Unsupport prompt_extend_method: {args.prompt_extend_method}") - print("done", flush=True) - - print("Step2: Init 14B t2i model...", end='', flush=True) - cfg = WAN_CONFIGS['t2i-14B'] - # cfg = WAN_CONFIGS['t2v-1.3B'] - wan_t2i = wan.WanT2V( - config=cfg, - checkpoint_dir=args.ckpt_dir, - device_id=0, - rank=0, - t5_fsdp=False, - dit_fsdp=False, - use_usp=False, - ) - print("done", flush=True) - - demo = gradio_interface() - demo.launch(server_name="0.0.0.0", share=False, server_port=7860) diff --git a/gradio/t2v_1.3B_singleGPU.py b/gradio/t2v_1.3B_singleGPU.py deleted file mode 100644 index 0a752d2..0000000 --- a/gradio/t2v_1.3B_singleGPU.py +++ /dev/null @@ -1,207 +0,0 @@ -# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. -import argparse -import os.path as osp -import os -import sys -import warnings - -import gradio as gr - -warnings.filterwarnings('ignore') - -# Model -sys.path.insert(0, os.path.sep.join(osp.realpath(__file__).split(os.path.sep)[:-2])) -import wan -from wan.configs import WAN_CONFIGS -from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander -from wan.utils.utils import cache_video - -# Global Var -prompt_expander = None -wan_t2v = None - - -# Button Func -def prompt_enc(prompt, tar_lang): - global prompt_expander - prompt_output = prompt_expander(prompt, tar_lang=tar_lang.lower()) - if prompt_output.status == False: - return prompt - else: - return prompt_output.prompt - - -def t2v_generation(txt2vid_prompt, resolution, sd_steps, guide_scale, - shift_scale, seed, n_prompt): - global wan_t2v - # print(f"{txt2vid_prompt},{resolution},{sd_steps},{guide_scale},{shift_scale},{seed},{n_prompt}") - - W = int(resolution.split("*")[0]) - H = int(resolution.split("*")[1]) - video = wan_t2v.generate( - txt2vid_prompt, - size=(W, H), - shift=shift_scale, - sampling_steps=sd_steps, - guide_scale=guide_scale, - n_prompt=n_prompt, - seed=seed, - offload_model=True) - - cache_video( - tensor=video[None], - save_file="example.mp4", - fps=16, - nrow=1, - normalize=True, - value_range=(-1, 1)) - - return "example.mp4" - - -# Interface -def gradio_interface(): - with gr.Blocks() as demo: - gr.Markdown(""" -
- Wan2.1 (T2V-1.3B) -
-
- Wan: Open and Advanced Large-Scale Video Generative Models. -
- """) - - with gr.Row(): - with gr.Column(): - txt2vid_prompt = gr.Textbox( - label="Prompt", - placeholder="Describe the video you want to generate", - ) - tar_lang = gr.Radio( - choices=["CH", "EN"], - label="Target language of prompt enhance", - value="CH") - run_p_button = gr.Button(value="Prompt Enhance") - - with gr.Accordion("Advanced Options", open=True): - resolution = gr.Dropdown( - label='Resolution(Width*Height)', - choices=[ - '480*832', - '832*480', - '624*624', - '704*544', - '544*704', - ], - value='480*832') - - with gr.Row(): - sd_steps = gr.Slider( - label="Diffusion steps", - minimum=1, - maximum=1000, - value=50, - step=1) - guide_scale = gr.Slider( - label="Guide scale", - minimum=0, - maximum=20, - value=6.0, - step=1) - with gr.Row(): - shift_scale = gr.Slider( - label="Shift scale", - minimum=0, - maximum=20, - value=8.0, - step=1) - seed = gr.Slider( - label="Seed", - minimum=-1, - maximum=2147483647, - step=1, - value=-1) - n_prompt = gr.Textbox( - label="Negative Prompt", - placeholder="Describe the negative prompt you want to add" - ) - - run_t2v_button = gr.Button("Generate Video") - - with gr.Column(): - result_gallery = gr.Video( - label='Generated Video', interactive=False, height=600) - - run_p_button.click( - fn=prompt_enc, - inputs=[txt2vid_prompt, tar_lang], - outputs=[txt2vid_prompt]) - - run_t2v_button.click( - fn=t2v_generation, - inputs=[ - txt2vid_prompt, resolution, sd_steps, guide_scale, shift_scale, - seed, n_prompt - ], - outputs=[result_gallery], - ) - - return demo - - -# Main -def _parse_args(): - parser = argparse.ArgumentParser( - description="Generate a video from a text prompt or image using Gradio") - parser.add_argument( - "--ckpt_dir", - type=str, - default="cache", - help="The path to the checkpoint directory.") - parser.add_argument( - "--prompt_extend_method", - type=str, - default="local_qwen", - choices=["dashscope", "local_qwen"], - help="The prompt extend method to use.") - parser.add_argument( - "--prompt_extend_model", - type=str, - default=None, - help="The prompt extend model to use.") - - args = parser.parse_args() - - return args - - -if __name__ == '__main__': - args = _parse_args() - - print("Step1: Init prompt_expander...", end='', flush=True) - if args.prompt_extend_method == "dashscope": - prompt_expander = DashScopePromptExpander( - model_name=args.prompt_extend_model, is_vl=False) - elif args.prompt_extend_method == "local_qwen": - prompt_expander = QwenPromptExpander( - model_name=args.prompt_extend_model, is_vl=False, device=0) - else: - raise NotImplementedError( - f"Unsupport prompt_extend_method: {args.prompt_extend_method}") - print("done", flush=True) - - print("Step2: Init 1.3B t2v model...", end='', flush=True) - cfg = WAN_CONFIGS['t2v-1.3B'] - wan_t2v = wan.WanT2V( - config=cfg, - checkpoint_dir=args.ckpt_dir, - device_id=0, - rank=0, - t5_fsdp=False, - dit_fsdp=False, - use_usp=False, - ) - print("done", flush=True) - - demo = gradio_interface() - demo.launch(server_name="0.0.0.0", share=False, server_port=7860) diff --git a/gradio/t2v_14B_singleGPU.py b/gradio/t2v_14B_singleGPU.py deleted file mode 100644 index 7e752d3..0000000 --- a/gradio/t2v_14B_singleGPU.py +++ /dev/null @@ -1,216 +0,0 @@ -# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. -import argparse -import os.path as osp -import os -import sys -import warnings - -import gradio as gr - -warnings.filterwarnings('ignore') - -# Model -sys.path.insert(0, os.path.sep.join(osp.realpath(__file__).split(os.path.sep)[:-2])) -import wan -from wan.configs import WAN_CONFIGS -from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander -from wan.utils.utils import cache_video - -# Global Var -prompt_expander = None -wan_t2v = None - - -# Button Func -def prompt_enc(prompt, tar_lang): - global prompt_expander - prompt_output = prompt_expander(prompt, tar_lang=tar_lang.lower()) - if prompt_output.status == False: - return prompt - else: - return prompt_output.prompt - - -def t2v_generation(txt2vid_prompt, resolution, sd_steps, guide_scale, - shift_scale, seed, n_prompt): - global wan_t2v - # print(f"{txt2vid_prompt},{resolution},{sd_steps},{guide_scale},{shift_scale},{seed},{n_prompt}") - - W = int(resolution.split("*")[0]) - H = int(resolution.split("*")[1]) - video = wan_t2v.generate( - txt2vid_prompt, - size=(W, H), - shift=shift_scale, - sampling_steps=sd_steps, - guide_scale=guide_scale, - n_prompt=n_prompt, - seed=seed, - offload_model=False) - - cache_video( - tensor=video[None], - save_file="example.mp4", - fps=16, - nrow=1, - normalize=True, - value_range=(-1, 1)) - - return "example.mp4" - - -# Interface -def gradio_interface(): - with gr.Blocks() as demo: - gr.Markdown(""" -
- Wan2.1 (T2V-14B) -
-
- Wan: Open and Advanced Large-Scale Video Generative Models. -
- """) - - with gr.Row(): - with gr.Column(): - txt2vid_prompt = gr.Textbox( - label="Prompt", - placeholder="Describe the video you want to generate", - ) - tar_lang = gr.Radio( - choices=["CH", "EN"], - label="Target language of prompt enhance", - value="CH") - run_p_button = gr.Button(value="Prompt Enhance") - - with gr.Accordion("Advanced Options", open=True): - resolution = gr.Dropdown( - label='Resolution(Width*Height)', - choices=[ - '720*1280', '1280*720', '960*960', '1088*832', - '832*1088', '480*832', '832*480', '624*624', - '704*544', '544*704' - ], - value='720*1280') - - with gr.Row(): - sd_steps = gr.Slider( - label="Diffusion steps", - minimum=1, - maximum=1000, - value=50, - step=1) - guide_scale = gr.Slider( - label="Guide scale", - minimum=0, - maximum=20, - value=5.0, - step=1) - with gr.Row(): - shift_scale = gr.Slider( - label="Shift scale", - minimum=0, - maximum=10, - value=5.0, - step=1) - seed = gr.Slider( - label="Seed", - minimum=-1, - maximum=2147483647, - step=1, - value=-1) - n_prompt = gr.Textbox( - label="Negative Prompt", - placeholder="Describe the negative prompt you want to add" - ) - - run_t2v_button = gr.Button("Generate Video") - - with gr.Column(): - result_gallery = gr.Video( - label='Generated Video', interactive=False, height=600) - - run_p_button.click( - fn=prompt_enc, - inputs=[txt2vid_prompt, tar_lang], - outputs=[txt2vid_prompt]) - - run_t2v_button.click( - fn=t2v_generation, - inputs=[ - txt2vid_prompt, resolution, sd_steps, guide_scale, shift_scale, - seed, n_prompt - ], - outputs=[result_gallery], - ) - - return demo - - -# Main -def _parse_args(): - parser = argparse.ArgumentParser( - description="Generate a video from a text prompt or image using Gradio") - parser.add_argument( - "--ckpt_dir", - type=str, - default="cache", - help="The path to the checkpoint directory.") - parser.add_argument( - "--prompt_extend_method", - type=str, - default="local_qwen", - choices=["dashscope", "local_qwen"], - help="The prompt extend method to use.") - parser.add_argument( - "--prompt_extend_model", - type=str, - default=None, - help="The prompt extend model to use.") - - args = parser.parse_args() - - return args - - -if __name__ == '__main__': - args = _parse_args() - - print("Step1: Init prompt_expander...", end='', flush=True) - prompt_expander = None - # if args.prompt_extend_method == "dashscope": - # prompt_expander = DashScopePromptExpander( - # model_name=args.prompt_extend_model, is_vl=False) - # elif args.prompt_extend_method == "local_qwen": - # prompt_expander = QwenPromptExpander( - # model_name=args.prompt_extend_model, is_vl=False, device=0) - # else: - # raise NotImplementedError( - # f"Unsupport prompt_extend_method: {args.prompt_extend_method}") - # print("done", flush=True) - - from mmgp import offload - - print("Step2: Init 14B t2v model...", end='', flush=True) - cfg = WAN_CONFIGS['t2v-14B'] - # cfg = WAN_CONFIGS['t2v-1.3B'] - - wan_t2v = wan.WanT2V( - config=cfg, - checkpoint_dir="../ckpts", - device_id=0, - rank=0, - t5_fsdp=False, - dit_fsdp=False, - use_usp=False, - ) - - pipe = {"transformer": wan_t2v.model, "text_encoder" : wan_t2v.text_encoder.model, "vae": wan_t2v.vae.model } # - # offload.profile(pipe, profile_no=4, budgets = {"transformer":100, "*":3000}, verboseLevel=2, quantizeTransformer = False, compile = "transformer") # - offload.profile(pipe, profile_no=4, budgets = {"transformer":100, "*":3000}, verboseLevel=2, quantizeTransformer = False) # - # offload.profile(pipe, profile_no=4, budgets = {"transformer":3000, "*":3000}, verboseLevel=2, quantizeTransformer = False) - - print("done", flush=True) - - demo = gradio_interface() - demo.launch(server_name="0.0.0.0", share=False, server_port=7860) diff --git a/preprocessing/matanyone/app.py b/preprocessing/matanyone/app.py index a58fa67..dd4b4ea 100644 --- a/preprocessing/matanyone/app.py +++ b/preprocessing/matanyone/app.py @@ -24,6 +24,7 @@ from .matanyone_wrapper import matanyone arg_device = "cuda" arg_sam_model_type="vit_h" arg_mask_save = False +model_loaded = False model = None matanyone_model = None @@ -409,36 +410,42 @@ def restart(): gr.update(visible=False), gr.update(visible=False, choices=[], value=[]), "", gr.update(visible=False) def load_unload_models(selected): + global model_loaded global model global matanyone_model if selected: - # args, defined in track_anything.py - sam_checkpoint_url_dict = { - 'vit_h': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth", - 'vit_l': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth", - 'vit_b': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth" - } - # os.path.join('.') + if model_loaded: + model.samcontroler.sam_controler.model.to(arg_device) + matanyone_model.to(arg_device) + else: + # args, defined in track_anything.py + sam_checkpoint_url_dict = { + 'vit_h': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth", + 'vit_l': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth", + 'vit_b': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth" + } + # os.path.join('.') - from mmgp import offload + from mmgp import offload - # sam_checkpoint = load_file_from_url(sam_checkpoint_url_dict[arg_sam_model_type], ".") - sam_checkpoint = None + # sam_checkpoint = load_file_from_url(sam_checkpoint_url_dict[arg_sam_model_type], ".") + sam_checkpoint = None - transfer_stream = torch.cuda.Stream() - with torch.cuda.stream(transfer_stream): - # initialize sams - model = MaskGenerator(sam_checkpoint, "cuda") - from .matanyone.model.matanyone import MatAnyone - matanyone_model = MatAnyone.from_pretrained("PeiqingYang/MatAnyone") - # pipe ={"mat" : matanyone_model, "sam" :model.samcontroler.sam_controler.model } - # offload.profile(pipe) - matanyone_model = matanyone_model.to(arg_device).eval() - matanyone_processor = InferenceCore(matanyone_model, cfg=matanyone_model.cfg) + transfer_stream = torch.cuda.Stream() + with torch.cuda.stream(transfer_stream): + # initialize sams + model = MaskGenerator(sam_checkpoint, arg_device) + from .matanyone.model.matanyone import MatAnyone + matanyone_model = MatAnyone.from_pretrained("PeiqingYang/MatAnyone") + # pipe ={"mat" : matanyone_model, "sam" :model.samcontroler.sam_controler.model } + # offload.profile(pipe) + matanyone_model = matanyone_model.to(arg_device).eval() + matanyone_processor = InferenceCore(matanyone_model, cfg=matanyone_model.cfg) + model_loaded = True else: import gc - model = None - matanyone_model = None + model.samcontroler.sam_controler.model.to("cpu") + matanyone_model.to("cpu") gc.collect() torch.cuda.empty_cache() @@ -451,10 +458,13 @@ def export_to_vace_video_input(foreground_video_output): return "V#" + str(time.time()), foreground_video_output def export_to_vace_video_mask(foreground_video_output, alpha_video_output): - gr.Info("Masked Video Input and Full Mask transferred to Vace For Stronger Inpainting") + gr.Info("Masked Video Input and Full Mask transferred to Vace For Inpainting") return "MV#" + str(time.time()), foreground_video_output, alpha_video_output -def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger): +def teleport_to_vace(): + return gr.Tabs(selected="video_gen"), gr.Dropdown(value="vace_1.3B") + +def display(tabs, model_choice, vace_video_input, vace_video_mask, video_prompt_video_guide_trigger): # my_tab.select(fn=load_unload_models, inputs=[], outputs=[]) media_url = "https://github.com/pq-yang/MatAnyone/releases/download/media/" @@ -576,18 +586,23 @@ def display(vace_video_input, vace_video_mask, video_prompt_video_guide_trigger) gr.Markdown("") # output video - with gr.Row(equal_height=True) as output_row: - with gr.Column(scale=2): - foreground_video_output = gr.Video(label="Masked Video Output", visible=False, elem_classes="video") - foreground_output_button = gr.Button(value="Black & White Video Output", visible=False, elem_classes="new_button") - export_to_vace_video_input_btn = gr.Button("Export to Vace Video Input Video For Inpainting", visible= False) - with gr.Column(scale=2): - alpha_video_output = gr.Video(label="B & W Mask Video Output", visible=False, elem_classes="video") - alpha_output_button = gr.Button(value="Alpha Mask Output", visible=False, elem_classes="new_button") - export_to_vace_video_mask_btn = gr.Button("Export to Vace Video Input and Video Mask for stronger Inpainting", visible= False) + with gr.Column() as output_row: #equal_height=True + with gr.Row(): + with gr.Column(scale=2): + foreground_video_output = gr.Video(label="Masked Video Output", visible=False, elem_classes="video") + foreground_output_button = gr.Button(value="Black & White Video Output", visible=False, elem_classes="new_button") + with gr.Column(scale=2): + alpha_video_output = gr.Video(label="B & W Mask Video Output", visible=False, elem_classes="video") + alpha_output_button = gr.Button(value="Alpha Mask Output", visible=False, elem_classes="new_button") + with gr.Row(): + with gr.Row(visible= False): + export_to_vace_video_input_btn = gr.Button("Export to Vace Video Input Video For Inpainting", visible= False) + with gr.Row(visible= True): + export_to_vace_video_mask_btn = gr.Button("Export to Vace Video Input and Video Mask", visible= False) export_to_vace_video_input_btn.click(fn=export_to_vace_video_input, inputs= [foreground_video_output], outputs= [video_prompt_video_guide_trigger, vace_video_input]) - export_to_vace_video_mask_btn.click(fn=export_to_vace_video_mask, inputs= [foreground_video_output, alpha_video_output], outputs= [video_prompt_video_guide_trigger, vace_video_input, vace_video_mask]) + export_to_vace_video_mask_btn.click(fn=export_to_vace_video_mask, inputs= [foreground_video_output, alpha_video_output], outputs= [video_prompt_video_guide_trigger, vace_video_input, vace_video_mask]).then( + fn=teleport_to_vace, inputs=[], outputs=[tabs, model_choice]) # first step: get the video information extract_frames_button.click( fn=get_frames_from_video, diff --git a/requirements.txt b/requirements.txt index 0b90776..b061536 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,7 +16,7 @@ gradio>=5.0.0 numpy>=1.23.5,<2 einops moviepy==1.0.3 -mmgp==3.3.4 +mmgp==3.4.0 peft==0.14.0 mutagen decord @@ -25,7 +25,6 @@ rembg[gpu]==2.0.65 matplotlib timm segment-anything -ffmpeg-python omegaconf hydra-core # rembg==2.0.65 \ No newline at end of file diff --git a/wan/image2video.py b/wan/image2video.py index ed08d44..b688676 100644 --- a/wan/image2video.py +++ b/wan/image2video.py @@ -48,7 +48,6 @@ class WanI2V: self, config, checkpoint_dir, - device_id=0, rank=0, t5_fsdp=False, dit_fsdp=False, @@ -58,6 +57,8 @@ class WanI2V: i2v720p= True, model_filename ="", text_encoder_filename="", + quantizeTransformer = False, + dtype = torch.bfloat16 ): r""" Initializes the image-to-video generation model components. @@ -82,22 +83,22 @@ class WanI2V: Enable initializing Transformer Model on CPU. Only works without FSDP or USP. init_on_cpu (`bool`, *optional*, defaults to True): """ - self.device = torch.device(f"cuda:{device_id}") + self.device = torch.device(f"cuda") self.config = config self.rank = rank self.use_usp = use_usp self.t5_cpu = t5_cpu - + self.dtype = dtype self.num_train_timesteps = config.num_train_timesteps self.param_dtype = config.param_dtype - shard_fn = partial(shard_model, device_id=device_id) + # shard_fn = partial(shard_model, device_id=device_id) self.text_encoder = T5EncoderModel( text_len=config.text_len, dtype=config.t5_dtype, device=torch.device('cpu'), checkpoint_path=text_encoder_filename, tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), - shard_fn=shard_fn if t5_fsdp else None, + shard_fn=None, ) self.vae_stride = config.vae_stride @@ -116,34 +117,16 @@ class WanI2V: logging.info(f"Creating WanModel from {model_filename}") from mmgp import offload - self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel, writable_tensors= False) #forcedConfigPath= "ckpts/config2.json", + self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel,do_quantize= quantizeTransformer, writable_tensors= False) + if self.dtype == torch.float16 and not "fp16" in model_filename: + self.model.to(self.dtype) + # offload.save_model(self.model, "i2v_720p_fp16.safetensors",do_quantize=True) + if self.dtype == torch.float16: + self.vae.model.to(self.dtype) + # offload.save_model(self.model, "wan2.1_Fun_InP_1.3B_bf16_bis.safetensors") self.model.eval().requires_grad_(False) - if t5_fsdp or dit_fsdp or use_usp: - init_on_cpu = False - - if use_usp: - from xfuser.core.distributed import \ - get_sequence_parallel_world_size - - from .distributed.xdit_context_parallel import (usp_attn_forward, - usp_dit_forward) - for block in self.model.blocks: - block.self_attn.forward = types.MethodType( - usp_attn_forward, block.self_attn) - self.model.forward = types.MethodType(usp_dit_forward, self.model) - self.sp_size = get_sequence_parallel_world_size() - else: - self.sp_size = 1 - - # if dist.is_initialized(): - # dist.barrier() - # if dit_fsdp: - # self.model = shard_fn(self.model) - # else: - # if not init_on_cpu: - # self.model.to(self.device) self.sample_neg_prompt = config.sample_neg_prompt @@ -229,16 +212,15 @@ class WanI2V: w = lat_w * self.vae_stride[2] clip_image_size = self.clip.model.image_size - img_interpolated = resize_lanczos(img, h, w).sub_(0.5).div_(0.5).unsqueeze(0).transpose(0,1).to(self.device) + img_interpolated = resize_lanczos(img, h, w).sub_(0.5).div_(0.5).unsqueeze(0).transpose(0,1).to(self.device, self.dtype) img = resize_lanczos(img, clip_image_size, clip_image_size) - img = img.sub_(0.5).div_(0.5).to(self.device) + img = img.sub_(0.5).div_(0.5).to(self.device, self.dtype) if img2!= None: - img_interpolated2 = resize_lanczos(img2, h, w).sub_(0.5).div_(0.5).unsqueeze(0).transpose(0,1).to(self.device) + img_interpolated2 = resize_lanczos(img2, h, w).sub_(0.5).div_(0.5).unsqueeze(0).transpose(0,1).to(self.device, self.dtype) img2 = resize_lanczos(img2, clip_image_size, clip_image_size) - img2 = img2.sub_(0.5).div_(0.5).to(self.device) + img2 = img2.sub_(0.5).div_(0.5).to(self.device, self.dtype) max_seq_len = lat_frames * lat_h * lat_w // ( self.patch_size[1] * self.patch_size[2]) - max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size seed = seed if seed >= 0 else random.randint(0, sys.maxsize) seed_g = torch.Generator(device=self.device) @@ -275,6 +257,9 @@ class WanI2V: context = [t.to(self.device) for t in context] context_null = [t.to(self.device) for t in context_null] + context = [u.to(self.dtype) for u in context] + context_null = [u.to(self.dtype) for u in context_null] + clip_context = self.clip.visual([img[:, None, :, :]]) if offload_model: self.clip.model.cpu() @@ -285,13 +270,13 @@ class WanI2V: mean2 = 0 enc= torch.concat([ img_interpolated, - torch.full( (3, frame_num-2, h, w), mean2, device=self.device, dtype= torch.bfloat16), + torch.full( (3, frame_num-2, h, w), mean2, device=self.device, dtype= self.dtype), img_interpolated2, ], dim=1).to(self.device) else: enc= torch.concat([ img_interpolated, - torch.zeros(3, frame_num-1, h, w, device=self.device, dtype= torch.bfloat16) + torch.zeros(3, frame_num-1, h, w, device=self.device, dtype= self.dtype) ], dim=1).to(self.device) lat_y = self.vae.encode([enc], VAE_tile_size, any_end_frame= any_end_frame and add_frames_for_end_image)[0] @@ -447,7 +432,7 @@ class WanI2V: callback(i, False) - x0 = [latent.to(self.device, dtype=torch.bfloat16)] + x0 = [latent.to(self.device, dtype=self.dtype)] if offload_model: self.model.cpu() diff --git a/wan/modules/attention.py b/wan/modules/attention.py index b6764bb..b795b06 100644 --- a/wan/modules/attention.py +++ b/wan/modules/attention.py @@ -5,6 +5,11 @@ from mmgp import offload import torch.nn.functional as F +try: + from xformers.ops import memory_efficient_attention +except ImportError: + memory_efficient_attention = None + try: import flash_attn_interface FLASH_ATTN_3_AVAILABLE = True @@ -123,13 +128,13 @@ def get_attention_modes(): ret = ["sdpa", "auto"] if flash_attn != None: ret.append("flash") - # if memory_efficient_attention != None: - # ret.append("xformers") + if memory_efficient_attention != None: + ret.append("xformers") if sageattn_varlen_wrapper != None: ret.append("sage") if sageattn != None and version("sageattention").startswith("2") : ret.append("sage2") - + return ret def get_supported_attention_modes(): @@ -338,6 +343,14 @@ def pay_attention( deterministic=deterministic).unflatten(0, (b, lq)) # output + + elif attn=="xformers": + x = memory_efficient_attention( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + ) #.unsqueeze(0) + return x.type(out_dtype) diff --git a/wan/modules/model.py b/wan/modules/model.py index 5af4ae8..0ba16ae 100644 --- a/wan/modules/model.py +++ b/wan/modules/model.py @@ -77,73 +77,6 @@ def rope_params_riflex(max_seq_len, dim, theta=10000, L_test=30, k=6): - -def rope_apply_(x, grid_sizes, freqs): - assert x.shape[0]==1 - - n, c = x.size(2), x.size(3) // 2 - - # split freqs - freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) - - f, h, w = grid_sizes[0] - seq_len = f * h * w - x_i = x[0, :seq_len, :, :] - - x_i = x_i.to(torch.float32) - x_i = x_i.reshape(seq_len, n, -1, 2) - x_i = torch.view_as_complex(x_i) - freqs_i = torch.cat([ - freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), - freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), - freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) - ], dim=-1) - freqs_i= freqs_i.reshape(seq_len, 1, -1) - - # apply rotary embedding - x_i *= freqs_i - x_i = torch.view_as_real(x_i).flatten(2) - x[0, :seq_len, :, :] = x_i.to(torch.bfloat16) - # x_i = torch.cat([x_i, x[0, seq_len:]]) - return x - -# @amp.autocast(enabled=False) -def rope_apply(x, grid_sizes, freqs): - n, c = x.size(2), x.size(3) // 2 - - # split freqs - freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) - - # loop over samples - output = [] - for i, (f, h, w) in enumerate(grid_sizes): - seq_len = f * h * w - - # precompute multipliers - # x_i = x[i, :seq_len] - x_i = x[i] - x_i = x_i[:seq_len, :, :] - - x_i = x_i.to(torch.float32) - x_i = x_i.reshape(seq_len, n, -1, 2) - x_i = torch.view_as_complex(x_i) - freqs_i = torch.cat([ - freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), - freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), - freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) - ], - dim=-1).reshape(seq_len, 1, -1) - - # apply rotary embedding - x_i *= freqs_i - x_i = torch.view_as_real(x_i).flatten(2) - x_i = x_i.to(torch.bfloat16) - x_i = torch.cat([x_i, x[i, seq_len:]]) - - # append to collection - output.append(x_i) - return torch.stack(output) #.float() - def relative_l1_distance(last_tensor, current_tensor): l1_distance = torch.abs(last_tensor - current_tensor).mean() norm = torch.abs(last_tensor).mean() @@ -256,8 +189,6 @@ class WanSelfAttention(nn.Module): k = k.view(b, s, n, d) v = self.v(x).view(b, s, n, d) del x - # rope_apply_(q, grid_sizes, freqs) - # rope_apply_(k, grid_sizes, freqs) qklist = [q,k] del q,k q,k = apply_rotary_emb(qklist, freqs, head_first=False) @@ -568,9 +499,9 @@ class Head(nn.Module): e(Tensor): Shape [B, C] """ # assert e.dtype == torch.float32 - + dtype = x.dtype e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1) - x = self.norm(x).to(torch.bfloat16) + x = self.norm(x).to(dtype) x *= (1 + e[1]) x += e[0] x = self.head(x) @@ -857,7 +788,7 @@ class WanModel(ModelMixin, ConfigMixin): # time embeddings e = self.time_embedding( sinusoidal_embedding_1d(self.freq_dim, t)) - e0 = self.time_projection(e).unflatten(1, (6, self.dim)).to(torch.bfloat16) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)).to(e.dtype) # context context_lens = None diff --git a/wan/modules/vae.py b/wan/modules/vae.py index 67dcd9a..8be1e6f 100644 --- a/wan/modules/vae.py +++ b/wan/modules/vae.py @@ -51,10 +51,11 @@ class RMS_norm(nn.Module): self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. def forward(self, x): + dtype = x.dtype x = F.normalize( x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias - x = x.to(torch.bfloat16) + x = x.to(dtype) return x class Upsample(nn.Upsample): @@ -208,6 +209,7 @@ class ResidualBlock(nn.Module): def forward(self, x, feat_cache=None, feat_idx=[0]): h = self.shortcut(x) + dtype = x.dtype for layer in self.residual: if isinstance(layer, CausalConv3d) and feat_cache is not None: idx = feat_idx[0] @@ -219,11 +221,11 @@ class ResidualBlock(nn.Module): cache_x.device), cache_x ], dim=2) - x = layer(x, feat_cache[idx]).to(torch.bfloat16) + x = layer(x, feat_cache[idx]).to(dtype) feat_cache[idx] = cache_x#.to("cpu") feat_idx[0] += 1 else: - x = layer(x).to(torch.bfloat16) + x = layer(x).to(dtype) return x + h @@ -323,6 +325,7 @@ class Encoder3d(nn.Module): CausalConv3d(out_dim, z_dim, 3, padding=1)) def forward(self, x, feat_cache=None, feat_idx=[0]): + dtype = x.dtype if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() @@ -333,7 +336,7 @@ class Encoder3d(nn.Module): cache_x.device), cache_x ], dim=2) - x = self.conv1(x, feat_cache[idx]).to(torch.bfloat16) + x = self.conv1(x, feat_cache[idx]).to(dtype) feat_cache[idx] = cache_x del cache_x feat_idx[0] += 1 diff --git a/wan/text2video.py b/wan/text2video.py index 035700a..b8140f1 100644 --- a/wan/text2video.py +++ b/wan/text2video.py @@ -47,14 +47,15 @@ class WanT2V: self, config, checkpoint_dir, - device_id=0, rank=0, t5_fsdp=False, dit_fsdp=False, use_usp=False, t5_cpu=False, model_filename = None, - text_encoder_filename = None + text_encoder_filename = None, + quantizeTransformer = False, + dtype = torch.bfloat16 ): r""" Initializes the Wan text-to-video generation model components. @@ -77,25 +78,24 @@ class WanT2V: t5_cpu (`bool`, *optional*, defaults to False): Whether to place T5 model on CPU. Only works without t5_fsdp. """ - self.device = torch.device(f"cuda:{device_id}") + self.device = torch.device(f"cuda") self.config = config self.rank = rank self.t5_cpu = t5_cpu - + self.dtype = dtype self.num_train_timesteps = config.num_train_timesteps self.param_dtype = config.param_dtype - shard_fn = partial(shard_model, device_id=device_id) self.text_encoder = T5EncoderModel( text_len=config.text_len, dtype=config.t5_dtype, device=torch.device('cpu'), checkpoint_path=text_encoder_filename, tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), - shard_fn=shard_fn if t5_fsdp else None) + shard_fn= None) self.vae_stride = config.vae_stride - self.patch_size = config.patch_size + self.patch_size = config.patch_size self.vae = WanVAE( @@ -105,31 +105,14 @@ class WanT2V: logging.info(f"Creating WanModel from {model_filename}") from mmgp import offload - - self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel, writable_tensors= False) - + self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel,do_quantize= quantizeTransformer, writable_tensors= False) + if self.dtype == torch.float16 and not "fp16" in model_filename: + self.model.to(self.dtype) + # offload.save_model(self.model, "t2v_fp16.safetensors",do_quantize=True) + if self.dtype == torch.float16: + self.vae.model.to(self.dtype) self.model.eval().requires_grad_(False) - if use_usp: - from xfuser.core.distributed import \ - get_sequence_parallel_world_size - - from .distributed.xdit_context_parallel import (usp_attn_forward, - usp_dit_forward) - for block in self.model.blocks: - block.self_attn.forward = types.MethodType( - usp_attn_forward, block.self_attn) - self.model.forward = types.MethodType(usp_dit_forward, self.model) - self.sp_size = get_sequence_parallel_world_size() - else: - self.sp_size = 1 - - # if dist.is_initialized(): - # dist.barrier() - # if dit_fsdp: - # self.model = shard_fn(self.model) - # else: - # self.model.to(self.device) self.sample_neg_prompt = config.sample_neg_prompt @@ -389,8 +372,10 @@ class WanT2V: 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 + target_shape[1]) + context = [u.to(self.dtype) for u in context] + context_null = [u.to(self.dtype) for u in context_null] noise = [ torch.randn( diff --git a/wgp.py b/wgp.py index d0dfc6d..b1a790b 100644 --- a/wgp.py +++ b/wgp.py @@ -20,7 +20,7 @@ from wan.modules.attention import get_attention_modes, get_supported_attention_m import torch import gc import traceback -import math +import math import typing import asyncio import inspect @@ -32,6 +32,8 @@ import zipfile import tempfile import atexit import shutil +import glob + global_queue_ref = [] AUTOSAVE_FILENAME = "queue.zip" PROMPT_VARS_MAX = 10 @@ -203,6 +205,7 @@ def process_prompt_and_add_tasks(state, model_choice): if isinstance(image_refs, list): image_refs = [ convert_image(tup[0]) for tup in image_refs ] + os.environ["U2NET_HOME"] = os.path.join(os.getcwd(), "ckpts", "rembg") from wan.utils.utils import resize_and_remove_background image_refs = resize_and_remove_background(image_refs, width, height, inputs["remove_background_image_ref"] ==1) @@ -921,7 +924,7 @@ def autoload_queue(state): update_global_queue_ref(original_queue) dataframe_update = update_queue_data(original_queue) else: - print(f"Autoload skipped: {AUTOSAVE_FILENAME} not found.") + # print(f"Autoload skipped: {AUTOSAVE_FILENAME} not found.") update_global_queue_ref([]) dataframe_update = update_queue_data([]) @@ -1090,19 +1093,13 @@ def _parse_args(): help="Lora preset to preload" ) - # parser.add_argument( - # "--i2v-settings", - # type=str, - # default="i2v_settings.json", - # help="Path to settings file for i2v" - # ) + parser.add_argument( + "--settings", + type=str, + default="settings", + help="Path to settings folder" + ) - # parser.add_argument( - # "--t2v-settings", - # type=str, - # default="t2v_settings.json", - # help="Path to settings file for t2v" - # ) # parser.add_argument( # "--lora-preset-i2v", @@ -1152,6 +1149,12 @@ def _parse_args(): help="Access advanced options by default" ) + parser.add_argument( + "--fp16", + action="store_true", + help="For using fp16 transformer model" + ) + parser.add_argument( "--server-port", type=str, @@ -1159,6 +1162,22 @@ def _parse_args(): help="Server port" ) + parser.add_argument( + "--theme", + type=str, + default="", + help="set UI Theme" + ) + + parser.add_argument( + "--perc-reserved-mem-max", + type=float, + default=0, + help="% of RAM allocated to Reserved RAM" + ) + + + parser.add_argument( "--server-name", type=str, @@ -1307,6 +1326,12 @@ transformer_choices_i2v=["ckpts/wan2.1_image2video_480p_14B_bf16.safetensors", " transformer_choices = transformer_choices_t2v + transformer_choices_i2v text_encoder_choices = ["ckpts/models_t5_umt5-xxl-enc-bf16.safetensors", "ckpts/models_t5_umt5-xxl-enc-quanto_int8.safetensors"] server_config_filename = "wgp_config.json" +if not os.path.isdir("settings"): + os.mkdir("settings") +if os.path.isfile("t2v_settings.json"): + for f in glob.glob(os.path.join(".", "*_settings.json*")): + target_file = os.path.join("settings", Path(f).parts[-1] ) + shutil.move(f, target_file) if not os.path.isfile(server_config_filename) and os.path.isfile("gradio_config.json"): shutil.move("gradio_config.json", server_config_filename) @@ -1321,10 +1346,11 @@ if not Path(server_config_filename).is_file(): "metadata_type": "metadata", "default_ui": "t2v", "boost" : 1, - "clear_file_list" : 0, + "clear_file_list" : 5, "vae_config": 0, "profile" : profile_type.LowRAM_LowVRAM, - "preload_model_policy": [] } + "preload_model_policy": [], + "UI_theme": "default" } with open(server_config_filename, "w", encoding="utf-8") as writer: writer.write(json.dumps(server_config)) @@ -1380,7 +1406,7 @@ def get_model_filename(model_type, quantization): return choices[0] def get_settings_file_name(model_filename): - return get_model_type(model_filename) + "_settings.json" + return os.path.join(args.settings, get_model_type(model_filename) + "_settings.json") def get_default_settings(filename): def get_default_prompt(i2v): @@ -1388,11 +1414,11 @@ def get_default_settings(filename): return "Several giant wooly mammoths approach treading through a snowy meadow, their long wooly fur lightly blows in the wind as they walk, snow covered trees and dramatic snow capped mountains in the distance, mid afternoon light with wispy clouds and a sun high in the distance creates a warm glow, the low camera view is stunning capturing the large furry mammal with beautiful photography, depth of field." else: return "A large orange octopus is seen resting on the bottom of the ocean floor, blending in with the sandy and rocky terrain. Its tentacles are spread out around its body, and its eyes are closed. The octopus is unaware of a king crab that is crawling towards it from behind a rock, its claws raised and ready to attack. The crab is brown and spiny, with long legs and antennae. The scene is captured from a wide angle, showing the vastness and depth of the ocean. The water is clear and blue, with rays of sunlight filtering through. The shot is sharp and crisp, with a high dynamic range. The octopus and the crab are in focus, while the background is slightly blurred, creating a depth of field effect." - i2v = "image2video" in file_name + i2v = "image2video" in filename or "Fun_InP" in filename defaults_filename = get_settings_file_name(filename) if not Path(defaults_filename).is_file(): ui_defaults = { - "prompts": get_default_prompt(i2v), + "prompt": get_default_prompt(i2v), "resolution": "832x480", "video_length": 81, "num_inference_steps": 30, @@ -1651,7 +1677,6 @@ def setup_loras(model_filename, transformer, lora_dir, lora_preselected_preset, if lora_dir != None: - import glob dir_loras = glob.glob( os.path.join(lora_dir , "*.sft") ) + glob.glob( os.path.join(lora_dir , "*.safetensors") ) dir_loras.sort() loras += [element for element in dir_loras if element not in loras ] @@ -1676,7 +1701,7 @@ def setup_loras(model_filename, transformer, lora_dir, lora_preselected_preset, return loras, loras_names, loras_presets, default_loras_choices, default_loras_multis_str, default_lora_preset_prompt, default_lora_preset -def load_t2v_model(model_filename, value): +def load_t2v_model(model_filename, value, quantizeTransformer = False, dtype = torch.bfloat16): cfg = WAN_CONFIGS['t2v-14B'] # cfg = WAN_CONFIGS['t2v-1.3B'] @@ -1685,20 +1710,21 @@ def load_t2v_model(model_filename, value): wan_model = wan.WanT2V( config=cfg, checkpoint_dir="ckpts", - device_id=0, rank=0, t5_fsdp=False, dit_fsdp=False, use_usp=False, model_filename=model_filename, - text_encoder_filename= text_encoder_filename + text_encoder_filename= text_encoder_filename, + quantizeTransformer = quantizeTransformer, + dtype = dtype ) pipe = {"transformer": wan_model.model, "text_encoder" : wan_model.text_encoder.model, "vae": wan_model.vae.model } return wan_model, pipe -def load_i2v_model(model_filename, value): +def load_i2v_model(model_filename, value, quantizeTransformer = False, dtype = torch.bfloat16): print(f"Loading '{model_filename}' model...") @@ -1707,14 +1733,15 @@ def load_i2v_model(model_filename, value): wan_model = wan.WanI2V( config=cfg, checkpoint_dir="ckpts", - device_id=0, rank=0, t5_fsdp=False, dit_fsdp=False, use_usp=False, i2v720p= True, model_filename=model_filename, - text_encoder_filename=text_encoder_filename + text_encoder_filename=text_encoder_filename, + quantizeTransformer = quantizeTransformer, + dtype = dtype ) pipe = {"transformer": wan_model.model, "text_encoder" : wan_model.text_encoder.model, "text_encoder_2": wan_model.clip.model, "vae": wan_model.vae.model } # @@ -1723,15 +1750,15 @@ def load_i2v_model(model_filename, value): wan_model = wan.WanI2V( config=cfg, checkpoint_dir="ckpts", - device_id=0, rank=0, t5_fsdp=False, dit_fsdp=False, use_usp=False, i2v720p= False, model_filename=model_filename, - text_encoder_filename=text_encoder_filename - + text_encoder_filename=text_encoder_filename, + quantizeTransformer = quantizeTransformer, + dtype = dtype ) pipe = {"transformer": wan_model.model, "text_encoder" : wan_model.text_encoder.model, "text_encoder_2": wan_model.clip.model, "vae": wan_model.vae.model } # else: @@ -1744,12 +1771,20 @@ def load_models(model_filename): global transformer_filename transformer_filename = model_filename + perc_reserved_mem_max = args.perc_reserved_mem_max + + major, minor = torch.cuda.get_device_capability(args.gpu if len(args.gpu) > 0 else None) + default_dtype = torch.float16 if major < 8 else torch.bfloat16 + if default_dtype == torch.float16 or args.fp16: + print("Switching to f16 model as GPU architecture doesn't support bf16") + if "quanto" in model_filename: + model_filename = model_filename.replace("quanto_int8", "quanto_fp16_int8") download_models(model_filename, text_encoder_filename) if test_class_i2v(model_filename): res720P = "720p" in model_filename - wan_model, pipe = load_i2v_model(model_filename, "720P" if res720P else "480P") + wan_model, pipe = load_i2v_model(model_filename, "720P" if res720P else "480P", quantizeTransformer = quantizeTransformer, dtype = default_dtype ) else: - wan_model, pipe = load_t2v_model(model_filename, "") + wan_model, pipe = load_t2v_model(model_filename, "", quantizeTransformer = quantizeTransformer, dtype = default_dtype) wan_model._model_file_name = model_filename kwargs = { "extraModelsToQuantize": None} if profile == 2 or profile == 4: @@ -1758,7 +1793,7 @@ def load_models(model_filename): # kwargs["partialPinning"] = True elif profile == 3: kwargs["budgets"] = { "*" : "70%" } - offloadobj = offload.profile(pipe, profile_no= profile, compile = compile, quantizeTransformer = quantizeTransformer, loras = "transformer", coTenantsMap= {}, **kwargs) + offloadobj = offload.profile(pipe, profile_no= profile, compile = compile, quantizeTransformer = quantizeTransformer, loras = "transformer", coTenantsMap= {}, perc_reserved_mem_max = perc_reserved_mem_max , convertWeightsFloatTo = default_dtype, **kwargs) if len(args.gpu) > 0: torch.set_default_device(args.gpu) @@ -1834,6 +1869,7 @@ def apply_changes( state, boost_choice = 1, clear_file_list = 0, preload_model_policy_choice = 1, + UI_theme_choice = "default" ): if args.lock_config: return @@ -1852,6 +1888,7 @@ def apply_changes( state, "boost" : boost_choice, "clear_file_list" : clear_file_list, "preload_model_policy" : preload_model_policy_choice, + "UI_theme" : UI_theme_choice } if Path(server_config_filename).is_file(): @@ -1874,7 +1911,7 @@ def apply_changes( state, if v != v_old: changes.append(k) - global attention_mode, profile, compile, transformer_filename, text_encoder_filename, vae_config, boost, lora_dir, reload_needed, preload_model_policy, transformer_quantization, transformer_types + global attention_mode, profile, compile, text_encoder_filename, vae_config, boost, lora_dir, reload_needed, preload_model_policy, transformer_quantization, transformer_types attention_mode = server_config["attention_mode"] profile = server_config["profile"] compile = server_config["compile"] @@ -1884,10 +1921,13 @@ def apply_changes( state, preload_model_policy = server_config["preload_model_policy"] transformer_quantization = server_config["transformer_quantization"] transformer_types = server_config["transformer_types"] - transformer_type = get_model_type(transformer_filename) - if not transformer_type in transformer_types: - transformer_type = transformer_types[0] if len(transformer_types) > 0 else model_types[0] - transformer_filename = get_model_filename(transformer_type, transformer_quantization) + model_filename = state["model_filename"] + model_transformer_type = get_model_type(model_filename) + + if not model_transformer_type in transformer_types: + model_transformer_type = transformer_types[0] if len(transformer_types) > 0 else model_types[0] + model_filename = get_model_filename(model_transformer_type, transformer_quantization) + state["model_filename"] = model_filename if all(change in ["attention_mode", "vae_config", "boost", "save_path", "metadata_choice", "clear_file_list"] for change in changes ): model_choice = gr.Dropdown() else: @@ -1990,6 +2030,15 @@ def refresh_gallery(state, msg): start_img_md = "" end_img_md = "" prompt = task["prompt"] + params = task["params"] + if "\n" in prompt and params.get("sliding_window_repeat", 0) > 0: + prompts = prompt.split("\n") + repeat_no= gen.get("repeat_no",1) + if repeat_no > len(prompts): + repeat_no = len(prompts) + repeat_no -= 1 + prompts[repeat_no]="" + prompts[repeat_no] + "" + prompt = "
".join(prompts) start_img_uri = task.get('start_image_data_base64') start_img_uri = start_img_uri[0] if start_img_uri !=None else None @@ -2463,15 +2512,7 @@ def generate_video( try: start_time = time.time() - # with tracker_lock: - # progress_tracker[task_id] = { - # 'current_step': 0, - # 'total_steps': num_inference_steps, - # 'start_time': start_time, - # 'last_update': start_time, - # 'repeats': repeat_generation, # f"{video_no}/{repeat_generation}", - # 'status': "Encoding Prompt" - # } + if trans.enable_teacache: trans.teacache_counter = 0 trans.num_steps = num_inference_steps @@ -2542,20 +2583,17 @@ def generate_video( gc.collect() torch.cuda.empty_cache() s = str(e) - keyword_list = ["vram", "VRAM", "memory","allocat"] - VRAM_crash= False - if any( keyword in s for keyword in keyword_list): - VRAM_crash = True - else: - stack = traceback.extract_stack(f=None, limit=5) - for frame in stack: - if any( keyword in frame.name for keyword in keyword_list): - VRAM_crash = True - break - + keyword_list = {"CUDA out of memory" : "VRAM", "Tried to allocate":"VRAM", "CUDA error: out of memory": "RAM", "CUDA error: too many resources requested": "RAM"} + crash_type = "" + for keyword, tp in keyword_list.items(): + if keyword in s: + crash_type = tp + break state["prompt"] = "" - if VRAM_crash: + if crash_type == "VRAM": 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." + elif crash_type == "RAM": + new_error = "The generation of the video has encountered an error: it is likely that you have unsufficient RAM and / or Reserved RAM allocation should be reduced using 'perc_reserved_mem_max' or using a different Profile." 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')[:-1] @@ -2929,12 +2967,13 @@ def refresh_lora_list(state, lset_name, loras_choices): pos = len(loras_presets) lset_name ="" - errors = getattr(wan_model.model, "_loras_errors", "") - if errors !=None and len(errors) > 0: - error_files = [path for path, _ in errors] - gr.Info("Error while refreshing Lora List, invalid Lora files: " + ", ".join(error_files)) - else: - gr.Info("Lora List has been refreshed") + if wan_model != None: + errors = getattr(wan_model.model, "_loras_errors", "") + if errors !=None and len(errors) > 0: + error_files = [path for path, _ in errors] + gr.Info("Error while refreshing Lora List, invalid Lora files: " + ", ".join(error_files)) + else: + gr.Info("Lora List has been refreshed") return gr.Dropdown(choices=lset_choices, value= lset_choices[pos][1]), gr.Dropdown(choices=new_loras_choices, value= lora_names_selected) @@ -3210,7 +3249,7 @@ def save_inputs( def download_loras(): from huggingface_hub import snapshot_download yield gr.Row(visible=True), "Please wait while the Loras are being downloaded", *[gr.Column(visible=False)] * 2 - lora_dir = get_lora_dir(get_model_filename("i2v"), quantizeTransformer) + lora_dir = get_lora_dir(get_model_filename("i2v", transformer_quantization)) log_path = os.path.join(lora_dir, "log.txt") if not os.path.isfile(log_path): tmp_path = os.path.join(lora_dir, "tmp_lora_dowload") @@ -4047,7 +4086,7 @@ def generate_video_tab(update_form = False, state_dict = None, ui_defaults = Non outputs=[modal_container] ) - return ( + return ( state, loras_choices, lset_name, state, queue_df, current_gen_column, gen_status, output, abort_btn, generate_btn, add_to_queue_btn, gen_info, queue_accordion, video_guide, video_mask, video_prompt_video_guide_trigger @@ -4068,9 +4107,7 @@ def generate_download_tab(lset_name,loras_choices, state): download_loras_btn.click(fn=download_loras, inputs=[], outputs=[download_status_row, download_status]).then(fn=refresh_lora_list, inputs=[state, lset_name,loras_choices], outputs=[lset_name, loras_choices]) -def generate_configuration_tab(header, model_choice): - state_dict = {} - state = gr.State(state_dict) +def generate_configuration_tab(state, blocks, header, model_choice): gr.Markdown("Please click Apply Changes at the bottom so that the changes are effective. Some choices below may be locked if the app has been launched by specifying a config preset.") with gr.Column(): model_list = [] @@ -4090,7 +4127,7 @@ def generate_configuration_tab(header, model_choice): quantization_choice = gr.Dropdown( choices=[ ("Int8 Quantization (recommended)", "int8"), - ("BF16 (no quantization)", "bf16"), + ("16 bits (no quantization)", "bf16"), ], value= transformer_quantization, label="Wan Transformer Model Quantization Type (if available)", @@ -4122,7 +4159,7 @@ def generate_configuration_tab(header, model_choice): ("Auto : pick sage2 > sage > sdpa depending on what is installed", "auto"), ("Scale Dot Product Attention: default, always available", "sdpa"), ("Flash" + check("flash")+ ": good quality - requires additional install (usually complex to set up on Windows without WSL)", "flash"), - # ("Xformers" + check("xformers")+ ": good quality - requires additional install (usually complex, may consume less VRAM to set up on Windows without WSL)", "xformers"), + ("Xformers" + check("xformers")+ ": good quality - requires additional install (usually complex, may consume less VRAM to set up on Windows without WSL)", "xformers"), ("Sage" + check("sage")+ ": 30% faster but slightly worse quality - requires additional install (usually complex to set up on Windows without WSL)", "sage"), ("Sage2" + check("sage2")+ ": 40% faster but slightly worse quality - requires additional install (usually complex to set up on Windows without WSL)", "sage2"), ], @@ -4201,10 +4238,19 @@ def generate_configuration_tab(header, model_choice): ("Keep the last 20 videos", 20), ("Keep the last 30 videos", 30), ], - value=server_config.get("clear_file_list", 0), + value=server_config.get("clear_file_list", 5), label="Keep Previously Generated Videos when starting a Generation Batch" ) + UI_theme_choice = gr.Dropdown( + choices=[ + ("Blue Sky", "default"), + ("Classic Gradio", "gradio"), + ], + value=server_config.get("UI_theme_choice", "default"), + label="User Interface Theme. You will need to restart the App the see new Theme." + ) + msg = gr.Markdown() apply_btn = gr.Button("Apply Changes") @@ -4224,6 +4270,7 @@ def generate_configuration_tab(header, model_choice): boost_choice, clear_file_list_choice, preload_model_policy_choice, + UI_theme_choice ], outputs= [msg , header, model_choice] ) @@ -4286,20 +4333,15 @@ def select_tab(tab_state, evt:gr.SelectData): elif new_tab_no == tab_video_mask_creator: if gen_in_progress: gr.Info("Unable to access this Tab while a Generation is in Progress. Please come back later") - tab_state["tab_auto"]=old_tab_no + tab_state["tab_no"] = 0 + return gr.Tabs(selected="video_gen") else: vmc_event_handler(True) tab_state["tab_no"] = new_tab_no -def select_tab_auto(tab_state): - old_tab_no = tab_state.pop("tab_auto", -1) - if old_tab_no>= 0: - tab_state["tab_auto"]=old_tab_no - return gr.Tabs(selected=old_tab_no) # !! doesnt work !! - return gr.Tab() - + return gr.Tabs() def create_demo(): - global vmc_event_handler + global vmc_event_handler css = """ #model_list{ background-color:black; @@ -4532,14 +4574,21 @@ def create_demo(): pointer-events: none; } """ - with gr.Blocks(css=css, theme=gr.themes.Soft(font=["Verdana"], primary_hue="sky", neutral_hue="slate", text_size="md"), title= "Wan2GP") as demo: + UI_theme = server_config.get("UI_theme", "default") + UI_theme = args.theme if len(args.theme) > 0 else UI_theme + if UI_theme == "gradio": + theme = None + else: + theme = gr.themes.Soft(font=["Verdana"], primary_hue="sky", neutral_hue="slate", text_size="md") + + with gr.Blocks(css=css, theme=theme, title= "Wan2GP") as demo: gr.Markdown("

WanGP v4.0 by DeepBeepMeep ") # (Updates)

") global model_list tab_state = gr.State({ "tab_no":0 }) with gr.Tabs(selected="video_gen", ) as main_tabs: - with gr.Tab("Video Generator", id="video_gen") as t2v_tab: + with gr.Tab("Video Generator", id="video_gen"): with gr.Row(): if args.lock_model: gr.Markdown("

" + get_model_name(transformer_filename) + "

") @@ -4551,23 +4600,23 @@ def create_demo(): with gr.Row(): header = gr.Markdown(generate_header(transformer_filename, compile, attention_mode), visible= True) with gr.Row(): - ( + ( state, loras_choices, lset_name, state, queue_df, current_gen_column, gen_status, output, abort_btn, generate_btn, add_to_queue_btn, gen_info, queue_accordion, video_guide, video_mask, video_prompt_type_video_trigger ) = generate_video_tab(model_choice=model_choice, header=header) - with gr.Tab("Informations"): + with gr.Tab("Informations", id="info"): generate_info_tab() with gr.Tab("Video Mask Creator", id="video_mask_creator") as video_mask_creator: from preprocessing.matanyone import app as matanyone_app vmc_event_handler = matanyone_app.get_vmc_event_handler() - matanyone_app.display(video_guide, video_mask, video_prompt_type_video_trigger) + matanyone_app.display(main_tabs, model_choice, video_guide, video_mask, video_prompt_type_video_trigger) if not args.lock_config: with gr.Tab("Downloads", id="downloads") as downloads_tab: generate_download_tab(lset_name, loras_choices, state) - with gr.Tab("Configuration"): - generate_configuration_tab(header, model_choice) + with gr.Tab("Configuration", id="configuration"): + generate_configuration_tab(state, demo, header, model_choice) with gr.Tab("About"): generate_about_tab() @@ -4589,7 +4638,7 @@ def create_demo(): trigger_mode="always_last" ) - main_tabs.select(fn=select_tab, inputs= [tab_state], outputs= None).then(fn=select_tab_auto, inputs= [tab_state], outputs=[main_tabs]) + main_tabs.select(fn=select_tab, inputs= [tab_state], outputs= main_tabs) return demo if __name__ == "__main__": From 5b6b3f28322924c7f3f1f617801467a44b07a20e Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Tue, 15 Apr 2025 09:00:22 +0200 Subject: [PATCH 68/69] updated mmgp version check --- wgp.py | 50 +++++++++++++++++++++++++------------------------- 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/wgp.py b/wgp.py index 0af6127..09dca96 100644 --- a/wgp.py +++ b/wgp.py @@ -40,7 +40,7 @@ global_queue_ref = [] AUTOSAVE_FILENAME = "queue.zip" PROMPT_VARS_MAX = 10 -target_mmgp_version = "3.3.4" +target_mmgp_version = "3.4.0" from importlib.metadata import version mmgp_version = version("mmgp") if mmgp_version != target_mmgp_version: @@ -52,29 +52,29 @@ task_id = 0 # progress_tracker = {} # tracker_lock = threading.Lock() -def download_ffmpeg(): - if os.name != 'nt': return - exes = ['ffmpeg.exe', 'ffprobe.exe', 'ffplay.exe'] - if all(os.path.exists(e) for e in exes): return - api_url = 'https://api.github.com/repos/GyanD/codexffmpeg/releases/latest' - r = requests.get(api_url, headers={'Accept': 'application/vnd.github+json'}) - assets = r.json().get('assets', []) - zip_asset = next((a for a in assets if 'essentials_build.zip' in a['name']), None) - if not zip_asset: return - zip_url = zip_asset['browser_download_url'] - zip_name = zip_asset['name'] - with requests.get(zip_url, stream=True) as resp: - total = int(resp.headers.get('Content-Length', 0)) - with open(zip_name, 'wb') as f, tqdm(total=total, unit='B', unit_scale=True) as pbar: - for chunk in resp.iter_content(chunk_size=8192): - f.write(chunk) - pbar.update(len(chunk)) - with zipfile.ZipFile(zip_name) as z: - for f in z.namelist(): - if f.endswith(tuple(exes)) and '/bin/' in f: - z.extract(f) - os.rename(f, os.path.basename(f)) - os.remove(zip_name) +# def download_ffmpeg(): +# if os.name != 'nt': return +# exes = ['ffmpeg.exe', 'ffprobe.exe', 'ffplay.exe'] +# if all(os.path.exists(e) for e in exes): return +# api_url = 'https://api.github.com/repos/GyanD/codexffmpeg/releases/latest' +# r = requests.get(api_url, headers={'Accept': 'application/vnd.github+json'}) +# assets = r.json().get('assets', []) +# zip_asset = next((a for a in assets if 'essentials_build.zip' in a['name']), None) +# if not zip_asset: return +# zip_url = zip_asset['browser_download_url'] +# zip_name = zip_asset['name'] +# with requests.get(zip_url, stream=True) as resp: +# total = int(resp.headers.get('Content-Length', 0)) +# with open(zip_name, 'wb') as f, tqdm(total=total, unit='B', unit_scale=True) as pbar: +# for chunk in resp.iter_content(chunk_size=8192): +# f.write(chunk) +# pbar.update(len(chunk)) +# with zipfile.ZipFile(zip_name) as z: +# for f in z.namelist(): +# if f.endswith(tuple(exes)) and '/bin/' in f: +# z.extract(f) +# os.rename(f, os.path.basename(f)) +# os.remove(zip_name) def format_time(seconds): if seconds < 60: @@ -4676,7 +4676,7 @@ def create_demo(): if __name__ == "__main__": atexit.register(autosave_queue) - download_ffmpeg() + # download_ffmpeg() # threading.Thread(target=runner, daemon=True).start() os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" server_port = int(args.server_port) From c8a538b42a1ae94eec1217d696bb7b61cdb7b642 Mon Sep 17 00:00:00 2001 From: DeepBeepMeep Date: Tue, 15 Apr 2025 13:42:45 +0200 Subject: [PATCH 69/69] removed ffmpeg dependency --- preprocessing/matanyone/app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/preprocessing/matanyone/app.py b/preprocessing/matanyone/app.py index dd4b4ea..f92dab9 100644 --- a/preprocessing/matanyone/app.py +++ b/preprocessing/matanyone/app.py @@ -4,7 +4,7 @@ import os import json import time import psutil -import ffmpeg +# import ffmpeg import imageio from PIL import Image