diff --git a/fastvideo/v1/sample/v1_fastvideo_inference.py b/fastvideo/v1/sample/v1_fastvideo_inference.py index 6a3667a64..69698f5bd 100644 --- a/fastvideo/v1/sample/v1_fastvideo_inference.py +++ b/fastvideo/v1/sample/v1_fastvideo_inference.py @@ -13,7 +13,9 @@ from fastvideo.v1.inference_args import InferenceArgs from fastvideo.v1.inference_args import prepare_inference_args from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel from fastvideo.v1.logger import logger - +from fastvideo.v1.models.loader.fsdp_load import get_param_names_mapping +from safetensors.torch import safe_open, save_file +import shutil def initialize_distributed_and_parallelism(inference_args: InferenceArgs): local_rank = int(os.environ.get("LOCAL_RANK", 0)) @@ -32,6 +34,110 @@ def initialize_distributed_and_parallelism(inference_args: InferenceArgs): sequence_model_parallel_size=inference_args.sp_size, tensor_model_parallel_size=inference_args.tp_size, ) + +def convert_and_save_lora_weights(input_folder, output_folder, lora_param_mapping_fn): + + # Check if the input folder exists + if not os.path.exists(input_folder): + raise ValueError(f"Input folder {input_folder} not found") + + # Create output folder if it doesn't exist + os.makedirs(output_folder, exist_ok=True) + + # Process safetensors file + safetensors_file = os.path.join(input_folder, "pytorch_lora_weights.safetensors") + if os.path.exists(safetensors_file): + # Load all tensors from the safetensors file + tensors = {} + with safe_open(safetensors_file, framework="pt", device="cpu") as f: + tensor_names = f.keys() + print(f"Found {len(tensor_names)} tensors in {safetensors_file}") + + for name in tensor_names: + tensors[name] = f.get_tensor(name) + + # Extract base names (without LoRA suffix) and LoRA suffixes + base_names_dict = {} + lora_suffixes = {} + + for name in tensor_names: + if name.endswith(".lora_A.weight") or name.endswith(".lora_B.weight"): + for suffix in [".lora_A.weight", ".lora_B.weight"]: + if name.endswith(suffix): + base_name = name[:-len(suffix)] + ".weight" + base_names_dict[name] = base_name + lora_suffixes[name] = suffix + break + else: + base_names_dict[name] = name + lora_suffixes[name] = "" + + # Apply the mapping function to get all mappings at once + name_mappings = {} + for full_name, base_name in base_names_dict.items(): + try: + # First remove "transformer." prefix if it exists + has_transformer_prefix = False + processed_base_name = base_name + if base_name.startswith("transformer."): + processed_base_name = base_name[12:] # Remove "transformer." (12 characters) + has_transformer_prefix = True + + # Apply the mapping function to the processed name + target_base_name, merge_index, total_splitted_params = lora_param_mapping_fn(processed_base_name) + + # Add back the "transformer." prefix if it was removed + if target_base_name is not None and has_transformer_prefix: + target_base_name = "transformer." + target_base_name + + if target_base_name is not None: + name_mappings[full_name] = (target_base_name, merge_index, total_splitted_params) + except Exception as e: + print(f"Error mapping parameter {base_name}: {e}") + + # Create the converted weights using the mappings + converted_weights = {} + for name, tensor in tensors.items(): + suffix = lora_suffixes.get(name, "") + + if name in name_mappings: + target_base_name, merge_index, _ = name_mappings[name] + + # Reconstruct full parameter name with LoRA suffix + if target_base_name.endswith('.weight'): + target_base_name_without_suffix = target_base_name[:-7] # 移除 .weight + new_key = target_base_name_without_suffix + suffix + else: + new_key = target_base_name + suffix + + if merge_index is not None: + new_key += f"_{merge_index}" + + converted_weights[new_key] = tensor + else: + # Keep original name if no mapping was found + converted_weights[name] = tensor + + print(f"Converted {len(converted_weights)} parameters") + + # Save the converted weights + output_safetensors_file = os.path.join(output_folder, "pytorch_lora_weights.safetensors") + save_file(converted_weights, output_safetensors_file) + print(f"Converted weights saved to {output_safetensors_file}") + else: + print(f"Warning: No safetensors file found in {input_folder}") + + # Copy other files (like config and optimizer state) without modification + for filename in os.listdir(input_folder): + if filename != "pytorch_lora_weights.safetensors": + input_path = os.path.join(input_folder, filename) + output_path = os.path.join(output_folder, filename) + + if os.path.isfile(input_path): + shutil.copy2(input_path, output_path) + print(f"Copied {filename} to output folder") + + print(f"LoRA conversion complete: {input_folder} → {output_folder}") def main(inference_args: InferenceArgs): initialize_distributed_and_parallelism(inference_args) @@ -44,14 +150,40 @@ def main(inference_args: InferenceArgs): prompts = [line.strip() for line in f.readlines()] else: prompts = [inference_args.prompt] - + # from IPython import embed; embed() + + # convert lora format + input_path = "data/Hunyuan-Black-Myth-Wukong-lora-weight/" + output_path = "data/Hunyuan-Black-Myth-Wukong-lora-weight_converted" + # from IPython import embed; embed() + lora_param_mapping_fn = get_param_names_mapping(engine.pipeline.transformer._param_names_mapping) + converted_weights = convert_and_save_lora_weights( + input_path, + output_path, + lora_param_mapping_fn, + ) + # Process each prompt for prompt in prompts: + lora_checkpoint = "data/Hunyuan-Black-Myth-Wukong-lora-weight_converted/" + if lora_checkpoint: + import json + print(f"Loading LoRA weights from lora_checkpoint: {lora_checkpoint}") + config_path = os.path.join(lora_checkpoint, "lora_config.json") + with open(config_path, "r") as f: + lora_config_dict = json.load(f) + rank = lora_config_dict["lora_params"]["lora_rank"] + lora_alpha = lora_config_dict["lora_params"]["lora_alpha"] + lora_scaling = lora_alpha / rank + engine.pipeline.load_lora_weights(lora_checkpoint, adapter_name="default") + from IPython import embed; embed() + engine.pipeline.set_adapters(["default"], [lora_scaling]) + print(f"Successfully Loaded LoRA weights from {lora_checkpoint}") outputs = engine.run( prompt=prompt, inference_args=inference_args, ) - + # img_attn_qkv, img_attn_proj, linear1, linear2 # Process outputs videos = rearrange(outputs["samples"], "b c t h w -> t b c h w") frames = []