nodes/aigc-apps-VideoX-Fun
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Support validations in all models && Update Saving Code && Preparing Code (#452)

Browse Source
This commit is contained in:
Bubbliiiing
2026-02-10 17:23:23 +08:00
committed by GitHub
parent ee44fbc950
commit 1ca4162447
92 changed files with 5118 additions and 6279 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/cogvideox_fun/train.py
+150 -173
View File
@@ -58,28 +58,30 @@ for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
from videox_fun.data.dataset_image_video import (ImageVideoControlDataset,
ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
from videox_fun.models import (AutoencoderKLCogVideoX,
CogVideoXTransformer3DModel, T5EncoderModel,
T5Tokenizer)
from videox_fun.pipeline import (CogVideoXFunPipeline,
CogVideoXFunControlPipeline,
CogVideoXFunInpaintPipeline)
CogVideoXTransformer3DModel, T5EncoderModel,
T5Tokenizer)
from videox_fun.pipeline import (CogVideoXFunControlPipeline,
CogVideoXFunInpaintPipeline,
CogVideoXFunPipeline)
from videox_fun.pipeline.pipeline_cogvideox_fun_inpaint import (
add_noise_to_reference_video, get_3d_rotary_pos_embed,
get_resize_crop_region_for_grid)
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import create_network, merge_lora, unmerge_lora
from videox_fun.utils.utils import (get_image_to_video_latent,
get_video_to_video_latent, save_videos_grid)
from videox_fun.utils.lora_utils import (create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -160,116 +162,106 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = CogVideoXTransformer3DModel.from_pretrained(
args.pretrained_model_name_or_path, subfolder="transformer"
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = DDIMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = DDIMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
if args.train_mode != "normal":
pipeline = CogVideoXFunInpaintPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
else:
pipeline = CogVideoXFunPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.train_mode != "normal":
pipeline = CogVideoXFunInpaintPipeline.from_pretrained(
args.pretrained_model_name_or_path,
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
torch_dtype=weight_dtype
)
else:
pipeline = CogVideoXFunPipeline.from_pretrained(
args.pretrained_model_name_or_path,
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
torch_dtype=weight_dtype
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -341,6 +333,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -905,7 +904,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -925,26 +924,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1771,25 +1750,24 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1797,25 +1775,24 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/cogvideox_fun/train.sh
+2 -2
View File
@@ -25,7 +25,7 @@ accelerate launch --mixed_precision="bf16" scripts/cogvideox_fun/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_cog" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -65,7 +65,7 @@ accelerate launch --mixed_precision="bf16" scripts/cogvideox_fun/train.py \
# --lr_scheduler="constant_with_warmup" \
# --lr_warmup_steps=100 \
# --seed=42 \
# --output_dir="output_dir" \
# --output_dir="output_dir_cog" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
scripts/cogvideox_fun/train_control.py
+117 -118
View File
@@ -35,7 +35,7 @@ from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.state import AcceleratorState
from accelerate.utils import ProjectConfiguration, set_seed
from diffusers import AutoencoderKL, DDPMScheduler
from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler
from diffusers.optimization import get_scheduler
from diffusers.training_utils import EMAModel
from diffusers.utils import check_min_version, deprecate, is_wandb_available
@@ -56,26 +56,28 @@ for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
from videox_fun.data.dataset_image_video import (ImageVideoControlDataset,
ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
from videox_fun.models import (AutoencoderKLCogVideoX,
CogVideoXTransformer3DModel, T5EncoderModel,
T5Tokenizer)
from videox_fun.pipeline import (CogVideoXFunPipeline,
CogVideoXFunControlPipeline,
CogVideoXFunInpaintPipeline)
CogVideoXTransformer3DModel, T5EncoderModel,
T5Tokenizer)
from videox_fun.pipeline import (CogVideoXFunControlPipeline,
CogVideoXFunInpaintPipeline,
CogVideoXFunPipeline)
from videox_fun.pipeline.pipeline_cogvideox_fun_inpaint import (
add_noise_to_reference_video, get_3d_rotary_pos_embed,
get_resize_crop_region_for_grid)
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import (get_image_to_video_latent,
get_video_to_video_latent, save_videos_grid)
from videox_fun.utils.utils import (calculate_dimensions,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -156,60 +158,79 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = CogVideoXTransformer3DModel.from_pretrained(
args.pretrained_model_name_or_path, subfolder="transformer"
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = DDIMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
pipeline = CogVideoXFunControlPipeline.from_pretrained(
args.pretrained_model_name_or_path,
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
torch_dtype=weight_dtype,
)
pipeline = pipeline.to(accelerator.device)
pipeline = CogVideoXFunControlPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
for i in range(len(args.validation_prompts)):
import cv2
cap = cv2.VideoCapture(args.validation_paths[i])
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
cap.release()
control_video = input_video,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "The video is not of a high quality, it has a low resolution. Watermark present in each frame. The background is solid. Strange body and strange trajectory. Distortion. ",
height = height,
width = width,
generator = generator,
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
control_video = input_video,
num_inference_steps = 25,
guidance_scale = 6.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -843,7 +864,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -863,26 +884,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1658,25 +1659,24 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1684,25 +1684,24 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/cogvideox_fun/train_control.sh
+2 -2
View File
@@ -25,7 +25,7 @@ accelerate launch --mixed_precision="bf16" scripts/cogvideox_fun/train_control.p
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=50 \
--seed=43 \
--output_dir="output_dir" \
--output_dir="output_dir_cog_control" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -64,7 +64,7 @@ accelerate launch --mixed_precision="bf16" scripts/cogvideox_fun/train_control.p
# --lr_scheduler="constant_with_warmup" \
# --lr_warmup_steps=50 \
# --seed=43 \
# --output_dir="output_dir" \
# --output_dir="output_dir_cog_control" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
scripts/cogvideox_fun/train_lora.py
+129 -172
View File
@@ -78,7 +78,8 @@ from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -161,116 +162,108 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = CogVideoXTransformer3DModel.from_pretrained(
args.pretrained_model_name_or_path, subfolder="transformer",
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = DDIMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = DDIMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
if args.train_mode != "normal":
pipeline = CogVideoXFunInpaintPipeline.from_pretrained(
args.pretrained_model_name_or_path,
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
torch_dtype=weight_dtype,
)
else:
pipeline = CogVideoXFunPipeline.from_pretrained(
args.pretrained_model_name_or_path,
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
torch_dtype=weight_dtype
)
if args.train_mode != "normal":
pipeline = CogVideoXFunInpaintPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
else:
pipeline = CogVideoXFunPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = pipeline.to(accelerator.device)
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
for i in range(len(args.validation_prompts)):
with torch.no_grad():
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 7,
generator = generator,
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -342,6 +335,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -916,7 +916,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -945,15 +945,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -966,34 +958,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1301,23 +1271,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=transformer3d.transformer_blocks)
transformer3d = shard_fn(transformer3d)
# Move text_encode and vae to gpu and cast to weight_dtype
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
@@ -1841,19 +1800,18 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1861,19 +1819,18 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/cogvideox_fun/train_lora.sh
+2 -2
View File
@@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/cogvideox_fun/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_cog_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -65,7 +65,7 @@ accelerate launch --mixed_precision="bf16" scripts/cogvideox_fun/train_lora.py \
# --checkpointing_steps=50 \
# --learning_rate=1e-04 \
# --seed=42 \
# --output_dir="output_dir" \
# --output_dir="output_dir_cog_lora" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
scripts/fantasytalking/train.py
+123 -124
View File
@@ -70,13 +70,18 @@ from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
from videox_fun.data.dataset_image_video import ImageVideoSampler, get_random_mask
from videox_fun.data.dataset_image_video import (ImageVideoSampler,
get_random_mask)
from videox_fun.data.dataset_video import VideoSpeechDataset
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel, FantasyTalkingAudioEncoder,
FantasyTalkingTransformer3DModel)
from videox_fun.models import (AutoencoderKLWan, CLIPModel,
FantasyTalkingAudioEncoder,
FantasyTalkingTransformer3DModel,
WanT5EncoderModel)
from videox_fun.pipeline import FantasyTalkingPipeline, WanFunPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -143,73 +148,86 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, audio_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
transformer3d_val = FantasyTalkingTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = FantasyTalkingPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
audio_encoder=audio_encoder,
)
pipeline = pipeline.to(accelerator.device)
pipeline = FantasyTalkingPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
for i in range(len(args.validation_prompts)):
start_image = Image.open(args.validation_image_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.video_sample_size * args.video_sample_size, width / height)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, clip_image = get_image_to_video_latent(args.validation_image_paths[i], None, video_length=video_length, sample_size=[height, width])
audio_path = args.validation_audio_paths[i]
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, clip_image = get_image_to_video_latent(args.validation_image_paths[i], None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
audio_path = args.validation_audio_paths[i]
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
video = input_video,
mask_video = input_video_mask,
clip_image = clip_image,
audio_path = audio_path,
shift = 5,
fps = 16
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
video = input_video,
mask_video = input_video_mask,
clip_image = clip_image,
audio_path = audio_path,
shift = 5,
fps = 16
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -895,7 +913,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -915,26 +933,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1350,8 +1348,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -1836,27 +1835,27 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
audio_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1864,27 +1863,27 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
audio_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/fantasytalking/train.sh
+1 -1
View File
@@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/fantasytalking/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_fantasytalking" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/flux/train.py
+3 -23
View File
@@ -262,7 +262,7 @@ def log_validation(vae, text_encoder, text_encoder_2, tokenizer, tokenizer_2, tr
text_encoder_2=text_encoder_2,
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -926,7 +926,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -946,26 +946,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1302,7 +1282,7 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
scripts/flux/train.sh
+1 -1
View File
@@ -20,7 +20,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_flux" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/flux/train_lora.py
+8 -51
View File
@@ -265,7 +265,7 @@ def log_validation(vae, text_encoder, text_encoder_2, tokenizer, tokenizer_2, tr
text_encoder_2=text_encoder_2,
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -929,7 +929,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -958,15 +958,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -979,34 +971,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1291,23 +1261,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=list(transformer3d.transformer_blocks) + list(transformer3d.single_transformer_blocks))
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -1315,8 +1274,6 @@ def main():
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.text_model.encoder.layers)
text_encoder = shard_fn(text_encoder)
# shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=text_encoder_2.encoder.block)
# text_encoder_2 = shard_fn(text_encoder_2)
# Move text_encode and vae to gpu and cast to weight_dtype
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
scripts/flux/train_lora.sh
+1 -1
View File
@@ -18,7 +18,7 @@ accelerate launch --mixed_precision="bf16" scripts/flux/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_flux_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/flux2/train.py
+2 -22
View File
@@ -329,7 +329,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -1018,7 +1018,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1038,26 +1038,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
scripts/flux2/train.sh
+1 -1
View File
@@ -20,7 +20,7 @@ accelerate launch --mixed_precision="bf16" scripts/flux2/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_flux2" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/flux2/train_lora.py
+8 -49
View File
@@ -332,7 +332,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, a
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -990,7 +990,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1019,15 +1019,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -1040,34 +1032,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1348,23 +1318,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=list(transformer3d.transformer_blocks) + list(transformer3d.single_transformer_blocks))
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
scripts/flux2/train_lora.sh
+1 -1
View File
@@ -18,7 +18,7 @@ accelerate launch --mixed_precision="bf16" scripts/flux2/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_lora" \
--output_dir="output_dir_flux2_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/flux2_fun/train_control.py
+2 -22
View File
@@ -333,7 +333,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -1040,7 +1040,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1060,26 +1060,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
scripts/flux2_fun/train_control_distill.py
+2 -22
View File
@@ -337,7 +337,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -1046,7 +1046,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1066,26 +1066,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
scripts/flux2_fun/train_control_distill.sh
+1 -1
View File
@@ -21,7 +21,7 @@ accelerate launch --mixed_precision="bf16" scripts/flux2_fun/train_control_disti
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_flux2_control_CFG_Distill" \
--output_dir="output_dir_flux2_control_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/hunyuanvideo/train.py
+104 -113
View File
@@ -167,67 +167,80 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, text_encoder_2, tokenizer, tokenizer_2, transformer3d, args, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = HunyuanVideoTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, 'transformer'),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler"
)
if args.train_mode != "normal":
raise NotImplementedError("train_mode is not implemented")
else:
pipeline = HunyuanVideoPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
text_encoder_2=accelerator.unwrap_model(text_encoder_2),
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
transformer=transformer3d_val,
scheduler=scheduler,
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler"
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
if args.train_mode != "normal":
raise NotImplementedError("train_mode is not implemented")
else:
pipeline = HunyuanVideoPipeline(
vae=vae,
text_encoder=text_encoder,
text_encoder_2=text_encoder_2,
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
raise NotImplementedError("train_mode is not implemented")
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
text_encoder_2.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
text_encoder_2.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
DEFAULT_PROMPT_TEMPLATE = {
"template": (
@@ -1026,7 +1039,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1046,26 +1059,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -2054,27 +2047,26 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
text_encoder_2,
tokenizer,
tokenizer_2,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
text_encoder_2,
tokenizer,
tokenizer_2,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2082,27 +2074,26 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
text_encoder_2,
tokenizer,
tokenizer_2,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
text_encoder_2,
tokenizer,
tokenizer_2,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/hunyuanvideo/train.sh
+1 -1
View File
@@ -25,7 +25,7 @@ accelerate launch --mixed_precision="bf16" scripts/hunyuanvideo/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_hunyuanvideo" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/hunyuanvideo/train_lora.py
+110 -131
View File
@@ -171,70 +171,81 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, text_encoder_2, tokenizer, tokenizer_2, transformer3d, network, args, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = HunyuanVideoTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, 'transformer'),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler"
)
if args.train_mode != "normal":
raise NotImplementedError("train_mode is not implemented")
else:
pipeline = HunyuanVideoPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
text_encoder_2=accelerator.unwrap_model(text_encoder_2),
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
transformer=transformer3d_val,
scheduler=scheduler,
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler"
)
pipeline = pipeline.to(accelerator.device)
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
if args.train_mode != "normal":
raise NotImplementedError("train_mode is not implemented")
else:
pipeline = HunyuanVideoPipeline(
vae=vae,
text_encoder=text_encoder,
text_encoder_2=text_encoder_2,
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
raise NotImplementedError("train_mode is not implemented")
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
text_encoder_2.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
text_encoder_2.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
DEFAULT_PROMPT_TEMPLATE = {
"template": (
@@ -1034,7 +1045,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1063,15 +1074,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -1084,34 +1087,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1429,22 +1410,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=list(transformer3d.transformer_blocks) + list(transformer3d.single_transformer_blocks))
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -1534,6 +1505,16 @@ def main():
else:
initial_global_step = 0
# function for saving/removing
def save_model(ckpt_file, unwrapped_nw):
os.makedirs(args.output_dir, exist_ok=True)
accelerator.print(f"\nsaving checkpoint: {ckpt_file}")
if isinstance(unwrapped_nw, dict):
from safetensors.torch import save_file
save_file(unwrapped_nw, ckpt_file, metadata={"format": "pt"})
return ckpt_file
unwrapped_nw.save_weights(ckpt_file, weight_dtype, None)
progress_bar = tqdm(
range(0, args.max_train_steps),
initial=initial_global_step,
@@ -2041,21 +2022,20 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
text_encoder_2,
tokenizer,
tokenizer_2,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
text_encoder_2,
tokenizer,
tokenizer_2,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2063,21 +2043,20 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
text_encoder_2,
tokenizer,
tokenizer_2,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
text_encoder_2,
tokenizer,
tokenizer_2,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/hunyuanvideo/train_lora.sh
+1 -1
View File
@@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/hunyuanvideo/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_lora" \
--output_dir="output_dir_hunyuanvideo_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/longcatvideo/train.py
+109 -188
View File
@@ -47,7 +47,8 @@ from einops import rearrange
from packaging import version
from PIL import Image
from torch.distributed.fsdp.fully_sharded_data_parallel import (
FullOptimStateDictConfig, FullStateDictConfig, ShardedStateDictConfig, ShardedOptimStateDictConfig)
FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig,
ShardedStateDictConfig)
from torch.utils.data import RandomSampler
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
@@ -62,20 +63,6 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
from videox_fun.data.dataset_image_video import (ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
LongCatVideoTransformer3DModel)
from videox_fun.pipeline import WanPipeline, WanI2VPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
@@ -84,9 +71,11 @@ from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
from videox_fun.data.dataset_image_video import (ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
from videox_fun.models import (AutoencoderKLLongCatVideo, CLIPModel, UMT5EncoderModel,
LongCatVideoTransformer3DModel)
from videox_fun.pipeline import WanI2VPipeline, WanPipeline
from videox_fun.models import (AutoencoderKLLongCatVideo, AutoencoderKLWan,
CLIPModel, LongCatVideoTransformer3DModel,
UMT5EncoderModel, WanT5EncoderModel)
from videox_fun.pipeline import (LongCatVideoPipeline, WanI2VPipeline,
WanPipeline)
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
@@ -177,117 +166,71 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = LongCatVideoTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, "dit"),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler"
)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler"
)
pipeline = LongCatVideoPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
else:
pipeline = WanPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.scale_factor_temporal * vae.config.scale_factor_temporal) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
for i in range(len(args.validation_prompts)):
sample = pipeline(
prompt = args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -974,7 +917,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -994,26 +937,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1068,8 +991,8 @@ def main():
if args.gradient_checkpointing:
transformer3d.enable_gradient_checkpointing()
elif args.selective_ac > 0:
from videox_fun.utils.ac_handle import apply_checkpointing, partial
from videox_fun.models.wan_transformer3d import WanAttentionBlock
from videox_fun.utils.ac_handle import apply_checkpointing, partial
apply_selective_ac = partial(apply_checkpointing, block=WanAttentionBlock)
apply_selective_ac(transformer3d, p=args.selective_ac)
@@ -1350,16 +1273,17 @@ def main():
new_examples['text'],
max_length=args.tokenizer_max_length,
padding="max_length",
add_special_tokens=True,
truncation=True,
add_special_tokens=True,
return_attention_mask=True,
return_tensors="pt"
)
encoder_hidden_states = text_encoder(
prompt_ids.input_ids, attention_mask=prompt_ids.attention_mask.to(latents.device)
)[0]
prompt_ids.input_ids, attention_mask=prompt_ids.attention_mask
).last_hidden_state
encoder_hidden_states = encoder_hidden_states.unsqueeze(1)
new_examples['encoder_attention_mask'] = prompt_ids.attention_mask
new_examples['encoder_hidden_states'] = encoder_hidden_states
new_examples['encoder_attention_mask'] = prompt_ids.attention_mask
return new_examples
@@ -1403,6 +1327,7 @@ def main():
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.encoder.block)
text_encoder = shard_fn(text_encoder)
@@ -1689,13 +1614,13 @@ def main():
max_length=args.tokenizer_max_length,
truncation=True,
add_special_tokens=True,
return_attention_mask=True,
return_tensors="pt"
)
text_input_ids = prompt_ids.input_ids
prompt_attention_mask = prompt_ids.attention_mask
text_input_ids = prompt_ids.input_ids.to(latents.device)
prompt_attention_mask = prompt_ids.attention_mask.to(latents.device)
seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
prompt_embeds = text_encoder(text_input_ids.to(latents.device), attention_mask=prompt_attention_mask.to(latents.device))[0]
prompt_embeds = text_encoder(text_input_ids, attention_mask=prompt_attention_mask).last_hidden_state
prompt_embeds = prompt_embeds.unsqueeze(1)
if args.low_vram and not args.enable_text_encoder_in_dataloader:
@@ -1850,26 +1775,24 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1877,26 +1800,24 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/longcatvideo/train.sh
+1 -1
View File
@@ -25,7 +25,7 @@ accelerate launch --mixed_precision="bf16" scripts/longcatvideo/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_longcat_full_finetune" \
--output_dir="output_dir_longcat" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/longcatvideo/train_lora.py
+100 -179
View File
@@ -68,9 +68,10 @@ from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
from videox_fun.data.dataset_image_video import (ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
from videox_fun.models import (AutoencoderKLLongCatVideo, CLIPModel, UMT5EncoderModel,
LongCatVideoTransformer3DModel)
from videox_fun.pipeline import WanI2VPipeline, WanPipeline
from videox_fun.models import (AutoencoderKLLongCatVideo, CLIPModel,
LongCatVideoTransformer3DModel,
UMT5EncoderModel)
from videox_fun.pipeline import LongCatVideoPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
@@ -161,118 +162,73 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, config, args, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = LongCatVideoTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, "dit"),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler"
)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler"
)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
pipeline = LongCatVideoPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
else:
pipeline = WanPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = pipeline.to(accelerator.device)
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
for i in range(len(args.validation_prompts)):
sample = pipeline(
prompt = args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.scale_factor_temporal * vae.config.scale_factor_temporal) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -951,7 +907,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -980,7 +936,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -993,34 +949,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1298,16 +1232,17 @@ def main():
new_examples['text'],
max_length=args.tokenizer_max_length,
padding="max_length",
add_special_tokens=True,
truncation=True,
add_special_tokens=True,
return_attention_mask=True,
return_tensors="pt"
)
encoder_hidden_states = text_encoder(
prompt_ids.input_ids, attention_mask=prompt_ids.attention_mask.to(latents.device)
)[0]
prompt_ids.input_ids, attention_mask=prompt_ids.attention_mask
).last_hidden_state
encoder_hidden_states = encoder_hidden_states.unsqueeze(1)
new_examples['encoder_attention_mask'] = prompt_ids.attention_mask
new_examples['encoder_hidden_states'] = encoder_hidden_states
new_examples['encoder_attention_mask'] = prompt_ids.attention_mask
return new_examples
@@ -1349,26 +1284,16 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.encoder.block)
text_encoder = shard_fn(text_encoder)
@@ -1719,13 +1644,13 @@ def main():
max_length=args.tokenizer_max_length,
truncation=True,
add_special_tokens=True,
return_attention_mask=True,
return_tensors="pt"
)
text_input_ids = prompt_ids.input_ids
prompt_attention_mask = prompt_ids.attention_mask
text_input_ids = prompt_ids.input_ids.to(latents.device)
prompt_attention_mask = prompt_ids.attention_mask.to(latents.device)
seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
prompt_embeds = text_encoder(text_input_ids.to(latents.device), attention_mask=prompt_attention_mask.to(latents.device))[0]
prompt_embeds = text_encoder(text_input_ids, attention_mask=prompt_attention_mask).last_hidden_state
prompt_embeds = prompt_embeds.unsqueeze(1)
if args.low_vram and not args.enable_text_encoder_in_dataloader:
@@ -1871,20 +1796,18 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1892,20 +1815,18 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/longcatvideo/train_lora.sh
+1 -1
View File
@@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/longcatvideo/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_longcat_lora" \
--output_dir="output_dir_longcatvideo_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/qwenimage/train.py
+3 -23
View File
@@ -151,7 +151,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -827,7 +827,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -847,26 +847,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1220,7 +1200,7 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers)
scripts/qwenimage/train.sh
+1 -1
View File
@@ -20,7 +20,7 @@ accelerate launch --mixed_precision="bf16" scripts/qwenimage/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_qwenimage" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/qwenimage/train_edit.py
+4 -24
View File
@@ -155,7 +155,7 @@ def log_validation(vae, text_encoder, tokenizer, processor, transformer3d, args,
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
processor=processor,
scheduler=scheduler,
)
@@ -164,7 +164,7 @@ def log_validation(vae, text_encoder, tokenizer, processor, transformer3d, args,
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
processor=processor,
scheduler=scheduler,
)
@@ -863,7 +863,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -883,26 +883,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1269,7 +1249,7 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers)
scripts/qwenimage/train_edit.sh
+1 -1
View File
@@ -20,7 +20,7 @@ accelerate launch --mixed_precision="bf16" scripts/qwenimage/train_edit.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_qwenimage_edit" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/qwenimage/train_edit_lora.py
+9 -50
View File
@@ -162,7 +162,7 @@ def log_validation(vae, text_encoder, tokenizer, processor, transformer3d, netwo
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
processor=processor,
scheduler=scheduler,
)
@@ -171,7 +171,7 @@ def log_validation(vae, text_encoder, tokenizer, processor, transformer3d, netwo
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
processor=processor,
scheduler=scheduler,
)
@@ -870,7 +870,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -899,15 +899,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -920,34 +912,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1262,23 +1232,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=transformer3d.transformer_blocks)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
scripts/qwenimage/train_edit_lora.sh
+1 -1
View File
@@ -18,7 +18,7 @@ accelerate launch --mixed_precision="bf16" scripts/qwenimage/train_edit_lora.py
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_lora" \
--output_dir="output_dir_qwenimage_edit_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/qwenimage/train_lora.py
+8 -49
View File
@@ -150,7 +150,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, a
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -826,7 +826,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -855,15 +855,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -876,34 +868,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1205,23 +1175,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=transformer3d.transformer_blocks)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
scripts/qwenimage/train_lora.sh
+1 -1
View File
@@ -18,7 +18,7 @@ accelerate launch --mixed_precision="bf16" scripts/qwenimage/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_lora" \
--output_dir="output_dir_qwenimage_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/qwenimage_fun/train_control.py
+3 -23
View File
@@ -151,7 +151,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -857,7 +857,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -877,26 +877,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1273,7 +1253,7 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import shard_model
scripts/qwenimage_instantx/train_control.py
+3 -23
View File
@@ -151,7 +151,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, cn_transformer,
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
controlnet=cn_transformer,
scheduler=scheduler,
)
@@ -865,7 +865,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -885,26 +885,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1315,7 +1295,7 @@ def main():
cn_transformer, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import shard_model
scripts/turbodiffusion/train_distill.py
+132 -153
View File
@@ -81,7 +81,9 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel,
WanTransformer3DModel)
from videox_fun.pipeline import WanI2VPipeline, WanPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -169,120 +171,109 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
else:
pipeline = WanPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video = input_video,
mask_video = input_video_mask,
num_inference_steps = len(args.denoising_step_list),
guidance_scale = 1.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = len(args.denoising_step_indices_list),
guidance_scale = 1.0,
).videos
video = input_video,
mask_video = input_video_mask,
num_inference_steps = len(args.denoising_step_list),
guidance_scale = 1.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = len(args.denoising_step_list),
guidance_scale = 1.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = len(args.denoising_step_indices_list),
guidance_scale = 1.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -354,6 +345,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--negative_prompt",
type=str,
@@ -912,6 +910,8 @@ def main():
os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
)
clip_image_encoder = clip_image_encoder.eval()
else:
clip_image_encoder = None
# Get Transformer
generator_transformer3d = TurboWanTransformer3DModel.from_pretrained(
@@ -987,7 +987,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1007,26 +1007,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1536,13 +1516,14 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler= accelerator_fake_score_transformer3d.prepare(
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
real_score_transformer3d = shard_fn(real_score_transformer3d)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
@@ -2268,20 +2249,19 @@ def main():
accelerator_fake_score_transformer3d.save_state(fake_score_save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2289,20 +2269,19 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/turbodiffusion/train_distill.sh
+1 -1
View File
@@ -28,7 +28,7 @@ accelerate launch --mixed_precision="bf16" scripts/turbodiffusion/train_distill.
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_distill_turbodiffusion" \
--output_dir="output_dir_turbodiffusion_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.1/train.py
+156 -172
View File
@@ -48,7 +48,8 @@ from omegaconf import OmegaConf
from packaging import version
from PIL import Image
from torch.distributed.fsdp.fully_sharded_data_parallel import (
FullOptimStateDictConfig, FullStateDictConfig, ShardedStateDictConfig, ShardedOptimStateDictConfig)
FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig,
ShardedStateDictConfig)
from torch.utils.data import RandomSampler
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
@@ -64,18 +65,20 @@ for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
from videox_fun.data.dataset_image_video import (ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
ImageVideoSampler,
get_random_mask)
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel)
from videox_fun.pipeline import WanPipeline, WanI2VPipeline
WanTransformer3DModel)
from videox_fun.pipeline import WanI2VPipeline, WanPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -163,116 +166,109 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
else:
pipeline = WanPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -344,6 +340,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -902,6 +905,8 @@ def main():
os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
)
clip_image_encoder = clip_image_encoder.eval()
else:
clip_image_encoder = None
# Get Transformer
transformer3d = WanTransformer3DModel.from_pretrained(
@@ -971,7 +976,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -991,26 +996,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1066,8 +1051,8 @@ def main():
if args.gradient_checkpointing:
transformer3d.enable_gradient_checkpointing()
elif args.selective_ac > 0:
from videox_fun.utils.ac_handle import apply_checkpointing, partial
from videox_fun.models.wan_transformer3d import WanAttentionBlock
from videox_fun.utils.ac_handle import apply_checkpointing, partial
apply_selective_ac = partial(apply_checkpointing, block=WanAttentionBlock)
apply_selective_ac(transformer3d, p=args.selective_ac)
@@ -1398,8 +1383,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -1866,27 +1852,26 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1894,27 +1879,26 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.1/train.sh
+2 -2
View File
@@ -26,7 +26,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.1" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -69,7 +69,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train.py \
# --lr_scheduler="constant_with_warmup" \
# --lr_warmup_steps=100 \
# --seed=42 \
# --output_dir="output_dir" \
# --output_dir="output_dir_wan2.1" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
scripts/wan2.1/train_distill.py
+131 -153
View File
@@ -80,7 +80,9 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel)
from videox_fun.pipeline import WanI2VPipeline, WanPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -168,120 +170,109 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
else:
pipeline = WanPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video = input_video,
mask_video = input_video_mask,
num_inference_steps = len(args.denoising_step_list),
guidance_scale = 1.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = len(args.denoising_step_indices_list),
guidance_scale = 1.0,
).videos
video = input_video,
mask_video = input_video_mask,
num_inference_steps = len(args.denoising_step_list),
guidance_scale = 1.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = len(args.denoising_step_list),
guidance_scale = 1.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = len(args.denoising_step_indices_list),
guidance_scale = 1.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -353,6 +344,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--negative_prompt",
type=str,
@@ -911,6 +909,8 @@ def main():
os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
)
clip_image_encoder = clip_image_encoder.eval()
else:
clip_image_encoder = None
# Get Transformer
generator_transformer3d = WanTransformer3DModel.from_pretrained(
@@ -986,7 +986,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1006,26 +1006,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1535,13 +1515,13 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler= accelerator_fake_score_transformer3d.prepare(
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
real_score_transformer3d = shard_fn(real_score_transformer3d)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
@@ -2267,20 +2247,19 @@ def main():
accelerator_fake_score_transformer3d.save_state(fake_score_save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2288,20 +2267,19 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.1/train_distill.sh
+1 -1
View File
@@ -27,7 +27,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_distill.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_distill" \
--output_dir="output_dir_wan2.1_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.1/train_distill_lora.py
+140 -182
View File
@@ -83,7 +83,9 @@ from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -169,119 +171,113 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, network, config, args, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
for i in range(len(args.validation_prompts)):
with torch.no_grad():
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = len(args.denoising_step_indices_list),
guidance_scale = 1.0,
).videos
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = len(args.denoising_step_indices_list),
guidance_scale = 1.0,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -353,6 +349,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--negative_prompt",
type=str,
@@ -933,6 +936,8 @@ def main():
os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
)
clip_image_encoder = clip_image_encoder.eval()
else:
clip_image_encoder = None
# Get Transformer
generator_transformer3d = WanTransformer3DModel.from_pretrained(
@@ -1028,7 +1033,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1057,15 +1062,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -1078,34 +1075,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1562,7 +1537,7 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler = accelerator_fake_score_transformer3d.prepare(
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
elif fsdp_stage != 0:
else:
generator_transformer3d.network = network
generator_transformer3d = generator_transformer3d.to(dtype=weight_dtype)
generator_transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
@@ -1573,21 +1548,6 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler = accelerator_fake_score_transformer3d.prepare(
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
fake_score_network, critic_optimizer, fake_score_lr_scheduler= accelerator_fake_score_transformer3d.prepare(
fake_score_network, critic_optimizer, fake_score_lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
generator_transformer3d = shard_fn(generator_transformer3d)
fake_score_transformer3d = shard_fn(fake_score_transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -2362,21 +2322,20 @@ def main():
accelerator_fake_score_transformer3d.save_state(fake_score_save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2384,21 +2343,20 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
generator_transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.1/train_distill_lora.sh
+1 -1
View File
@@ -25,7 +25,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_distill_lora.py
--learning_rate=1e-05 \
--learning_rate_critic=1e-06 \
--seed=42 \
--output_dir="output_dir_distill" \
--output_dir="output_dir_wan2.1_distill_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.1/train_lora.py
+140 -170
View File
@@ -75,7 +75,9 @@ from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -161,119 +163,113 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, network, config, args, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.train_mode != "normal":
pipeline = WanI2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
for i in range(len(args.validation_prompts)):
with torch.no_grad():
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -345,6 +341,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -906,6 +909,8 @@ def main():
os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
)
clip_image_encoder = clip_image_encoder.eval()
else:
clip_image_encoder = None
# Get Transformer
transformer3d = WanTransformer3DModel.from_pretrained(
@@ -971,7 +976,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1000,7 +1005,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -1013,34 +1018,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1370,23 +1353,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -1910,21 +1882,20 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1932,21 +1903,20 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.1/train_lora.sh
+2 -2
View File
@@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.1_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -67,7 +67,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1/train_lora.py \
# --checkpointing_steps=50 \
# --learning_rate=1e-04 \
# --seed=42 \
# --output_dir="output_dir" \
# --output_dir="output_dir_wan2.1_lora" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
scripts/wan2.1_fun/train.py
+145 -162
View File
@@ -77,7 +77,9 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel)
from videox_fun.pipeline import WanFunInpaintPipeline, WanFunPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -146,116 +148,109 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = WanFunInpaintPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
else:
pipeline = WanFunPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.train_mode != "normal":
pipeline = WanFunInpaintPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanFunPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -314,6 +309,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -866,6 +868,8 @@ def main():
os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
)
clip_image_encoder = clip_image_encoder.eval()
else:
clip_image_encoder = None
# Get Transformer
transformer3d = WanTransformer3DModel.from_pretrained(
@@ -935,7 +939,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -955,26 +959,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1395,8 +1379,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -1869,27 +1854,26 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1897,27 +1881,26 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.1_fun/train.sh
+2 -2
View File
@@ -26,7 +26,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_fun/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.1_fun" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -69,7 +69,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_fun/train.py \
# --lr_scheduler="constant_with_warmup" \
# --lr_warmup_steps=100 \
# --seed=42 \
# --output_dir="output_dir" \
# --output_dir="output_dir_wan2.1_fun" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
scripts/wan2.1_fun/train_control.py
+109 -113
View File
@@ -77,7 +77,8 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel)
from videox_fun.pipeline import WanFunControlPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -112,64 +113,80 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
transformer3d_val = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = WanFunControlPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
pipeline = pipeline.to(accelerator.device)
pipeline = WanFunControlPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
for i in range(len(args.validation_prompts)):
import cv2
cap = cv2.VideoCapture(args.validation_paths[i])
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
cap.release()
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
control_video = input_video,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
control_video = input_video,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -868,7 +885,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -888,26 +905,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1399,8 +1396,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -1904,27 +1902,26 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1932,27 +1929,26 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.1_fun/train_control.sh
+1 -1
View File
@@ -26,7 +26,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_fun/train_control.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.1_fun_control" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.1_fun/train_control_lora.py
+104 -132
View File
@@ -80,7 +80,8 @@ from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -113,70 +114,84 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, network, config, args, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
transformer3d_val = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = WanFunControlPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
pipeline = pipeline.to(accelerator.device)
pipeline = WanFunControlPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
for i in range(len(args.validation_prompts)):
import cv2
cap = cv2.VideoCapture(args.validation_paths[i])
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
cap.release()
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
control_video = input_video,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
control_video = input_video,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -874,7 +889,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -903,15 +918,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -924,34 +931,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1390,23 +1375,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -1967,21 +1941,20 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1989,21 +1962,20 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.1_fun/train_control_lora.sh
+1 -1
View File
@@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_fun/train_control_lora
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.1_fun_control_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.1_fun/train_lora.py
+140 -178
View File
@@ -76,7 +76,9 @@ from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -142,119 +144,113 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, network, config, args, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = WanFunInpaintPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanFunPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.train_mode != "normal":
pipeline = WanFunInpaintPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
else:
pipeline = WanFunPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
for i in range(len(args.validation_prompts)):
with torch.no_grad():
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -313,6 +309,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -868,6 +871,8 @@ def main():
os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
)
clip_image_encoder = clip_image_encoder.eval()
else:
clip_image_encoder = None
# Get Transformer
transformer3d = WanTransformer3DModel.from_pretrained(
@@ -933,7 +938,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -962,15 +967,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -983,34 +980,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1378,23 +1353,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -1924,21 +1888,20 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1946,21 +1909,20 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.1_fun/train_lora.sh
+2 -2
View File
@@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_fun/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.1_fun_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -68,7 +68,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_fun/train_lora.py \
# --checkpointing_steps=50 \
# --learning_rate=1e-04 \
# --seed=42 \
# --output_dir="output_dir" \
# --output_dir="output_dir_wan2.1_fun_lora" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
scripts/wan2.1_vace/train.py
+112 -119
View File
@@ -78,7 +78,8 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel,
VaceWanTransformer3DModel, WanT5EncoderModel)
from videox_fun.pipeline import WanVacePipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -111,71 +112,86 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
transformer3d_val = VaceWanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = WanVacePipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = WanVacePipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
for i in range(len(args.validation_prompts)):
import cv2
cap = cv2.VideoCapture(args.validation_paths[i])
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
cap.release()
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
inpaint_video, inpaint_video_mask, clip_image = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
control_video, _, _, _ = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
inpaint_video, inpaint_video_mask, clip_image = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[height, width])
control_video, _, _, _ = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video = inpaint_video,
mask_video = inpaint_video_mask,
control_video = control_video,
subject_ref_images = None,
vace_context_scale = 1,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video = inpaint_video,
mask_video = inpaint_video_mask,
control_video = control_video,
subject_ref_images = None,
vace_context_scale = 1,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -854,7 +870,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -874,26 +890,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1395,8 +1391,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -1999,27 +1996,25 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2027,27 +2022,25 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.1_vace/train.sh
+1 -1
View File
@@ -25,7 +25,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_vace/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.1_vace" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.2/train.py
+179 -210
View File
@@ -48,7 +48,8 @@ from omegaconf import OmegaConf
from packaging import version
from PIL import Image
from torch.distributed.fsdp.fully_sharded_data_parallel import (
FullOptimStateDictConfig, FullStateDictConfig, ShardedStateDictConfig, ShardedOptimStateDictConfig)
FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig,
ShardedStateDictConfig)
from torch.utils.data import RandomSampler
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
@@ -64,18 +65,20 @@ for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
from videox_fun.data.dataset_image_video import (ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, WanT5EncoderModel,
Wan2_2Transformer3DModel)
from videox_fun.pipeline import Wan2_2Pipeline, Wan2_2I2VPipeline
ImageVideoSampler,
get_random_mask)
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
Wan2_2Transformer3DModel, WanT5EncoderModel)
from videox_fun.pipeline import Wan2_2I2VPipeline, Wan2_2Pipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -163,152 +166,132 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
return images
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -380,6 +363,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -1009,7 +999,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1029,26 +1019,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1432,8 +1402,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -1921,26 +1892,25 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1948,26 +1918,25 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2/train.sh
+2 -2
View File
@@ -26,7 +26,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -73,7 +73,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train.py \
# --lr_scheduler="constant_with_warmup" \
# --lr_warmup_steps=100 \
# --seed=42 \
# --output_dir="output_dir" \
# --output_dir="output_dir_wan2.2" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
scripts/wan2.2/train_animate.py
+166 -163
View File
@@ -62,22 +62,22 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.data import (ImageVideoDataset, ImageVideoSampler,
VideoAnimateDataset, get_random_mask,
process_pose_file, process_pose_params)
from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
from videox_fun.data import (VideoAnimateDataset,
ImageVideoDataset,
ImageVideoSampler,
get_random_mask,
process_pose_file,
process_pose_params)
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, CLIPModel, WanT5EncoderModel,
Wan2_2Transformer3DModel_Animate)
from videox_fun.pipeline import Wan2_2FunControlPipeline
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
CLIPModel, Wan2_2Transformer3DModel_Animate,
WanT5EncoderModel)
from videox_fun.pipeline import Wan2_2AnimatePipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image,
get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -157,100 +157,132 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
pipeline = Wan2_2AnimatePipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
pipeline = pipeline.to(accelerator.device)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = Wan2_2FunControlPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
images = []
for i in range(len(args.validation_prompts)):
src_root_path = args.validation_paths[i]
src_pose_path = os.path.join(src_root_path, "src_pose.mp4")
src_face_path = os.path.join(src_root_path, "src_face.mp4")
src_ref_path = os.path.join(src_root_path, "src_ref.png")
src_bg_path = os.path.join(src_root_path, "src_bg.mp4")
src_mask_path = os.path.join(src_root_path, "src_mask.mp4")
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
import cv2
cap = cv2.VideoCapture(src_pose_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
cap.release()
control_video = input_video,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
width, height = calculate_dimensions(args.video_sample_size * args.video_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
fps = 16
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
pose_video, _, _, _ = get_video_to_video_latent(src_pose_path, video_length=video_length, sample_size=[height, width], fps=fps, ref_image=None)
face_video, _, _, _ = get_video_to_video_latent(src_face_path, video_length=video_length, sample_size=[512, 512], fps=fps, ref_image=None)
ref_image = get_image(src_ref_path)
return images
if os.path.exists(src_bg_path):
bg_video, _, _, _ = get_video_to_video_latent(src_bg_path, video_length=video_length, sample_size=[height, width], fps=fps, ref_image=None)
mask_video, _, _, _ = get_video_to_video_latent(src_mask_path, video_length=video_length, sample_size=[height, width], fps=fps, ref_image=None)
mask_video = mask_video[:, :1]
replace_flag = True
else:
bg_video = None
mask_video = None
replace_flag = False
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
guidance_scale = 4.5,
num_inference_steps = 25,
pose_video = pose_video,
face_video = face_video,
ref_image = ref_image,
bg_video = bg_video,
mask_video = mask_video,
replace_flag = replace_flag,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -450,9 +482,6 @@ def parse_args():
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
),
)
parser.add_argument(
"--report_model_info", action="store_true", help="Whether or not to report more info about model (such as norm, grad)."
)
parser.add_argument(
"--mixed_precision",
type=str,
@@ -577,12 +606,6 @@ def parse_args():
default=512,
help="Sample size of the video.",
)
parser.add_argument(
"--image_sample_size",
type=int,
default=512,
help="Sample size of the image.",
)
parser.add_argument(
"--fix_sample_size",
nargs=2, type=int, default=None,
@@ -940,7 +963,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -960,26 +983,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1118,7 +1121,6 @@ def main():
if args.fix_sample_size is not None and args.enable_bucket:
args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size)
args.image_sample_size = max(max(args.fix_sample_size), args.image_sample_size)
args.training_with_video_token_length = False
args.random_hw_adapt = False
@@ -1423,8 +1425,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -1951,26 +1954,26 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1978,26 +1981,26 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2/train_animate.sh
+1 -1
View File
@@ -25,7 +25,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_animate.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2_animate" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.2/train_animate_lora.py
+154 -170
View File
@@ -77,13 +77,14 @@ from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
CLIPModel, Wan2_2Transformer3DModel,
Wan2_2Transformer3DModel_Animate,
WanT5EncoderModel)
from videox_fun.pipeline import (Wan2_2FunControlPipeline, Wan2_2I2VPipeline,
Wan2_2Pipeline)
from videox_fun.pipeline import Wan2_2AnimatePipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image,
get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -163,104 +164,134 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
pipeline = Wan2_2AnimatePipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
pipeline = pipeline.to(accelerator.device)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = Wan2_2FunControlPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
images = []
for i in range(len(args.validation_prompts)):
src_root_path = args.validation_paths[i]
src_pose_path = os.path.join(src_root_path, "src_pose.mp4")
src_face_path = os.path.join(src_root_path, "src_face.mp4")
src_ref_path = os.path.join(src_root_path, "src_ref.png")
src_bg_path = os.path.join(src_root_path, "src_bg.mp4")
src_mask_path = os.path.join(src_root_path, "src_mask.mp4")
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
import cv2
cap = cv2.VideoCapture(src_pose_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
cap.release()
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
width, height = calculate_dimensions(args.video_sample_size * args.video_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
fps = 16
control_video = input_video,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
pose_video, _, _, _ = get_video_to_video_latent(src_pose_path, video_length=video_length, sample_size=[height, width], fps=fps, ref_image=None)
face_video, _, _, _ = get_video_to_video_latent(src_face_path, video_length=video_length, sample_size=[512, 512], fps=fps, ref_image=None)
ref_image = get_image(src_ref_path)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if os.path.exists(src_bg_path):
bg_video, _, _, _ = get_video_to_video_latent(src_bg_path, video_length=video_length, sample_size=[height, width], fps=fps, ref_image=None)
mask_video, _, _, _ = get_video_to_video_latent(src_mask_path, video_length=video_length, sample_size=[height, width], fps=fps, ref_image=None)
mask_video = mask_video[:, :1]
replace_flag = True
else:
bg_video = None
mask_video = None
replace_flag = False
return images
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
guidance_scale = 4.5,
num_inference_steps = 25,
pose_video = pose_video,
face_video = face_video,
ref_image = ref_image,
bg_video = bg_video,
mask_video = mask_video,
replace_flag = replace_flag,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -604,12 +635,6 @@ def parse_args():
default=512,
help="Sample size of the video.",
)
parser.add_argument(
"--image_sample_size",
type=int,
default=512,
help="Sample size of the image.",
)
parser.add_argument(
"--fix_sample_size",
nargs=2, type=int, default=None,
@@ -949,7 +974,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -978,15 +1003,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -999,34 +1016,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1112,7 +1107,6 @@ def main():
if args.fix_sample_size is not None and args.enable_bucket:
args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size)
args.image_sample_size = max(max(args.fix_sample_size), args.image_sample_size)
args.training_with_video_token_length = False
args.random_hw_adapt = False
@@ -1417,23 +1411,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -1962,40 +1945,41 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2/train_animate_lora.sh
+1 -1
View File
@@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_animate_lora.py
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_animate_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.2/train_distill.py
+159 -190
View File
@@ -76,11 +76,14 @@ from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
from videox_fun.data.dataset_image_video import (ImageVideoDataset,
ImageVideoSampler,
TextDataset, get_random_mask)
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel, AutoencoderKLWan3_8,
Wan2_2Transformer3DModel)
from videox_fun.pipeline import WanI2VPipeline, WanPipeline
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
CLIPModel, Wan2_2Transformer3DModel,
WanT5EncoderModel)
from videox_fun.pipeline import Wan2_2I2VPipeline, Wan2_2Pipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -168,152 +171,133 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
return images
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 4,
guidance_scale = 1.0,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 4,
guidance_scale = 1.0,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -385,6 +369,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--negative_prompt",
type=str,
@@ -1026,7 +1017,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1046,26 +1037,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1575,13 +1546,13 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler= accelerator_fake_score_transformer3d.prepare(
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
real_score_transformer3d = shard_fn(real_score_transformer3d)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
@@ -2367,19 +2338,18 @@ def main():
accelerator_fake_score_transformer3d.save_state(fake_score_save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2387,19 +2357,18 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2/train_distill.sh
+1 -1
View File
@@ -27,7 +27,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_distill.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_distill" \
--output_dir="output_dir_wan2.2_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.2/train_distill_lora.py
+169 -224
View File
@@ -79,12 +79,14 @@ from videox_fun.data.dataset_image_video import (ImageVideoDataset,
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
CLIPModel, Wan2_2Transformer3DModel,
WanT5EncoderModel)
from videox_fun.pipeline import WanI2VPipeline, WanPipeline
from videox_fun.pipeline import Wan2_2I2VPipeline, Wan2_2Pipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -170,158 +172,136 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, config, args, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
return images
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 4,
guidance_scale = 1.0,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 4,
guidance_scale = 1.0,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -393,6 +373,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--negative_prompt",
type=str,
@@ -1073,7 +1060,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1102,15 +1089,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -1123,34 +1102,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1607,7 +1564,7 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler = accelerator_fake_score_transformer3d.prepare(
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
elif fsdp_stage != 0:
else:
generator_transformer3d.network = network
generator_transformer3d = generator_transformer3d.to(dtype=weight_dtype)
generator_transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
@@ -1618,21 +1575,6 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler = accelerator_fake_score_transformer3d.prepare(
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
fake_score_network, critic_optimizer, fake_score_lr_scheduler= accelerator_fake_score_transformer3d.prepare(
fake_score_network, critic_optimizer, fake_score_lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
generator_transformer3d = shard_fn(generator_transformer3d)
fake_score_transformer3d = shard_fn(fake_score_transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -1640,6 +1582,11 @@ def main():
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
real_score_transformer3d = shard_fn(real_score_transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
# Move text_encode and vae to gpu and cast to weight_dtype
@@ -2467,20 +2414,19 @@ def main():
accelerator_fake_score_transformer3d.save_state(fake_score_save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
generator_transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
generator_transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2488,20 +2434,19 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
generator_transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
generator_transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2/train_distill_lora.sh
+1 -1
View File
@@ -25,7 +25,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_distill_lora.py
--learning_rate=1e-05 \
--learning_rate_critic=1e-06 \
--seed=42 \
--output_dir="output_dir_distill" \
--output_dir="output_dir_wan2.2_distill_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.2/train_lora.py
+164 -217
View File
@@ -75,7 +75,9 @@ from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -161,155 +163,136 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, config, args, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -381,6 +364,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -1015,7 +1005,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1044,15 +1034,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -1065,34 +1047,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1424,23 +1384,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -1986,20 +1935,19 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2007,20 +1955,19 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2/train_lora.sh
+2 -2
View File
@@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -72,7 +72,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_lora.py \
# --checkpointing_steps=50 \
# --learning_rate=1e-04 \
# --seed=42 \
# --output_dir="output_dir" \
# --output_dir="output_dir_wan2.2_lora" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
scripts/wan2.2/train_s2v.py
+167 -203
View File
@@ -21,6 +21,7 @@ import logging
import math
import os
import pickle
import random
import shutil
import sys
@@ -32,7 +33,6 @@ import torch.nn.functional as F
import torch.utils.checkpoint
import torchvision.transforms.functional as TF
import transformers
import random
from accelerate import Accelerator, FullyShardedDataParallelPlugin
from accelerate.logging import get_logger
from accelerate.state import AcceleratorState
@@ -49,7 +49,8 @@ from omegaconf import OmegaConf
from packaging import version
from PIL import Image
from torch.distributed.fsdp.fully_sharded_data_parallel import (
FullOptimStateDictConfig, FullStateDictConfig, ShardedStateDictConfig, ShardedOptimStateDictConfig)
FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig,
ShardedStateDictConfig)
from torch.utils.data import RandomSampler
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
@@ -65,19 +66,23 @@ for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
ASPECT_RATIO_RANDOM_CROP_512,
ASPECT_RATIO_RANDOM_CROP_PROB,
AspectRatioBatchImageVideoSampler,
RandomSampler, get_closest_ratio)
from videox_fun.data.dataset_image_video import (ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
ImageVideoSampler,
get_random_mask)
from videox_fun.data.dataset_video import VideoSpeechControlDataset
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, WanT5EncoderModel, WanAudioEncoder,
Wan2_2Transformer3DModel_S2V)
from videox_fun.pipeline import Wan2_2S2VPipeline, Wan2_2I2VPipeline
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
Wan2_2Transformer3DModel_S2V, WanAudioEncoder,
WanT5EncoderModel)
from videox_fun.pipeline import Wan2_2I2VPipeline, Wan2_2S2VPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -131,144 +136,108 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, audio_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = Wan2_2S2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
pipeline = Wan2_2S2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
audio_encoder=audio_encoder,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
return images
for i in range(len(args.validation_prompts)):
start_image = Image.open(args.validation_image_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.video_sample_size * args.video_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
pose_video, _, _, _ = get_video_to_video_latent(None, video_length=video_length, sample_size=(height, width), ref_image=None)
ref_image = get_image_latent(args.validation_image_paths[i], sample_size=(height, width))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
audio_path = args.validation_audio_paths[i],
pose_video = pose_video,
ref_image = ref_image,
init_first_frame = False,
num_inference_steps = 25,
guidance_scale = 4.5,
fps = 16,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -340,6 +309,20 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_image_paths",
type=str,
default=None,
nargs="+",
help=("A set of images evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_audio_paths",
type=str,
default=None,
nargs="+",
help=("A set of audios evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -680,7 +663,7 @@ def parse_args():
parser.add_argument(
"--boundary_type",
type=str,
default="low",
default="full",
help=(
'The format of training data. Support `"low"` and `"high"`'
),
@@ -976,7 +959,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -996,26 +979,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1453,8 +1416,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -2006,26 +1970,26 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
audio_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2033,26 +1997,26 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
audio_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2/train_s2v.sh
+3 -1
View File
@@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_s2v.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2_s2v" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
@@ -35,5 +35,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_s2v.py \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--boundary_type="full" \
--control_ref_image="random" \
--low_vram \
--trainable_modules "."
scripts/wan2.2/train_s2v_lora.py
+146 -204
View File
@@ -89,7 +89,8 @@ from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -145,144 +146,110 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, audio_encoder, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel_S2V.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = Wan2_2S2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
pipeline = Wan2_2S2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
audio_encoder=audio_encoder,
)
pipeline = pipeline.to(accelerator.device)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
return images
for i in range(len(args.validation_prompts)):
start_image = Image.open(args.validation_image_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.video_sample_size * args.video_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
pose_video, _, _, _ = get_video_to_video_latent(None, video_length=video_length, sample_size=(height, width), ref_image=None)
ref_image = get_image_latent(args.validation_image_paths[i], sample_size=(height, width))
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
audio_path = args.validation_audio_paths[i],
pose_video = pose_video,
ref_image = ref_image,
init_first_frame = False,
num_inference_steps = 25,
guidance_scale = 4.5,
fps = 16,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
@@ -354,6 +321,20 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_image_paths",
type=str,
default=None,
nargs="+",
help=("A set of images evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_audio_paths",
type=str,
default=None,
nargs="+",
help=("A set of audios evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -711,7 +692,7 @@ def parse_args():
parser.add_argument(
"--boundary_type",
type=str,
default="low",
default="full",
help=(
'The format of training data. Support `"low"` and `"high"`'
),
@@ -999,7 +980,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1028,15 +1009,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -1049,34 +1022,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1461,23 +1412,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -2034,19 +1974,20 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
audio_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2054,19 +1995,20 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
audio_encoder,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2/train_s2v_lora.sh
+1 -1
View File
@@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_s2v_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2_s2v_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.2_fun/train.py
+171 -212
View File
@@ -73,11 +73,14 @@ from videox_fun.data.bucket_sampler import (ASPECT_RATIO_512,
from videox_fun.data.dataset_image_video import (ImageVideoDataset,
ImageVideoSampler,
get_random_mask)
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, CLIPModel, WanT5EncoderModel,
Wan2_2Transformer3DModel)
from videox_fun.pipeline import Wan2_2Pipeline, Wan2_2I2VPipeline
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
CLIPModel, Wan2_2Transformer3DModel,
WanT5EncoderModel)
from videox_fun.pipeline import Wan2_2I2VPipeline, Wan2_2Pipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -146,161 +149,131 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
return images
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -359,6 +332,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -988,7 +968,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1008,26 +988,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1449,8 +1409,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -1926,26 +1887,25 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1953,26 +1913,25 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2_fun/train.sh
+1 -1
View File
@@ -26,7 +26,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2_fun/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2_fun" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.2_fun/train_control.py
+135 -145
View File
@@ -73,11 +73,13 @@ from videox_fun.data.dataset_image_video import (ImageVideoControlDataset,
get_random_mask,
process_pose_file,
process_pose_params)
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, CLIPModel, WanT5EncoderModel,
Wan2_2Transformer3DModel)
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
CLIPModel, Wan2_2Transformer3DModel,
WanT5EncoderModel)
from videox_fun.pipeline import Wan2_2FunControlPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -148,98 +150,107 @@ logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = Wan2_2FunControlPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = Wan2_2FunControlPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
for i in range(len(args.validation_prompts)):
import cv2
cap = cv2.VideoCapture(args.validation_paths[i])
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
cap.release()
control_video = input_video,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
inpaint_video, inpaint_video_mask, clip_image = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[height, width])
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
control_video = input_video,
video = inpaint_video,
mask_video = inpaint_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -955,7 +966,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -975,26 +986,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1504,8 +1495,9 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
@@ -2071,26 +2063,25 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2098,26 +2089,25 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2_fun/train_control.sh
+1 -1
View File
@@ -23,7 +23,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2_fun/train_control.py \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-05 \
--lr_scheduler="constant_with_warmup" \
--output_dir="output_dir_wan2.2_fun_control" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
scripts/wan2.2_fun/train_control_lora.py
+127 -163
View File
@@ -81,7 +81,8 @@ from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -150,105 +151,111 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, config, args, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
pipeline = Wan2_2FunControlPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = Wan2_2FunControlPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
for i in range(len(args.validation_prompts)):
import cv2
cap = cv2.VideoCapture(args.validation_paths[i])
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
cap.release()
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
inpaint_video, inpaint_video_mask, clip_image = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[height, width])
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, ref_image, clip_image = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
control_video = input_video,
video = inpaint_video,
mask_video = inpaint_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
control_video = input_video,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -963,7 +970,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -992,15 +999,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -1013,34 +1012,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1497,23 +1474,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -2136,20 +2102,19 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2157,20 +2122,19 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2_fun/train_control_lora.sh
+1 -1
View File
@@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2_fun/train_control_lora
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2_fun_control_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.2_fun/train_lora.py
+163 -217
View File
@@ -76,7 +76,9 @@ from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
if is_wandb_available():
import wandb
@@ -143,155 +145,135 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, config, args, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
pipeline = merge_lora(
pipeline, None, 1, accelerator.device, state_dict=accelerator.unwrap_model(network).state_dict(), transformer_only=True
)
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
for i in range(len(args.validation_prompts)):
with torch.no_grad():
if args.train_mode != "normal":
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
video_length = 1
input_video, input_video_mask, _ = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
guidance_scale = 6.0,
generator = generator,
video = input_video,
mask_video = input_video_mask,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
with torch.autocast("cuda", dtype=weight_dtype):
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
sample = pipeline(
args.validation_prompts[i],
num_frames = 1,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-image-{i}.gif"))
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
for i in range(len(args.validation_prompts)):
if args.train_mode != "normal":
start_image = Image.open(args.validation_paths[i])
width, height = start_image.width, start_image.height
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = 25,
guidance_scale = 4.5,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.900)
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
else:
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
num_inference_steps = 25,
guidance_scale = 4.5,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
transformer3d.to(accelerator.device, dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -350,6 +332,13 @@ def parse_args():
nargs="+",
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--validation_paths",
type=str,
default=None,
nargs="+",
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
)
parser.add_argument(
"--output_dir",
type=str,
@@ -978,7 +967,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1007,15 +996,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -1028,34 +1009,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1424,23 +1383,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
@@ -1974,20 +1922,19 @@ def main():
accelerator.save_state(accelerator_save_path)
logger.info(f"Saved state to {accelerator_save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -1995,20 +1942,19 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
config,
args,
accelerator,
weight_dtype,
global_step,
)
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
network,
args,
config,
accelerator,
weight_dtype,
global_step,
)
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2_fun/train_lora.sh
+1 -1
View File
@@ -24,7 +24,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2_fun/train_lora.py \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2_fun_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/wan2.2_vace_fun/train.py
+135 -159
View File
@@ -74,11 +74,13 @@ from videox_fun.data.dataset_image_video import (ImageVideoControlDataset,
padding_image,
process_pose_file,
process_pose_params)
from videox_fun.models import (AutoencoderKLWan, CLIPModel, AutoencoderKLWan3_8,
VaceWanTransformer3DModel, WanT5EncoderModel)
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
CLIPModel, VaceWanTransformer3DModel,
WanT5EncoderModel)
from videox_fun.pipeline import Wan2_2VaceFunPipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import (get_image_to_video_latent,
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
@@ -111,107 +113,108 @@ check_min_version("0.18.0.dev0")
logger = get_logger(__name__, log_level="INFO")
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
def log_validation(vae, text_encoder, tokenizer, transformer3d, args, config, accelerator, weight_dtype, global_step):
try:
logger.info("Running validation... ")
if args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
transformer3d_2_val = None
else:
if args.boundary_type == "low":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
if is_deepspeed:
origin_config = transformer3d.config
transformer3d.config = accelerator.unwrap_model(transformer3d).config
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = VaceWanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = VaceWanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
pipeline = Wan2_2VaceFunPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2_val = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
transformer3d_2_val.load_state_dict(accelerator.unwrap_model(transformer3d).state_dict())
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
if args.seed is None:
generator = None
else:
rank_seed = args.seed + accelerator.process_index
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
pipeline = Wan2_2VaceFunPipeline(
vae=accelerator.unwrap_model(vae).to(weight_dtype),
text_encoder=accelerator.unwrap_model(text_encoder),
tokenizer=tokenizer,
transformer=transformer3d_val,
transformer_2=transformer3d_2_val,
scheduler=scheduler,
clip_image_encoder=clip_image_encoder,
)
pipeline = pipeline.to(accelerator.device)
for i in range(len(args.validation_prompts)):
import cv2
cap = cv2.VideoCapture(args.validation_paths[i])
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
cap.release()
if args.seed is None:
generator = None
else:
generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
inpaint_video, inpaint_video_mask, clip_image = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[height, width])
control_video, _, _, _ = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = height,
width = width,
generator = generator,
images = []
for i in range(len(args.validation_prompts)):
with torch.no_grad():
with torch.autocast("cuda", dtype=weight_dtype):
video_length = int(args.video_sample_n_frames // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
inpaint_video, inpaint_video_mask, clip_image = get_image_to_video_latent(None, None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
control_video, _, _, _ = get_video_to_video_latent(args.validation_paths[i], video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "bad detailed",
height = args.video_sample_size,
width = args.video_sample_size,
generator = generator,
video = inpaint_video,
mask_video = inpaint_video_mask,
control_video = control_video,
subject_ref_images = None,
vace_context_scale = 1,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(
sample,
os.path.join(
args.output_dir,
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.gif"
)
)
video = inpaint_video,
mask_video = inpaint_video_mask,
control_video = control_video,
subject_ref_images = None,
vace_context_scale = 1,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
save_videos_grid(sample, os.path.join(args.output_dir, f"sample/sample-{global_step}-{i}.gif"))
del pipeline
del transformer3d_val
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return images
del pipeline
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if is_deepspeed:
transformer3d.config = origin_config
except Exception as e:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
print(f"Eval error with info {e}")
return None
print(f"Eval error on rank {accelerator.process_index} with info {e}")
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if not args.enable_text_encoder_in_dataloader:
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -906,7 +909,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -926,26 +929,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1448,16 +1431,13 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
text_encoder = shard_fn(text_encoder)
# shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
# transformer3d = shard_fn(transformer3d)
if args.use_ema:
ema_transformer3d.to(accelerator.device)
@@ -2070,27 +2050,25 @@ def main():
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
if accelerator.is_main_process:
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -2098,27 +2076,25 @@ def main():
if global_step >= args.max_train_steps:
break
if accelerator.is_main_process:
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
clip_image_encoder,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer3d.store(transformer3d.parameters())
ema_transformer3d.copy_to(transformer3d.parameters())
log_validation(
vae,
text_encoder,
tokenizer,
transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original transformer3d parameters.
ema_transformer3d.restore(transformer3d.parameters())
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
scripts/wan2.2_vace_fun/train.sh
+1 -1
View File
@@ -26,7 +26,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.2_vace_fun/train.py \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2_vace_fun" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
scripts/z_image/train.py
+2 -22
View File
@@ -203,7 +203,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -869,7 +869,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -889,26 +889,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
scripts/z_image/train_distill.py
+332 -257
View File
@@ -25,6 +25,7 @@ import pickle
import random
import shutil
import sys
from functools import partial
from typing import (Any, Callable, Dict, List, NamedTuple, Optional, Tuple,
Union)
@@ -55,7 +56,7 @@ from PIL import Image
from torch.distributed.fsdp.fully_sharded_data_parallel import (
FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig,
ShardedStateDictConfig)
from torch.utils.data import Dataset, RandomSampler, BatchSampler
from torch.utils.data import BatchSampler, Dataset, RandomSampler
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from tqdm.auto import tqdm
@@ -84,10 +85,12 @@ from videox_fun.models import (AutoencoderKL, AutoProcessor, AutoTokenizer,
Qwen3ForCausalLM, QwenImageTransformer2DModel,
ZImageTransformer2DModel)
from videox_fun.pipeline import ZImagePipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
from videox_fun.utils import (DiscreteSampling, RectifiedFlow_TrigFlowWrapper,
calculate_dimensions,
convert_peft_lora_to_kohya_lora, create_network,
get_image_latent, get_image_to_video_latent,
merge_lora, sample_trigflow_timesteps,
save_videos_grid, unmerge_lora)
if is_wandb_available():
import wandb
@@ -208,7 +211,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -603,12 +606,6 @@ def parse_args():
default=[],
help='Enter a list of trainable modules with lower learning rate'
)
parser.add_argument(
'--tokenizer_max_length',
type=int,
default=512,
help='Max length of tokenizer'
)
parser.add_argument(
"--use_deepspeed", action="store_true", help="Whether or not to use deepspeed."
)
@@ -618,47 +615,6 @@ def parse_args():
parser.add_argument(
"--low_vram", action="store_true", help="Whether enable low_vram mode."
)
parser.add_argument(
"--prompt_template_encode",
type=str,
default="<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n",
help=(
'The prompt template for text encoder.'
),
)
parser.add_argument(
"--prompt_template_encode_start_idx",
type=int,
default=34,
help=(
'The start idx for prompt template.'
),
)
parser.add_argument(
"--train_mode",
type=str,
default="normal",
help=(
'The format of training data. Support `"normal"`'
' (default), `"i2v"`.'
),
)
parser.add_argument(
"--abnormal_norm_clip_start",
type=int,
default=1000,
help=(
'When do we start doing additional processing on abnormal gradients. '
),
)
parser.add_argument(
"--initial_grad_norm_ratio",
type=int,
default=5,
help=(
'The initial gradient is relative to the multiple of the max_grad_norm. '
),
)
parser.add_argument(
"--weighting_scheme",
type=str,
@@ -678,11 +634,17 @@ def parse_args():
default=1.29,
help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
)
parser.add_argument(
"--guidance_scale",
"--use_trigflow",
action="store_true",
help="whether to use trigflow in training.",
)
parser.add_argument(
"--sigma_max",
type=float,
default=3.5,
help="the FLUX.1 dev variant is a guidance distilled model",
default=80.0,
help="The max value of sigma in trigflow.",
)
parser.add_argument(
"--gen_update_interval",
@@ -834,6 +796,7 @@ def main():
weight_dtype = torch.bfloat16
args.mixed_precision = accelerator.mixed_precision
args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list]
# Load scheduler, tokenizer and models.
noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
@@ -956,7 +919,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -976,26 +939,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1252,16 +1195,12 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
if fsdp_stage != 0:
from functools import partial
if fsdp_stage != 0 or zero_stage != 0:
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=list(real_score_transformer3d.layers))
real_score_transformer3d = shard_fn(real_score_transformer3d)
if fsdp_stage != 0:
from functools import partial
if fsdp_stage != 0 or zero_stage != 0:
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=list(text_encoder.model.layers))
text_encoder = shard_fn(text_encoder)
@@ -1361,9 +1300,170 @@ def main():
vae_stream_1 = None
vae_stream_2 = None
# Calculate the index we need】
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
# RectifiedFlow Mode
denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)]
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
# TrigFlow Mode
scaling = RectifiedFlow_TrigFlowWrapper(1, args.train_sampling_steps)
sample_trigflow_timesteps_D = partial(
sample_trigflow_timesteps,
P_mean=0.0,
P_std=1.6
)
def denoise(model, xt, timestep, prompt_embeds, noise_scheduler=None, trigflow_scaling=None, multiply_c_in=True):
"""
Unified denoise function supporting both TrigFlow and Rectified Flow
Args:
model: Diffusion model
xt: Noised input (B, C, T, H, W) or (B, C, H, W)
timestep: Timesteps (B,) or (B, 1)
prompt_embeds: Text condition embeddings
noise_scheduler: Noise scheduler (required for Rectified Flow)
trigflow_scaling: TrigFlow scaling function (required for TrigFlow)
multiply_c_in: Whether to multiply c_in with input (TrigFlow only)
Returns:
x0_pred: Predicted clean data
flow_pred: Predicted velocity/flow field
"""
use_trigflow = getattr(args, 'use_trigflow', False)
original_dtype = xt.dtype
device = xt.device
if use_trigflow:
# TrigFlow path
if trigflow_scaling is None:
raise ValueError("trigflow_scaling is required when using trigflow")
ndim = xt.ndim
trigflow_t = timestep
if trigflow_t.ndim == 1:
trigflow_t = trigflow_t.view(-1, 1)
if ndim == 4:
trigflow_t_expanded = trigflow_t.view(-1, 1, 1, 1)
elif ndim == 5:
trigflow_t_expanded = trigflow_t.view(-1, 1, 1, 1, 1)
else:
raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.")
# Get TrigFlow preconditioning coefficients
c_skip, c_out, c_in, c_noise = trigflow_scaling(trigflow_t_expanded)
# Precondition input
if multiply_c_in:
model_input = (xt * c_in).to(xt.dtype)
else:
model_input = xt.to(xt.dtype)
timestep_normalized = c_noise.squeeze(1).squeeze(1).squeeze(1).squeeze(1)
# Model inference
model_output = model(
x=model_input,
cap_feats=prompt_embeds,
t=(1000 - timestep_normalized) / 1000,
)[0]
flow_pred = -model_output.double()
# EDM-style x0 reconstruction
x0_pred = c_skip * xt + c_out * flow_pred
else:
# Rectified Flow path
if noise_scheduler is None:
raise ValueError("scheduler is required for Rectified Flow")
xt_double = xt.double()
timestep = timestep.to(device).double()
timesteps = noise_scheduler.timesteps.to(device).double()
sigmas = noise_scheduler.sigmas.to(device).double()
timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
sigma_t = sigmas[timestep_id]
ndim = xt.ndim
if ndim == 4:
sigma_t_expanded = sigma_t.view(-1, 1, 1, 1)
elif ndim == 5:
sigma_t_expanded = sigma_t.view(-1, 1, 1, 1, 1)
else:
raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.")
model_output = model(
x=xt,
cap_feats=prompt_embeds,
t=(1000 - timestep) / 1000,
)[0]
flow_pred = -model_output.double()
x0_pred = xt_double - sigma_t_expanded * flow_pred
return x0_pred.to(original_dtype), flow_pred
def add_noise(x0, noise, timesteps, noise_scheduler=None):
"""
Unified add noise function supporting both TrigFlow and Rectified Flow
Args:
x0: Clean data
noise: Gaussian noise
timesteps: Timesteps
Returns:
xt: Noised data
"""
use_trigflow = getattr(args, 'use_trigflow', False)
if use_trigflow:
# TrigFlow path: xt = cos(t) * x0 + sin(t) * noise
trigflow_t = timesteps
ndim = x0.ndim
if trigflow_t.ndim == 1:
trigflow_t = trigflow_t.view(-1, 1)
if ndim == 4:
trigflow_t_expanded = trigflow_t.view(-1, 1, 1, 1)
elif ndim == 5:
trigflow_t_expanded = trigflow_t.view(-1, 1, 1, 1, 1)
else:
raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.")
cos_t = torch.cos(trigflow_t_expanded)
sin_t = torch.sin(trigflow_t_expanded)
return cos_t * x0 + sin_t * noise
else:
# Rectified Flow path: xt = (1 - sigma) * x0 + sigma * noise
if noise_scheduler is None:
raise ValueError("noise_scheduler are required for Rectified Flow")
def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype)
schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
timesteps = timesteps.to(accelerator.device)
step_indices = [
torch.argmin(torch.abs(schedule_timesteps - t)).item()
for t in timesteps
]
step_indices = torch.tensor(step_indices, device=accelerator.device)
sigma = sigmas[step_indices].flatten()
while len(sigma.shape) < n_dim:
sigma = sigma.unsqueeze(-1)
return sigma
sigmas = get_sigmas(timesteps, n_dim=x0.ndim, dtype=x0.dtype)
return (1.0 - sigmas) * x0 + sigmas * noise
for epoch in range(first_epoch, args.num_train_epochs):
train_dmd_loss = 0.0
@@ -1420,12 +1520,14 @@ def main():
else:
with torch.no_grad():
prompt_embeds = encode_prompt(
batch['text'], device=accelerator.device,
batch['text'],
device=accelerator.device,
text_encoder=text_encoder,
tokenizer=tokenizer,
)
neg_prompt_embeds = encode_prompt(
["低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。构图混乱。文字模糊，扭曲。"], device=accelerator.device,
["亮度过高，过曝，严重的色彩失真，低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。构图混乱。文字模糊，扭曲。"],
device=accelerator.device,
text_encoder=text_encoder,
tokenizer=tokenizer,
)
@@ -1436,148 +1538,105 @@ def main():
if args.low_vram:
real_score_transformer3d = real_score_transformer3d.to(accelerator.device)
if getattr(args, 'use_trigflow', False):
# Create discrete denoising steps
t_max = torch.arctan(torch.tensor(args.sigma_max))
denoising_step_list = torch.linspace(t_max.item(), 0.0, args.train_sampling_steps)
denoising_step_list = denoising_step_list[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)]
else:
image_seq_len = int(target_shape[-1] // 2 * target_shape[-2] // 2)
mu = calculate_shift(
image_seq_len,
noise_scheduler.config.get("base_image_seq_len", 256),
noise_scheduler.config.get("max_image_seq_len", 4096),
noise_scheduler.config.get("base_shift", 0.5),
noise_scheduler.config.get("max_shift", 1.15),
)
noise_scheduler.sigma_min = 0.0
noise_scheduler.set_timesteps(args.train_sampling_steps, device=accelerator.device, mu=mu)
denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)]
# ==================== Generator Update (DMD) ====================
with accelerator.accumulate(generator_transformer3d):
def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype)
schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
timesteps = timesteps.to(accelerator.device)
step_indices = [
torch.argmin(torch.abs(schedule_timesteps - t)).item()
for t in timesteps
]
step_indices = torch.tensor(step_indices, device=accelerator.device)
sigma = sigmas[step_indices].flatten()
while len(sigma.shape) < n_dim:
sigma = sigma.unsqueeze(-1)
return sigma
def add_noise(latents, noise, timesteps):
sigmas = get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype)
return (1.0 - sigmas) * latents + sigmas * noise
def generate_and_sync_list(num_denoising_steps, device):
indices = torch.randint(low=0, high=num_denoising_steps, size=(1,), generator=torch_rng, device=device)
if dist.is_initialized():
dist.broadcast(indices, src=0)
return indices.tolist()
def convert_flow_pred_to_x0(
scheduler,
flow_pred: torch.Tensor,
xt: torch.Tensor,
timestep: torch.Tensor
) -> torch.Tensor:
"""
Convert flow matching's prediction to x0 prediction.
Supports both 4D [B, C, H, W] and 5D [B, C, F, H, W] inputs.
"""
original_dtype = flow_pred.dtype
device = flow_pred.device
flow_pred = flow_pred.double()
xt = xt.double()
timesteps = scheduler.timesteps.to(device).double()
sigmas = scheduler.sigmas.to(device).double()
timestep = timestep.to(device).double()
timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
sigma_t = sigmas[timestep_id]
ndim = flow_pred.ndim
if ndim == 4:
sigma_t = sigma_t.view(-1, 1, 1, 1)
elif ndim == 5:
sigma_t = sigma_t.view(-1, 1, 1, 1, 1)
else:
raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.")
x0_pred = xt - sigma_t * flow_pred
return x0_pred.to(original_dtype)
# --- Main Training Logic ---
bsz, channel, num_frames, height, width = target_shape
if step % args.gen_update_interval == 0:
generator_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype)
num_denoising_steps = len(denoising_step_list)
final_step_index = generate_and_sync_list(num_denoising_steps, device=generator_noise.device)[0]
# Precompute seq_len once (same for all steps)
for index, current_timestep in enumerate(denoising_step_list):
# Multi-step denoising (backward simulation)
for index in range(num_denoising_steps):
is_final_step = (index == final_step_index)
timestep = torch.full(
generator_noise.shape[:1],
current_timestep,
device=generator_noise.device,
dtype=torch.int64
)
current_t = denoising_step_list[index].expand(bsz).to(accelerator.device)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
context_manager = torch.no_grad() if not is_final_step else contextlib.nullcontext()
with context_manager:
generator_pred = generator_transformer3d(
x=generator_noise,
cap_feats=prompt_embeds,
t=(1000 - timestep) / 1000,
)[0]
generator_pred = -generator_pred
generator_pred = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=generator_pred,
generator_pred, _ = denoise(
model=generator_transformer3d,
xt=generator_noise,
timestep=timestep
timestep=current_t,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling,
multiply_c_in=False if index == 0 else True,
)
if is_final_step:
break
next_timestep = denoising_step_list[index + 1] * torch.ones(
generator_noise.shape[:1], dtype=torch.long, device=generator_noise.device
)
generator_noise = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
next_timestep
)
# Add noise for next step
if index < num_denoising_steps - 1:
next_t = denoising_step_list[index + 1].expand(bsz).to(accelerator.device)
generator_noise = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
next_t,
noise_scheduler=noise_scheduler
)
if getattr(args, 'use_trigflow', False):
# Sample timesteps for discriminator (D distribution)
generator_timestep = sample_trigflow_timesteps_D(bsz, device=accelerator.device)
else:
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
# Add noise to generated samples
generator_denoised_input = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
generator_timestep
generator_timestep,
noise_scheduler=noise_scheduler
).detach().to(accelerator.device, dtype=weight_dtype)
# Compute fake score
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device), torch.no_grad():
fake_score_main_cond = fake_score_transformer3d(
x=generator_denoised_input,
cap_feats=prompt_embeds,
t=(1000 - generator_timestep) / 1000
)[0]
fake_score_main_cond = -fake_score_main_cond
fake_score_main_cond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=fake_score_main_cond,
fake_score_main_cond, _ = denoise(
model=fake_score_transformer3d,
xt=generator_denoised_input,
timestep=generator_timestep
timestep=generator_timestep,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
if args.fake_guidance_scale != 0.0:
fake_score_main_uncond = fake_score_transformer3d(
x=generator_denoised_input,
cap_feats=neg_prompt_embeds,
t=(1000 - generator_timestep) / 1000
)[0]
fake_score_main_uncond = -fake_score_main_uncond
fake_score_main_uncond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=fake_score_main_uncond,
fake_score_main_uncond, _ = denoise(
model=fake_score_transformer3d,
xt=generator_denoised_input,
timestep=generator_timestep
timestep=generator_timestep,
prompt_embeds=neg_prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
fake_score_main = fake_score_main_uncond + (
fake_score_main_cond - fake_score_main_uncond
@@ -1585,32 +1644,24 @@ def main():
else:
fake_score_main = fake_score_main_cond
# Compute real score
real_score_main_cond = real_score_transformer3d(
x=generator_denoised_input,
cap_feats=prompt_embeds,
t=(1000 - generator_timestep) / 1000
)[0]
real_score_main_cond = -real_score_main_cond
real_score_main_cond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=real_score_main_cond,
# Compute real score (teacher)
real_score_main_cond, _ = denoise(
model=real_score_transformer3d,
xt=generator_denoised_input,
timestep=generator_timestep
timestep=generator_timestep,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
if args.real_guidance_scale != 0.0:
real_score_main_uncond = real_score_transformer3d(
x=generator_denoised_input,
cap_feats=neg_prompt_embeds,
t=(1000 - generator_timestep) / 1000
)[0]
real_score_main_uncond = -real_score_main_uncond
real_score_main_uncond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=real_score_main_uncond,
real_score_main_uncond, _ = denoise(
model=real_score_transformer3d,
xt=generator_denoised_input,
timestep=generator_timestep
timestep=generator_timestep,
prompt_embeds=neg_prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
real_score_main = real_score_main_uncond + (
@@ -1622,7 +1673,7 @@ def main():
# DMD loss
fake_to_real_grad = fake_score_main - real_score_main
generator_to_real_norm = generator_pred - real_score_main
normalizer = torch.abs(generator_to_real_norm).mean(dim=[1, 2, 3, 4], keepdim=True)
normalizer = torch.abs(generator_to_real_norm).mean(dim=[1, 2, 3, 4], keepdim=True).clip(min=1e-5)
fake_to_real_grad = fake_to_real_grad / normalizer
fake_to_real_grad = torch.nan_to_num(fake_to_real_grad)
@@ -1657,61 +1708,58 @@ def main():
num_denoising_steps = len(denoising_step_list)
final_step_index = generate_and_sync_list(num_denoising_steps, device=fake_score_critic_noise.device)[0]
for index, current_timestep in enumerate(denoising_step_list):
for index in range(num_denoising_steps):
is_final_step = (index == final_step_index)
timestep = torch.full(
fake_score_critic_noise.shape[:1],
current_timestep,
device=fake_score_critic_noise.device,
dtype=torch.int64
)
current_t = denoising_step_list[index].expand(bsz).to(accelerator.device)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_pred = generator_transformer3d(
x=fake_score_critic_noise,
cap_feats=prompt_embeds,
t=(1000 - timestep) / 1000
)[0]
fake_score_denoised_pred = -fake_score_denoised_pred
fake_score_denoised_pred = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=fake_score_denoised_pred,
fake_score_denoised_pred, _ = denoise(
model=generator_transformer3d,
xt=fake_score_critic_noise,
timestep=timestep
timestep=current_t,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling,
multiply_c_in=False if index == 0 else True,
)
if is_final_step:
break
next_timestep = denoising_step_list[index + 1] * torch.ones(
fake_score_critic_noise.shape[:1],
dtype=torch.long,
device=fake_score_critic_noise.device
)
fake_score_critic_noise = add_noise(
fake_score_denoised_pred,
torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng),
next_timestep
)
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if index < num_denoising_steps - 1:
next_t = denoising_step_list[index + 1].expand(bsz).to(accelerator.device)
fake_score_critic_noise = add_noise(
fake_score_denoised_pred,
torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng),
next_t,
noise_scheduler=noise_scheduler
)
# Sample timesteps for critic
if getattr(args, 'use_trigflow', False):
# Sample timesteps for discriminator (D distribution)
critic_timestep = sample_trigflow_timesteps_D(bsz, device=accelerator.device)
else:
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)
fake_score_denoised_input = add_noise(
fake_score_denoised_pred,
critic_noise,
critic_timestep
critic_timestep,
noise_scheduler=noise_scheduler
)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
cap_feats=prompt_embeds,
t=(1000 - critic_timestep) / 1000
)[0]
fake_score_denoised_output = -fake_score_denoised_output
fake_score_pred, _ = denoise(
model=fake_score_transformer3d,
xt=fake_score_denoised_input,
timestep=critic_timestep,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
noise_pred = noise_pred.float()
@@ -1725,10 +1773,37 @@ def main():
final_loss = masked_loss.mean()
return final_loss
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
# Compute weighting based on sin(t) (following rCM)
if getattr(args, 'use_trigflow', False):
ndim = fake_score_denoised_input.ndim
if critic_timestep.ndim == 1:
critic_timestep_view = critic_timestep.view(-1, 1)
else:
critic_timestep_view = critic_timestep
if ndim == 4:
critic_t_expanded = critic_timestep_view.view(-1, 1, 1, 1)
elif ndim == 5:
critic_t_expanded = critic_timestep_view.view(-1, 1, 1, 1, 1)
sin_t = torch.sin(critic_t_expanded)
weighting = 1.0 / (sin_t ** 2 + 1e-8)
else:
weighting = None
denoising_loss = custom_mse_loss(
fake_score_pred,
fake_score_denoised_pred,
weighting=weighting
)
avg_denoising_loss = accelerator_fake_score_transformer3d.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
if args.low_vram:
generator_transformer3d = generator_transformer3d.to("cpu")
torch.cuda.empty_cache()
accelerator_fake_score_transformer3d.backward(denoising_loss)
if accelerator_fake_score_transformer3d.sync_gradients:
accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
scripts/z_image/train_distill_lora.py
+334 -261
View File
@@ -25,6 +25,7 @@ import pickle
import random
import shutil
import sys
from functools import partial
from typing import (Any, Callable, Dict, List, NamedTuple, Optional, Tuple,
Union)
@@ -55,7 +56,7 @@ from PIL import Image
from torch.distributed.fsdp.fully_sharded_data_parallel import (
FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig,
ShardedStateDictConfig)
from torch.utils.data import Dataset, RandomSampler, BatchSampler
from torch.utils.data import BatchSampler, Dataset, RandomSampler
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from tqdm.auto import tqdm
@@ -84,13 +85,12 @@ from videox_fun.models import (AutoencoderKL, AutoProcessor, AutoTokenizer,
Qwen3ForCausalLM, QwenImageTransformer2DModel,
ZImageTransformer2DModel)
from videox_fun.pipeline import ZImagePipeline
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.lora_utils import (convert_peft_lora_to_kohya_lora,
create_network, merge_lora,
unmerge_lora)
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
save_videos_grid)
from videox_fun.utils import (DiscreteSampling, RectifiedFlow_TrigFlowWrapper,
calculate_dimensions,
convert_peft_lora_to_kohya_lora, create_network,
get_image_latent, get_image_to_video_latent,
merge_lora, sample_trigflow_timesteps,
save_videos_grid, unmerge_lora)
if is_wandb_available():
import wandb
@@ -102,36 +102,6 @@ def filter_kwargs(cls, kwargs):
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
return filtered_kwargs
def linear_decay(initial_value, final_value, total_steps, current_step):
if current_step >= total_steps:
return final_value
current_step = max(0, current_step)
step_size = (final_value - initial_value) / total_steps
current_value = initial_value + step_size * current_step
return current_value
def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=None):
u = torch.normal(mean=0.0, std=1.0, size=shape, device=device, generator=generator)
t = 1 / (1 + torch.exp(-u)) * (high - low) + low
return torch.clip(t.to(torch.int32), low, high - 1)
def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
a1, b1 = 8.73809524e-05, 1.89833333
a2, b2 = 0.00016927, 0.45666666
if image_seq_len > 4300:
mu = a2 * image_seq_len + b2
return float(mu)
m_200 = a2 * image_seq_len + b2
m_10 = a1 * image_seq_len + b1
a = (m_200 - m_10) / 190.0
b = m_200 - 200.0 * a
mu = a * num_steps + b
return float(mu)
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
@@ -211,7 +181,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, a
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -659,6 +629,17 @@ def parse_args():
help=("The module is trained in loras. "),
)
parser.add_argument(
"--use_trigflow",
action="store_true",
help="whether to use trigflow in training.",
)
parser.add_argument(
"--sigma_max",
type=float,
default=80.0,
help="The max value of sigma in trigflow.",
)
parser.add_argument(
"--gen_update_interval",
type=int,
@@ -810,6 +791,7 @@ def main():
args.mixed_precision = accelerator.mixed_precision
# Load scheduler, tokenizer and models.
args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list]
noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler"
@@ -948,7 +930,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -977,15 +959,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -998,34 +972,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1228,7 +1180,7 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler = accelerator_fake_score_transformer3d.prepare(
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
elif fsdp_stage != 0:
else:
generator_transformer3d.network = network
generator_transformer3d = generator_transformer3d.to(dtype=weight_dtype)
generator_transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
@@ -1239,24 +1191,13 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler = accelerator_fake_score_transformer3d.prepare(
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
fake_score_network, critic_optimizer, fake_score_lr_scheduler= accelerator_fake_score_transformer3d.prepare(
fake_score_network, critic_optimizer, fake_score_lr_scheduler
)
if fsdp_stage != 0:
from functools import partial
if fsdp_stage != 0 or zero_stage != 0:
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=list(real_score_transformer3d.layers))
real_score_transformer3d = shard_fn(real_score_transformer3d)
if fsdp_stage != 0:
from functools import partial
if fsdp_stage != 0 or zero_stage != 0:
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=list(text_encoder.model.layers))
text_encoder = shard_fn(text_encoder)
@@ -1368,10 +1309,170 @@ def main():
vae_stream_1 = None
vae_stream_2 = None
# Calculate the index we need】
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
args.denoising_step_indices_list = [int(tmp) for tmp in args.denoising_step_indices_list]
# RectifiedFlow Mode
denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)]
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
# TrigFlow Mode
scaling = RectifiedFlow_TrigFlowWrapper(1, args.train_sampling_steps)
sample_trigflow_timesteps_D = partial(
sample_trigflow_timesteps,
P_mean=0.0,
P_std=1.6
)
def denoise(model, xt, timestep, prompt_embeds, noise_scheduler=None, trigflow_scaling=None, multiply_c_in=True):
"""
Unified denoise function supporting both TrigFlow and Rectified Flow
Args:
model: Diffusion model
xt: Noised input (B, C, T, H, W) or (B, C, H, W)
timestep: Timesteps (B,) or (B, 1)
prompt_embeds: Text condition embeddings
noise_scheduler: Noise scheduler (required for Rectified Flow)
trigflow_scaling: TrigFlow scaling function (required for TrigFlow)
multiply_c_in: Whether to multiply c_in with input (TrigFlow only)
Returns:
x0_pred: Predicted clean data
flow_pred: Predicted velocity/flow field
"""
use_trigflow = getattr(args, 'use_trigflow', False)
original_dtype = xt.dtype
device = xt.device
if use_trigflow:
# TrigFlow path
if trigflow_scaling is None:
raise ValueError("trigflow_scaling is required when using trigflow")
ndim = xt.ndim
trigflow_t = timestep
if trigflow_t.ndim == 1:
trigflow_t = trigflow_t.view(-1, 1)
if ndim == 4:
trigflow_t_expanded = trigflow_t.view(-1, 1, 1, 1)
elif ndim == 5:
trigflow_t_expanded = trigflow_t.view(-1, 1, 1, 1, 1)
else:
raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.")
# Get TrigFlow preconditioning coefficients
c_skip, c_out, c_in, c_noise = trigflow_scaling(trigflow_t_expanded)
# Precondition input
if multiply_c_in:
model_input = (xt * c_in).to(xt.dtype)
else:
model_input = xt.to(xt.dtype)
timestep_normalized = c_noise.squeeze(1).squeeze(1).squeeze(1).squeeze(1)
# Model inference
model_output = model(
x=model_input,
cap_feats=prompt_embeds,
t=(1000 - timestep_normalized) / 1000,
)[0]
flow_pred = -model_output.double()
# EDM-style x0 reconstruction
x0_pred = c_skip * xt + c_out * flow_pred
else:
# Rectified Flow path
if noise_scheduler is None:
raise ValueError("scheduler is required for Rectified Flow")
xt_double = xt.double()
timestep = timestep.to(device).double()
timesteps = noise_scheduler.timesteps.to(device).double()
sigmas = noise_scheduler.sigmas.to(device).double()
timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
sigma_t = sigmas[timestep_id]
ndim = xt.ndim
if ndim == 4:
sigma_t_expanded = sigma_t.view(-1, 1, 1, 1)
elif ndim == 5:
sigma_t_expanded = sigma_t.view(-1, 1, 1, 1, 1)
else:
raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.")
model_output = model(
x=xt,
cap_feats=prompt_embeds,
t=(1000 - timestep) / 1000,
)[0]
flow_pred = -model_output.double()
x0_pred = xt_double - sigma_t_expanded * flow_pred
return x0_pred.to(original_dtype), flow_pred
def add_noise(x0, noise, timesteps, noise_scheduler=None):
"""
Unified add noise function supporting both TrigFlow and Rectified Flow
Args:
x0: Clean data
noise: Gaussian noise
timesteps: Timesteps
Returns:
xt: Noised data
"""
use_trigflow = getattr(args, 'use_trigflow', False)
if use_trigflow:
# TrigFlow path: xt = cos(t) * x0 + sin(t) * noise
trigflow_t = timesteps
ndim = x0.ndim
if trigflow_t.ndim == 1:
trigflow_t = trigflow_t.view(-1, 1)
if ndim == 4:
trigflow_t_expanded = trigflow_t.view(-1, 1, 1, 1)
elif ndim == 5:
trigflow_t_expanded = trigflow_t.view(-1, 1, 1, 1, 1)
else:
raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.")
cos_t = torch.cos(trigflow_t_expanded)
sin_t = torch.sin(trigflow_t_expanded)
return cos_t * x0 + sin_t * noise
else:
# Rectified Flow path: xt = (1 - sigma) * x0 + sigma * noise
if noise_scheduler is None:
raise ValueError("noise_scheduler are required for Rectified Flow")
def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype)
schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
timesteps = timesteps.to(accelerator.device)
step_indices = [
torch.argmin(torch.abs(schedule_timesteps - t)).item()
for t in timesteps
]
step_indices = torch.tensor(step_indices, device=accelerator.device)
sigma = sigmas[step_indices].flatten()
while len(sigma.shape) < n_dim:
sigma = sigma.unsqueeze(-1)
return sigma
sigmas = get_sigmas(timesteps, n_dim=x0.ndim, dtype=x0.dtype)
return (1.0 - sigmas) * x0 + sigmas * noise
for epoch in range(first_epoch, args.num_train_epochs):
train_dmd_loss = 0.0
@@ -1428,12 +1529,14 @@ def main():
else:
with torch.no_grad():
prompt_embeds = encode_prompt(
batch['text'], device=accelerator.device,
batch['text'],
device=accelerator.device,
text_encoder=text_encoder,
tokenizer=tokenizer,
)
neg_prompt_embeds = encode_prompt(
["亮度过高，过曝，严重的色彩失真，低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。构图混乱。文字模糊，扭曲。"], device=accelerator.device,
["亮度过高，过曝，严重的色彩失真，低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。构图混乱。文字模糊，扭曲。"],
device=accelerator.device,
text_encoder=text_encoder,
tokenizer=tokenizer,
)
@@ -1444,147 +1547,105 @@ def main():
if args.low_vram:
real_score_transformer3d = real_score_transformer3d.to(accelerator.device)
if getattr(args, 'use_trigflow', False):
# Create discrete denoising steps
t_max = torch.arctan(torch.tensor(args.sigma_max))
denoising_step_list = torch.linspace(t_max.item(), 0.0, args.train_sampling_steps)
denoising_step_list = denoising_step_list[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)]
else:
image_seq_len = int(target_shape[-1] // 2 * target_shape[-2] // 2)
mu = calculate_shift(
image_seq_len,
noise_scheduler.config.get("base_image_seq_len", 256),
noise_scheduler.config.get("max_image_seq_len", 4096),
noise_scheduler.config.get("base_shift", 0.5),
noise_scheduler.config.get("max_shift", 1.15),
)
noise_scheduler.sigma_min = 0.0
noise_scheduler.set_timesteps(args.train_sampling_steps, device=accelerator.device, mu=mu)
denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)]
# ==================== Generator Update (DMD) ====================
with accelerator.accumulate(generator_transformer3d):
def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype)
schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
timesteps = timesteps.to(accelerator.device)
step_indices = [
torch.argmin(torch.abs(schedule_timesteps - t)).item()
for t in timesteps
]
step_indices = torch.tensor(step_indices, device=accelerator.device)
sigma = sigmas[step_indices].flatten()
while len(sigma.shape) < n_dim:
sigma = sigma.unsqueeze(-1)
return sigma
def add_noise(latents, noise, timesteps):
sigmas = get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype)
return (1.0 - sigmas) * latents + sigmas * noise
def generate_and_sync_list(num_denoising_steps, device):
indices = torch.randint(low=0, high=num_denoising_steps, size=(1,), generator=torch_rng, device=device)
if dist.is_initialized():
dist.broadcast(indices, src=0)
return indices.tolist()
def convert_flow_pred_to_x0(
scheduler,
flow_pred: torch.Tensor,
xt: torch.Tensor,
timestep: torch.Tensor
) -> torch.Tensor:
"""
Convert flow matching's prediction to x0 prediction.
Supports both 4D [B, C, H, W] and 5D [B, C, F, H, W] inputs.
"""
original_dtype = flow_pred.dtype
device = flow_pred.device
flow_pred = flow_pred.double()
xt = xt.double()
timesteps = scheduler.timesteps.to(device).double()
sigmas = scheduler.sigmas.to(device).double()
timestep = timestep.to(device).double()
timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
sigma_t = sigmas[timestep_id]
ndim = flow_pred.ndim
if ndim == 4:
sigma_t = sigma_t.view(-1, 1, 1, 1)
elif ndim == 5:
sigma_t = sigma_t.view(-1, 1, 1, 1, 1)
else:
raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.")
x0_pred = xt - sigma_t * flow_pred
return x0_pred.to(original_dtype)
# --- Main Training Logic ---
bsz, channel, num_frames, height, width = target_shape
if step % args.gen_update_interval == 0:
generator_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype)
num_denoising_steps = len(denoising_step_list)
final_step_index = generate_and_sync_list(num_denoising_steps, device=generator_noise.device)[0]
# Precompute seq_len once (same for all steps)
for index, current_timestep in enumerate(denoising_step_list):
# Multi-step denoising (backward simulation)
for index in range(num_denoising_steps):
is_final_step = (index == final_step_index)
timestep = torch.full(
generator_noise.shape[:1],
current_timestep,
device=generator_noise.device,
dtype=torch.int64
)
current_t = denoising_step_list[index].expand(bsz).to(accelerator.device)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
context_manager = torch.no_grad() if not is_final_step else contextlib.nullcontext()
with context_manager:
generator_pred = generator_transformer3d(
x=generator_noise,
cap_feats=prompt_embeds,
t=(1000 - timestep) / 1000,
)[0]
generator_pred = -generator_pred
generator_pred = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=generator_pred,
generator_pred, _ = denoise(
model=generator_transformer3d,
xt=generator_noise,
timestep=timestep
timestep=current_t,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling,
multiply_c_in=False if index == 0 else True,
)
if is_final_step:
break
next_timestep = denoising_step_list[index + 1] * torch.ones(
generator_noise.shape[:1], dtype=torch.long, device=generator_noise.device
)
generator_noise = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
next_timestep
)
# Add noise for next step
if index < num_denoising_steps - 1:
next_t = denoising_step_list[index + 1].expand(bsz).to(accelerator.device)
generator_noise = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
next_t,
noise_scheduler=noise_scheduler
)
if getattr(args, 'use_trigflow', False):
# Sample timesteps for discriminator (D distribution)
generator_timestep = sample_trigflow_timesteps_D(bsz, device=accelerator.device)
else:
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
# Add noise to generated samples
generator_denoised_input = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
generator_timestep
generator_timestep,
noise_scheduler=noise_scheduler
).detach().to(accelerator.device, dtype=weight_dtype)
# Compute fake score
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device), torch.no_grad():
fake_score_main_cond = fake_score_transformer3d(
x=generator_denoised_input,
cap_feats=prompt_embeds,
t=(1000 - generator_timestep) / 1000
)[0]
fake_score_main_cond = -fake_score_main_cond
fake_score_main_cond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=fake_score_main_cond,
fake_score_main_cond, _ = denoise(
model=fake_score_transformer3d,
xt=generator_denoised_input,
timestep=generator_timestep
timestep=generator_timestep,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
if args.fake_guidance_scale != 0.0:
fake_score_main_uncond = fake_score_transformer3d(
x=generator_denoised_input,
cap_feats=neg_prompt_embeds,
t=(1000 - generator_timestep) / 1000
)[0]
fake_score_main_uncond = -fake_score_main_uncond
fake_score_main_uncond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=fake_score_main_uncond,
fake_score_main_uncond, _ = denoise(
model=fake_score_transformer3d,
xt=generator_denoised_input,
timestep=generator_timestep
timestep=generator_timestep,
prompt_embeds=neg_prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
fake_score_main = fake_score_main_uncond + (
fake_score_main_cond - fake_score_main_uncond
@@ -1592,32 +1653,24 @@ def main():
else:
fake_score_main = fake_score_main_cond
# Compute real score
real_score_main_cond = real_score_transformer3d(
x=generator_denoised_input,
cap_feats=prompt_embeds,
t=(1000 - generator_timestep) / 1000
)[0]
real_score_main_cond = -real_score_main_cond
real_score_main_cond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=real_score_main_cond,
# Compute real score (teacher)
real_score_main_cond, _ = denoise(
model=real_score_transformer3d,
xt=generator_denoised_input,
timestep=generator_timestep
timestep=generator_timestep,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
if args.real_guidance_scale != 0.0:
real_score_main_uncond = real_score_transformer3d(
x=generator_denoised_input,
cap_feats=neg_prompt_embeds,
t=(1000 - generator_timestep) / 1000
)[0]
real_score_main_uncond = -real_score_main_uncond
real_score_main_uncond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=real_score_main_uncond,
real_score_main_uncond, _ = denoise(
model=real_score_transformer3d,
xt=generator_denoised_input,
timestep=generator_timestep
timestep=generator_timestep,
prompt_embeds=neg_prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
real_score_main = real_score_main_uncond + (
@@ -1629,7 +1682,7 @@ def main():
# DMD loss
fake_to_real_grad = fake_score_main - real_score_main
generator_to_real_norm = generator_pred - real_score_main
normalizer = torch.abs(generator_to_real_norm).mean(dim=[1, 2, 3, 4], keepdim=True)
normalizer = torch.abs(generator_to_real_norm).mean(dim=[1, 2, 3, 4], keepdim=True).clip(min=1e-5)
fake_to_real_grad = fake_to_real_grad / normalizer
fake_to_real_grad = torch.nan_to_num(fake_to_real_grad)
@@ -1664,61 +1717,58 @@ def main():
num_denoising_steps = len(denoising_step_list)
final_step_index = generate_and_sync_list(num_denoising_steps, device=fake_score_critic_noise.device)[0]
for index, current_timestep in enumerate(denoising_step_list):
for index in range(num_denoising_steps):
is_final_step = (index == final_step_index)
timestep = torch.full(
fake_score_critic_noise.shape[:1],
current_timestep,
device=fake_score_critic_noise.device,
dtype=torch.int64
)
current_t = denoising_step_list[index].expand(bsz).to(accelerator.device)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_pred = generator_transformer3d(
x=fake_score_critic_noise,
cap_feats=prompt_embeds,
t=(1000 - timestep) / 1000
)[0]
fake_score_denoised_pred = -fake_score_denoised_pred
fake_score_denoised_pred = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=fake_score_denoised_pred,
fake_score_denoised_pred, _ = denoise(
model=generator_transformer3d,
xt=fake_score_critic_noise,
timestep=timestep
timestep=current_t,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling,
multiply_c_in=False if index == 0 else True,
)
if is_final_step:
break
next_timestep = denoising_step_list[index + 1] * torch.ones(
fake_score_critic_noise.shape[:1],
dtype=torch.long,
device=fake_score_critic_noise.device
)
fake_score_critic_noise = add_noise(
fake_score_denoised_pred,
torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng),
next_timestep
)
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if index < num_denoising_steps - 1:
next_t = denoising_step_list[index + 1].expand(bsz).to(accelerator.device)
fake_score_critic_noise = add_noise(
fake_score_denoised_pred,
torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng),
next_t,
noise_scheduler=noise_scheduler
)
# Sample timesteps for critic
if getattr(args, 'use_trigflow', False):
# Sample timesteps for discriminator (D distribution)
critic_timestep = sample_trigflow_timesteps_D(bsz, device=accelerator.device)
else:
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)
fake_score_denoised_input = add_noise(
fake_score_denoised_pred,
critic_noise,
critic_timestep
critic_timestep,
noise_scheduler=noise_scheduler
)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
cap_feats=prompt_embeds,
t=(1000 - critic_timestep) / 1000
)[0]
fake_score_denoised_output = -fake_score_denoised_output
fake_score_pred, _ = denoise(
model=fake_score_transformer3d,
xt=fake_score_denoised_input,
timestep=critic_timestep,
prompt_embeds=prompt_embeds,
noise_scheduler=noise_scheduler,
trigflow_scaling=scaling
)
def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
noise_pred = noise_pred.float()
@@ -1732,7 +1782,30 @@ def main():
final_loss = masked_loss.mean()
return final_loss
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
# Compute weighting based on sin(t) (following rCM)
if getattr(args, 'use_trigflow', False):
ndim = fake_score_denoised_input.ndim
if critic_timestep.ndim == 1:
critic_timestep_view = critic_timestep.view(-1, 1)
else:
critic_timestep_view = critic_timestep
if ndim == 4:
critic_t_expanded = critic_timestep_view.view(-1, 1, 1, 1)
elif ndim == 5:
critic_t_expanded = critic_timestep_view.view(-1, 1, 1, 1, 1)
sin_t = torch.sin(critic_t_expanded)
weighting = 1.0 / (sin_t ** 2 + 1e-8)
else:
weighting = None
denoising_loss = custom_mse_loss(
fake_score_pred,
fake_score_denoised_pred,
weighting=weighting
)
avg_denoising_loss = accelerator_fake_score_transformer3d.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
scripts/z_image/train_lora.py
+8 -49
View File
@@ -206,7 +206,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, network, args, a
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -850,7 +850,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -879,15 +879,7 @@ def main():
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
@@ -900,34 +892,12 @@ def main():
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
network_state_dict = accelerate_state_dict
network_state_dict = {}
for key in accelerate_state_dict:
if "network" in key:
network_state_dict[key.replace("network.", "")] = accelerate_state_dict[key].to(weight_dtype)
save_file(network_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
safetensor_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model.safetensors")
if args.use_peft_lora:
network_state_dict = get_peft_model_state_dict(accelerator.unwrap_model(models[-1]))
save_model(safetensor_save_path, network_state_dict)
network_state_dict_kohya = convert_peft_lora_to_kohya_lora(network_state_dict)
safetensor_kohya_format_save_path = os.path.join(output_dir, f"lora_diffusion_pytorch_model_compatible_with_comfyui.safetensors")
save_model(safetensor_kohya_format_save_path, network_state_dict_kohya)
else:
save_model(safetensor_save_path, accelerator.unwrap_model(models[-1]))
if not args.use_deepspeed:
for _ in range(len(weights)):
weights.pop()
@@ -1207,23 +1177,12 @@ def main():
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
elif fsdp_stage != 0:
else:
transformer3d.network = network
transformer3d = transformer3d.to(dtype=weight_dtype)
transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
transformer3d, optimizer, train_dataloader, lr_scheduler
)
else:
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
network, optimizer, train_dataloader, lr_scheduler
)
if zero_stage != 0 and not args.use_peft_lora:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=list(transformer3d.layers))
transformer3d = shard_fn(transformer3d)
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
scripts/z_image_fun/train_control.py
+2 -22
View File
@@ -207,7 +207,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -928,7 +928,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -948,26 +948,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
scripts/z_image_fun/train_control_distill.py
+16 -24
View File
@@ -213,7 +213,7 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d,
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
scheduler=scheduler,
)
pipeline = pipeline.to(accelerator.device)
@@ -987,7 +987,7 @@ def main():
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
@@ -1007,26 +1007,6 @@ def main():
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
elif zero_stage == 3:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
else:
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
@@ -1420,14 +1400,14 @@ def main():
fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler
)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=list(real_score_transformer3d.layers))
real_score_transformer3d = shard_fn(real_score_transformer3d)
if fsdp_stage != 0:
if fsdp_stage != 0 or zero_stage != 0:
from functools import partial
from videox_fun.dist import set_multi_gpus_devices, shard_model
@@ -1662,6 +1642,18 @@ def main():
if args.low_vram:
real_score_transformer3d = real_score_transformer3d.to(accelerator.device)
image_seq_len = int(target_shape[-1] // 2 * target_shape[-2] // 2)
mu = calculate_shift(
image_seq_len,
noise_scheduler.config.get("base_image_seq_len", 256),
noise_scheduler.config.get("max_image_seq_len", 4096),
noise_scheduler.config.get("base_shift", 0.5),
noise_scheduler.config.get("max_shift", 1.15),
)
noise_scheduler.sigma_min = 0.0
noise_scheduler.set_timesteps(args.train_sampling_steps, device=accelerator.device, mu=mu)
denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)]
with accelerator.accumulate(generator_transformer3d):
def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype)
videox_fun/models/cogvideox_transformer3d.py
+2 -1
View File
@@ -182,8 +182,9 @@ class CogVideoXPatchEmbed(nn.Module):
post_time_compression_frames,
self.spatial_interpolation_scale,
self.temporal_interpolation_scale,
output_type="pt",
)
pos_embedding = torch.from_numpy(pos_embedding).flatten(0, 1)
pos_embedding = pos_embedding.flatten(0, 1)
joint_pos_embedding = torch.zeros(
1, self.max_text_seq_length + num_patches, self.embed_dim, requires_grad=False
)
videox_fun/utils/__init__.py
+7 -4
View File
@@ -11,10 +11,13 @@ from .fp8_optimization import (autocast_model_forward,
from .group_offload import (register_auto_device_hook,
safe_enable_group_offload,
safe_remove_group_offloading)
from .lora_utils import merge_lora, unmerge_lora
from .utils import (filter_kwargs, get_autocast_dtype, get_image_latent,
get_image_to_video_latent, get_video_to_video_latent,
save_videos_grid)
from .lora_utils import (convert_peft_lora_to_kohya_lora, create_network,
merge_lora, unmerge_lora)
from .trigflow_sampler import (RectifiedFlow_TrigFlowWrapper,
sample_trigflow_timesteps)
from .utils import (calculate_dimensions, filter_kwargs, get_autocast_dtype,
get_image_latent, get_image_to_video_latent,
get_video_to_video_latent, save_videos_grid)
# The pai_fuser is an internally developed acceleration package, which can be used on PAI.
if importlib.util.find_spec("paifuser") is not None:
videox_fun/utils/trigflow_sampler.py
+23
View File
@@ -0,0 +1,23 @@
import torch
# Copied from https://github.com/NVlabs/rcm/blob/main/rcm/utils/denoiser_scaling.py
class RectifiedFlow_TrigFlowWrapper:
def __init__(self, sigma_data: float = 1.0, t_scaling_factor: float = 1.0):
assert abs(sigma_data - 1.0) < 1e-6, "sigma_data must be 1.0 for RectifiedFlowScaling"
self.t_scaling_factor = t_scaling_factor
def __call__(self, trigflow_t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
trigflow_t = trigflow_t.to(torch.float64)
c_skip = 1 / (torch.cos(trigflow_t) + torch.sin(trigflow_t))
c_out = -1 * torch.sin(trigflow_t) / (torch.cos(trigflow_t) + torch.sin(trigflow_t))
c_in = 1 / (torch.cos(trigflow_t) + torch.sin(trigflow_t))
c_noise = (torch.sin(trigflow_t) / (torch.cos(trigflow_t) + torch.sin(trigflow_t))) * self.t_scaling_factor
return c_skip, c_out, c_in, c_noise
# Sample timesteps
def sample_trigflow_timesteps(batch_size, device, P_mean=0.0, P_std=1.6):
"""Sample timesteps for training"""
sigma = torch.randn(batch_size, device=device)
sigma = (sigma * P_std + P_mean).exp()
timesteps = torch.arctan(sigma)
return timesteps