nodes/sylym-comfy_vid2vid
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

first modification

Browse Source
This commit is contained in:
sylym
2023-03-28 00:34:18 +08:00
parent 686791ef70
commit 8a005a02ee
3 changed files with 67 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train_tuneavideo.py
+35 -127
View File
@@ -1,29 +1,20 @@
import argparse
import inspect
import math
import os
from typing import Dict, Optional, Tuple
from omegaconf import OmegaConf
from typing import Optional, Tuple
import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch.utils.data import Dataset
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from diffusers import AutoencoderKL, DDPMScheduler
from diffusers.optimization import get_scheduler
from diffusers.utils import check_min_version
from diffusers.utils.import_utils import is_xformers_available
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer
from tuneavideo.models.unet import UNet3DConditionModel
from tuneavideo.data.dataset import TuneAVideoDataset
from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline
from einops import rearrange
import shutil
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
@@ -31,13 +22,27 @@ check_min_version("0.10.0.dev0")
logger = get_logger(__name__, log_level="INFO")
class TuneAVideoDataset(Dataset):
def __init__(self):
self.prompt_ids = None
self.pixel_values = None
def __len__(self):
return 1
def __getitem__(self, index):
def main(
pretrained_model_path: str,
pretrained_vae_path: str,
output_dir: str,
train_data: Dict,
inference_data: Dict = None,
example = {
"pixel_values": self.pixel_values,
"prompt_ids": self.prompt_ids
}
return example
def train(
model,
noise_scheduler,
samples,
context,
device,
trainable_modules: Tuple[str] = (
"attn1.to_q",
"attn2.to_q",
@@ -56,13 +61,8 @@ def main(
max_grad_norm: float = 1.0,
gradient_accumulation_steps: int = 1,
gradient_checkpointing: bool = True,
checkpointing_steps: int = 5000,
resume_from_checkpoint: Optional[str] = None,
mixed_precision: Optional[str] = "fp16",
use_8bit_adam: bool = False,
enable_xformers_memory_efficient_attention: bool = True,
seed: Optional[int] = None,
gpu_id: str = "0",
):
*_, config = inspect.getargvalues(inspect.currentframe())
@@ -76,34 +76,13 @@ def main(
if seed is not None:
set_seed(seed)
# Handle the output folder creation
if accelerator.is_main_process:
os.makedirs(output_dir, exist_ok=True)
OmegaConf.save(config, os.path.join(output_dir, 'config.yaml'))
# Load scheduler, tokenizer and models.
noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder")
vae = AutoencoderKL.from_pretrained(pretrained_vae_path, subfolder="vae")
unet = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet").to(f"cuda:{gpu_id}")
# Freeze vae and text_encoder
vae.requires_grad_(False)
text_encoder.requires_grad_(False)
unet = model.model.model.diffusion_model.to(device)
unet.requires_grad_(False)
for name, module in unet.named_modules():
if name.endswith(tuple(trainable_modules)):
for params in module.parameters():
params.requires_grad = True
if enable_xformers_memory_efficient_attention:
if is_xformers_available():
unet.enable_xformers_memory_efficient_attention()
else:
raise ValueError("xformers is not available. Make sure it is installed correctly")
if gradient_checkpointing:
unet.enable_gradient_checkpointing()
@@ -112,18 +91,7 @@ def main(
learning_rate * gradient_accumulation_steps * train_batch_size * accelerator.num_processes
)
# Initialize the optimizer
if use_8bit_adam:
try:
import bitsandbytes as bnb
except ImportError:
raise ImportError(
"Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
)
optimizer_cls = bnb.optim.AdamW8bit
else:
optimizer_cls = torch.optim.AdamW
optimizer_cls = torch.optim.AdamW
optimizer = optimizer_cls(
unet.parameters(),
@@ -134,12 +102,11 @@ def main(
)
# Get the training dataset
train_dataset = TuneAVideoDataset(**train_data)
train_dataset = TuneAVideoDataset()
# Preprocessing the dataset
train_dataset.prompt_ids = tokenizer(
train_dataset.prompt, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
).input_ids[0]
train_dataset.prompt_ids = context
train_dataset.pixel_values = samples
# DataLoaders creation:
train_dataloader = torch.utils.data.DataLoader(
@@ -167,10 +134,6 @@ def main(
elif accelerator.mixed_precision == "bf16":
weight_dtype = torch.bfloat16
# Move text_encode and vae to gpu and cast to weight_dtype
text_encoder.to(f"cuda:{gpu_id}", dtype=weight_dtype)
vae.to(f"cuda:{gpu_id}", dtype=weight_dtype)
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
# Afterwards we recalculate our number of training epochs
@@ -185,23 +148,6 @@ def main(
global_step = 0
first_epoch = 0
# Potentially load in the weights and states from a previous save
if resume_from_checkpoint:
if resume_from_checkpoint != "latest":
path = os.path.basename(resume_from_checkpoint)
else:
# Get the most recent checkpoint
dirs = os.listdir(output_dir)
dirs = [d for d in dirs if d.startswith("checkpoint")]
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
path = dirs[-1]
accelerator.print(f"Resuming from checkpoint {path}")
accelerator.load_state(os.path.join(output_dir, path))
global_step = int(path.split("-")[1])
first_epoch = global_step // num_update_steps_per_epoch
resume_step = global_step % num_update_steps_per_epoch
# Only show the progress bar once on each machine.
progress_bar = tqdm(range(global_step, max_train_steps), disable=not accelerator.is_local_main_process)
progress_bar.set_description("Steps")
@@ -210,20 +156,10 @@ def main(
unet.train()
train_loss = 0.0
for step, batch in enumerate(train_dataloader):
# Skip steps until we reach the resumed step
if resume_from_checkpoint and epoch == first_epoch and step < resume_step:
if step % gradient_accumulation_steps == 0:
progress_bar.update(1)
continue
with accelerator.accumulate(unet):
# Convert videos to latent space
pixel_values = batch["pixel_values"].to(weight_dtype).to(f"cuda:{gpu_id}")
video_length = pixel_values.shape[1]
pixel_values = rearrange(pixel_values, "b f c h w -> (b f) c h w")
latents = vae.encode(pixel_values).latent_dist.sample()
latents = rearrange(latents, "(b f) c h w -> b c f h w", f=video_length)
latents = latents * 0.18215
latents = batch["pixel_values"].type(weight_dtype).to(device)
# Sample noise that we'll add to the latents
noise = torch.randn_like(latents)
@@ -237,7 +173,7 @@ def main(
noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
# Get the text embedding for conditioning
encoder_hidden_states = text_encoder(batch["prompt_ids"].to(f"cuda:{gpu_id}"))[0]
encoder_hidden_states = batch["prompt_ids"].type(weight_dtype).to(device)
# Get the target for loss depending on the prediction type
if noise_scheduler.prediction_type == "epsilon":
@@ -248,7 +184,9 @@ def main(
raise ValueError(f"Unknown prediction type {noise_scheduler.prediction_type}")
# Predict the noise residual and compute loss
model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
noisy_latents = rearrange(noisy_latents.squeeze(0), "c f h w -> f c h w")
model_pred = unet(noisy_latents, timesteps, encoder_hidden_states)
model_pred = rearrange(model_pred.unsqueeze(0), "b f c h w -> b c f h w")
loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
# Gather the losses across all processes for logging (if we use distributed training).
@@ -270,12 +208,6 @@ def main(
accelerator.log({"train_loss": train_loss}, step=global_step)
train_loss = 0.0
if global_step % checkpointing_steps == 0:
if accelerator.is_main_process:
save_path = os.path.join(output_dir, f"checkpoint-{global_step}")
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -284,31 +216,7 @@ def main(
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
if accelerator.is_main_process:
unet = accelerator.unwrap_model(unet)
pipeline = TuneAVideoPipeline.from_pretrained(
pretrained_model_path,
text_encoder=text_encoder,
vae=vae,
unet=unet,
)
pipeline.save_pretrained(output_dir)
unet = accelerator.unwrap_model(unet)
accelerator.end_training()
# 删除冗余文件夹
shutil.rmtree(os.path.join(output_dir, "scheduler"))
shutil.rmtree(os.path.join(output_dir, "text_encoder"))
shutil.rmtree(os.path.join(output_dir, "vae"))
shutil.rmtree(os.path.join(output_dir, "tokenizer"))
# 删除冗余文件
os.remove(os.path.join(output_dir, "config.yaml"))
os.remove(os.path.join(output_dir, "model_index.json"))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="./configs/Continuous_frame_challenge.yaml")
parser.add_argument("--cuda", type=str, default="0")
args = parser.parse_args()
main(**OmegaConf.load(args.config), gpu_id=args.cuda)
model.model.model.diffusion_model = unet.to(torch.device("cpu"))
return model