diff --git a/easyanimate/utils/IDDIM.py b/easyanimate/utils/IDDIM.py deleted file mode 100644 index ec99e87..0000000 --- a/easyanimate/utils/IDDIM.py +++ /dev/null @@ -1,51 +0,0 @@ -# Modified from OpenAI's diffusion repos -# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py -# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion -# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py -from . import gaussian_diffusion as gd -from .respace import SpacedDiffusion, space_timesteps - - -def IDDPM( - timestep_respacing, - noise_schedule="linear", - use_kl=False, - sigma_small=False, - predict_xstart=False, - learn_sigma=True, - pred_sigma=True, - rescale_learned_sigmas=False, - diffusion_steps=1000, - snr=False, - return_startx=False, -): - betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps) - if use_kl: - loss_type = gd.LossType.RESCALED_KL - elif rescale_learned_sigmas: - loss_type = gd.LossType.RESCALED_MSE - else: - loss_type = gd.LossType.MSE - if timestep_respacing is None or timestep_respacing == "": - timestep_respacing = [diffusion_steps] - return SpacedDiffusion( - use_timesteps=space_timesteps(diffusion_steps, timestep_respacing), - betas=betas, - model_mean_type=( - gd.ModelMeanType.START_X if predict_xstart else gd.ModelMeanType.EPSILON - ), - model_var_type=( - (gd.ModelVarType.LEARNED_RANGE if learn_sigma else ( - gd.ModelVarType.FIXED_LARGE - if not sigma_small - else gd.ModelVarType.FIXED_SMALL - ) - ) - if pred_sigma - else None - ), - loss_type=loss_type, - snr=snr, - return_startx=return_startx, - # rescale_timesteps=rescale_timesteps, - ) \ No newline at end of file diff --git a/scripts/train_t2iv.py b/scripts/train_t2iv.py index 9479018..1371e1d 100644 --- a/scripts/train_t2iv.py +++ b/scripts/train_t2iv.py @@ -20,6 +20,7 @@ import gc import logging import math import os +import pickle import shutil import sys @@ -50,7 +51,6 @@ from tqdm.auto import tqdm from transformers import T5EncoderModel, T5Tokenizer from transformers.utils import ContextManagers -import pickle import datasets current_file_path = os.path.abspath(__file__) @@ -61,9 +61,9 @@ for project_root in project_roots: from easyanimate.data.bucket_sampler import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, ASPECT_RATIO_RANDOM_CROP_PROB, - AspectRatioBatchImageSampler, RandomSampler, + AspectRatioBatchImageSampler, AspectRatioBatchImageVideoSampler, - get_closest_ratio) + RandomSampler, get_closest_ratio) from easyanimate.data.dataset_image import CC15M from easyanimate.data.dataset_image_video import (ImageVideoDataset, ImageVideoSampler) @@ -73,7 +73,8 @@ from easyanimate.models.transformer3d import Transformer3DModel from easyanimate.pipeline.pipeline_easyanimate import EasyAnimatePipeline from easyanimate.pipeline.pipeline_pixart_magvit import \ PixArtAlphaMagvitPipeline -from easyanimate.utils.IDDIM import IDDPM +from easyanimate.utils import gaussian_diffusion as gd +from easyanimate.utils.respace import SpacedDiffusion, space_timesteps from easyanimate.utils.utils import save_videos_grid if is_wandb_available(): @@ -592,7 +593,11 @@ def main(): # Load scheduler, tokenizer and models. # noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") - train_diffusion = IDDPM(str(args.train_sampling_steps), learn_sigma=True, pred_sigma=True, snr=args.snr_loss) + train_diffusion = SpacedDiffusion( + use_timesteps=space_timesteps(1000, str(args.train_sampling_steps)), betas=gd.get_named_beta_schedule("linear", 1000), + model_mean_type=(gd.ModelMeanType.EPSILON), model_var_type=((gd.ModelVarType.LEARNED_RANGE)), + loss_type=gd.LossType.MSE, snr=args.snr_loss, return_startx=False, + ) tokenizer = T5Tokenizer.from_pretrained( args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision ) diff --git a/scripts/train_t2iv_lora.py b/scripts/train_t2iv_lora.py index b509b0a..ac002d0 100644 --- a/scripts/train_t2iv_lora.py +++ b/scripts/train_t2iv_lora.py @@ -20,6 +20,7 @@ import gc import logging import math import os +import pickle import shutil import sys @@ -50,7 +51,6 @@ from tqdm.auto import tqdm from transformers import T5EncoderModel, T5Tokenizer from transformers.utils import ContextManagers -import pickle import datasets current_file_path = os.path.abspath(__file__) @@ -61,33 +61,26 @@ for project_root in project_roots: from easyanimate.data.bucket_sampler import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, ASPECT_RATIO_RANDOM_CROP_PROB, - AspectRatioBatchImageSampler, RandomSampler, + AspectRatioBatchImageSampler, AspectRatioBatchImageVideoSampler, - get_closest_ratio) + AspectRatioBatchSampler, + RandomSampler, get_closest_ratio) from easyanimate.data.dataset_image import CC15M from easyanimate.data.dataset_image_video import (ImageVideoDataset, ImageVideoSampler) +from easyanimate.data.dataset_video import VideoDataset, WebVid10M from easyanimate.models.autoencoder_magvit import AutoencoderKLMagvit from easyanimate.models.transformer2d import Transformer2DModel from easyanimate.models.transformer3d import Transformer3DModel from easyanimate.pipeline.pipeline_easyanimate import EasyAnimatePipeline +from easyanimate.pipeline.pipeline_easyanimate_inpaint import \ + EasyAnimateInpaintPipeline from easyanimate.pipeline.pipeline_pixart_magvit import \ PixArtAlphaMagvitPipeline -from easyanimate.utils.IDDIM import IDDPM -from easyanimate.utils.utils import save_videos_grid - -from easyanimate.data.bucket_sampler import (ASPECT_RATIO_512, - ASPECT_RATIO_RANDOM_CROP_512, - ASPECT_RATIO_RANDOM_CROP_PROB, - AspectRatioBatchSampler, - get_closest_ratio) -from easyanimate.data.dataset_video import VideoDataset, WebVid10M -from easyanimate.models.autoencoder_magvit import AutoencoderKLMagvit -from easyanimate.models.transformer3d import Transformer3DModel -from easyanimate.pipeline.pipeline_easyanimate import EasyAnimatePipeline -from easyanimate.pipeline.pipeline_easyanimate_inpaint import EasyAnimateInpaintPipeline -from easyanimate.utils.IDDIM import IDDPM -from easyanimate.utils.lora_utils import create_network, merge_lora, unmerge_lora +from easyanimate.utils import gaussian_diffusion as gd +from easyanimate.utils.lora_utils import (create_network, merge_lora, + unmerge_lora) +from easyanimate.utils.respace import SpacedDiffusion, space_timesteps from easyanimate.utils.utils import save_videos_grid if is_wandb_available(): @@ -619,7 +612,11 @@ def main(): # Load scheduler, tokenizer and models. # noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") - train_diffusion = IDDPM(str(args.train_sampling_steps), learn_sigma=True, pred_sigma=True, snr=args.snr_loss) + train_diffusion = SpacedDiffusion( + use_timesteps=space_timesteps(1000, str(args.train_sampling_steps)), betas=gd.get_named_beta_schedule("linear", 1000), + model_mean_type=(gd.ModelMeanType.EPSILON), model_var_type=((gd.ModelVarType.LEARNED_RANGE)), + loss_type=gd.LossType.MSE, snr=args.snr_loss, return_startx=False, + ) tokenizer = T5Tokenizer.from_pretrained( args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision )