From 7c5b0a6ccdfbcf75355f5ddc98a144dfe1cd312e Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Thu, 24 Oct 2024 21:35:27 +0000 Subject: [PATCH] Remove all npu code --- opensora/acceleration/__init__.py | 0 opensora/acceleration/communications.py | 133 --- opensora/acceleration/parallel_states.py | 57 -- opensora/dataset/t2v_datasets.py | 22 - .../model/causal_vae/modeling_causalvae.py | 24 +- .../causalvideovae/model/modules/attention.py | 11 +- .../causalvideovae/model/modules/conv.py | 28 +- .../model/modules/resnet_block.py | 16 +- .../model/modules/updownsample.py | 97 +-- .../diffusion/opensora/modeling_inpaint.py | 9 - .../diffusion/opensora/modeling_opensora.py | 32 +- opensora/models/diffusion/opensora/modules.py | 344 +++----- opensora/models/diffusion/opensora/rope.py | 15 +- .../models/diffusion/udit/modeling_udit.py | 9 - opensora/models/diffusion/udit/modules.py | 169 ++-- opensora/sample/pipeline_opensora.py | 817 ++++++++++++++++++ opensora/sample/sample_t2v_sp.py | 57 +- opensora/train/train_t2v_diffusers.py | 6 +- scripts/slurm/placeholder | 0 scripts/text_condition/sample_inpaint.sh | 18 - scripts/text_condition/sample_inpaint_ddp.sh | 21 - scripts/text_condition/sample_inpaint_sp.sh | 21 - scripts/text_condition/train_inpaint.sh | 74 -- scripts/train_data/merge_data_on_npu.txt | 1 - 24 files changed, 1045 insertions(+), 936 deletions(-) delete mode 100644 opensora/acceleration/__init__.py delete mode 100644 opensora/acceleration/communications.py delete mode 100644 opensora/acceleration/parallel_states.py create mode 100644 opensora/sample/pipeline_opensora.py delete mode 100644 scripts/slurm/placeholder delete mode 100644 scripts/text_condition/sample_inpaint.sh delete mode 100644 scripts/text_condition/sample_inpaint_ddp.sh delete mode 100644 scripts/text_condition/sample_inpaint_sp.sh delete mode 100644 scripts/text_condition/train_inpaint.sh delete mode 100644 scripts/train_data/merge_data_on_npu.txt diff --git a/opensora/acceleration/__init__.py b/opensora/acceleration/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/opensora/acceleration/communications.py b/opensora/acceleration/communications.py deleted file mode 100644 index 6f69a73a0..000000000 --- a/opensora/acceleration/communications.py +++ /dev/null @@ -1,133 +0,0 @@ -import torch -import torch.distributed as dist -from einops import rearrange -from opensora.acceleration.parallel_states import hccl_info, lccl_info, enable_LCCL -try: - from lcalib.functional import lcal_all2allvc -except: - lcal_all2allvc = None - -def broadcast(input_: torch.Tensor): - sp_size = hccl_info.world_size - src = hccl_info.rank // sp_size * sp_size - dist.broadcast(input_, src=src, group=hccl_info.group) - -_COUNT = 0 -def _all_to_all( - input_: torch.Tensor, - scatter_dim: int, - gather_dim: int, -): - group = hccl_info.group - sp_size = hccl_info.world_size - input_list = [t.contiguous() for t in torch.tensor_split(input_, sp_size, scatter_dim)] - output_list = [torch.empty_like(input_list[0]) for _ in range(sp_size)] - dist.all_to_all(output_list, input_list, group=group) - return torch.cat(output_list, dim=gather_dim).contiguous() - -def _single_all_to_all( - input_: torch.Tensor, - scatter_dim: int, - gather_dim: int, - enable_HCCL=False, -): - if enable_LCCL: - sp_size = lccl_info.world_size - else: - sp_size = hccl_info.world_size - inp_shape = list(input_.shape) - inp_shape[scatter_dim] = inp_shape[scatter_dim] // sp_size - if scatter_dim < 1: - input_t = input_.reshape( - [sp_size, inp_shape[scatter_dim]] + \ - inp_shape[scatter_dim + 1:] - ) - else: - # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! - input_t = input_.reshape( - [-1, sp_size, inp_shape[scatter_dim]] + \ - inp_shape[scatter_dim + 1:] - ).transpose(0, 1).contiguous() - - output = torch.empty_like(input_t) - if enable_LCCL and not enable_HCCL: - matrix_count = torch.ones([sp_size, sp_size], dtype=torch.int64, device=input_t.device) * ( - input_t.numel() // sp_size) - lcal_all2allvc(input_t, output, matrix_count, lccl_info.group) - else: - dist.all_to_all_single(output, input_t, group=hccl_info.group) - # if scattering the seq-dim, transpose the heads back to the original dimension - if scatter_dim < 1: - output = output.transpose(0, 1).contiguous() - - return output.reshape( - inp_shape[: gather_dim] + [inp_shape[gather_dim] * sp_size, ] + inp_shape[gather_dim + 1:]) - - -class _AllToAll(torch.autograd.Function): - """All-to-all communication. - - Args: - input_: input matrix - process_group: communication group - scatter_dim: scatter dimension - gather_dim: gather dimension - """ - - @staticmethod - def forward(ctx, input_, scatter_dim, gather_dim, all_to_all_func): - ctx.scatter_dim = scatter_dim - ctx.gather_dim = gather_dim - ctx.all_to_all = all_to_all_func - output = ctx.all_to_all(input_, scatter_dim, gather_dim) - return output - - @staticmethod - def backward(ctx, grad_output): - grad_output = ctx.all_to_all( - grad_output, - ctx.gather_dim, - ctx.scatter_dim, - ) - return ( - grad_output, - None, - None, - None, - ) - -def all_to_all_SBH( - input_: torch.Tensor, - scatter_dim: int = 1, - gather_dim: int = 0, -): - return _AllToAll.apply(input_, scatter_dim, gather_dim, _single_all_to_all) - -def all_to_all_BSND( - input_: torch.Tensor, - scatter_dim: int = 2, - gather_dim: int = 1, -): - return _AllToAll.apply(input_, scatter_dim, gather_dim, _all_to_all) - - -def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num): - def all_to_all(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask): - hidden_states = _single_all_to_all(hidden_states, scatter_dim=2, gather_dim=0, enable_HCCL=True) - encoder_hidden_states = _single_all_to_all(encoder_hidden_states, scatter_dim=1, gather_dim=0, enable_HCCL=True) - attention_mask = _single_all_to_all(attention_mask, scatter_dim=1, gather_dim=0, enable_HCCL=True) - encoder_attention_mask = _single_all_to_all(encoder_attention_mask, scatter_dim=1, gather_dim=0, enable_HCCL=True) - return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask - - sp_size = hccl_info.world_size - frame = hidden_states.shape[2] - assert frame % sp_size == 0, "frame should be a multiple of sp_size" - - encoder_hidden_states = rearrange(encoder_hidden_states, 'b 1 (n x) h -> b n x h', - n=sp_size, x=encoder_hidden_states.shape[2]//sp_size).contiguous() - hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask = all_to_all(hidden_states, - encoder_hidden_states, - attention_mask.repeat(1, sp_size, 1, 1), - encoder_attention_mask.repeat(1, sp_size, 1)) - - return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num \ No newline at end of file diff --git a/opensora/acceleration/parallel_states.py b/opensora/acceleration/parallel_states.py deleted file mode 100644 index e3f96e121..000000000 --- a/opensora/acceleration/parallel_states.py +++ /dev/null @@ -1,57 +0,0 @@ -import torch -import torch_npu -import torch.distributed as dist -import os -try: - from lcalib.functional import lcal_initialize - enable_LCCL = True -except: - lcal_initialize = None - enable_LCCL = False -class COMM_INFO: - def __init__(self): - self.group = None - self.world_size = 0 - self.rank = -1 - -lccl_info = COMM_INFO() -hccl_info = COMM_INFO() -_SEQUENCE_PARALLEL_STATE = False -def initialize_sequence_parallel_state(sequence_parallel_size): - global _SEQUENCE_PARALLEL_STATE - if sequence_parallel_size > 1: - _SEQUENCE_PARALLEL_STATE = True - initialize_sequence_parallel_group(sequence_parallel_size) - -def set_sequence_parallel_state(state): - global _SEQUENCE_PARALLEL_STATE - _SEQUENCE_PARALLEL_STATE = state - -def get_sequence_parallel_state(): - return _SEQUENCE_PARALLEL_STATE - -def initialize_sequence_parallel_group(sequence_parallel_size): - """Initialize the sequence parallel group.""" - rank = int(os.getenv('RANK', '0')) - world_size = int(os.getenv("WORLD_SIZE", '1')) - assert world_size % sequence_parallel_size == 0, "world_size must be divisible by sequence_parallel_size" - # hccl - hccl_info.world_size = sequence_parallel_size - hccl_info.rank = rank - num_sequence_parallel_groups: int = world_size // sequence_parallel_size - for i in range(num_sequence_parallel_groups): - ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size) - group = dist.new_group(ranks) - if rank in ranks: - hccl_info.group = group - - if enable_LCCL: - assert sequence_parallel_size == 8, "sequence_parallel_size should be 8 when enable_LCCL is True" - rank %= sequence_parallel_size - lccl_info.world_size = sequence_parallel_size - lccl_info.group = lcal_initialize(rank, sequence_parallel_size) - lccl_info.rank = rank - -def destroy_sequence_parallel_group(): - """Destroy the sequence parallel group.""" - dist.destroy_process_group() diff --git a/opensora/dataset/t2v_datasets.py b/opensora/dataset/t2v_datasets.py index 4ed8444a2..c67c83e80 100644 --- a/opensora/dataset/t2v_datasets.py +++ b/opensora/dataset/t2v_datasets.py @@ -1,12 +1,6 @@ import time import traceback -try: - import torch_npu - from opensora.npu_config import npu_config -except: - torch_npu = None - npu_config = None import glob import json import os, io, csv, math, random @@ -181,11 +175,6 @@ class T2V_dataset(Dataset): return self.get_image(idx) def get_video(self, idx): - # npu_config.print_msg(f"current idx is {idx}") - # video = random.choice([random_video_noise(65, 3, 336, 448), random_video_noise(65, 3, 1024, 1024), random_video_noise(65, 3, 360, 480)]) - # # print('random shape', video.shape) - # input_ids = torch.ones(1, 120).to(torch.long).squeeze(0) - # cond_mask = torch.cat([torch.ones(1, 60).to(torch.long), torch.ones(1, 60).to(torch.long)], dim=1).squeeze(0) video_path = dataset_prog.cap_list[idx]['path'] assert os.path.exists(video_path), f"file {video_path} do not exist!" @@ -405,17 +394,6 @@ class T2V_dataset(Dataset): cap_lists += sub_list return cap_lists - # def get_img_cap_list(self): - # use_image_num = self.use_image_num if self.use_image_num != 0 else 1 - # if npu_config is None: - # img_cap_lists = self.read_jsons(self.image_data, postfix=".jpg") - # img_cap_lists = [img_cap_lists[i: i + use_image_num] for i in range(0, len(img_cap_lists), use_image_num)] - # else: - # img_cap_lists = npu_config.try_load_pickle("img_cap_lists_all", - # lambda: self.read_jsons(self.image_data, postfix=".jpg")) - # img_cap_lists = [img_cap_lists[i: i + use_image_num] for i in range(0, len(img_cap_lists), use_image_num)] - # img_cap_lists = img_cap_lists[npu_config.get_local_rank()::npu_config.N_NPU_PER_NODE] - # return img_cap_lists[:-1] # drop last to avoid error length def get_cap_list(self): cap_lists = self.read_jsons(self.data) diff --git a/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py b/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py index a6f20cc2b..86c19c4ff 100755 --- a/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py +++ b/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py @@ -1,9 +1,4 @@ -try: - import torch_npu - from opensora.npu_config import npu_config -except: - torch_npu = None - npu_config = None + from ..modeling_videobase import VideoBaseAE from ..modules import Normalize from ..modules.ops import nonlinearity @@ -147,11 +142,7 @@ class Encoder(nn.Module): h = self.mid.block_1(h) h = self.mid.attn_1(h) h = self.mid.block_2(h) - - if npu_config is None: - h = self.norm_out(h) - else: - h = npu_config.run_group_norm(self.norm_out, h) + h = self.norm_out(h) h = nonlinearity(h) h = self.conv_out(h) return h @@ -264,16 +255,9 @@ class Decoder(nn.Module): h = self.up[i_level].upsample(h) if hasattr(self.up[i_level], "time_upsample"): h = self.up[i_level].time_upsample(h) - if npu_config is None: - h = self.norm_out(h) - else: - h = npu_config.run_group_norm(self.norm_out, h) + h = self.norm_out(h) h = nonlinearity(h) - if npu_config is None: - h = self.conv_out(h) - else: - h_dtype = h.dtype - h = npu_config.run_conv3d(self.conv_out, h, h_dtype) + h = self.conv_out(h) return h diff --git a/opensora/models/causalvideovae/model/modules/attention.py b/opensora/models/causalvideovae/model/modules/attention.py index 73c4c3793..c3b9bb19b 100755 --- a/opensora/models/causalvideovae/model/modules/attention.py +++ b/opensora/models/causalvideovae/model/modules/attention.py @@ -1,10 +1,4 @@ import torch.nn as nn -try: - import torch_npu - from opensora.npu_config import npu_config -except: - torch_npu = None - npu_config = None from .normalize import Normalize from .conv import CausalConv3d import torch @@ -97,10 +91,7 @@ class AttnBlock3DFix(nn.Module): def forward(self, x): h_ = x - if npu_config is None: - h_ = self.norm(h_) - else: - h_ = npu_config.run_group_norm(self.norm, h_) + h_ = self.norm(h_) q = self.q(h_) k = self.k(h_) v = self.v(h_) diff --git a/opensora/models/causalvideovae/model/modules/conv.py b/opensora/models/causalvideovae/model/modules/conv.py index 71d6a4705..a4c6b78c0 100755 --- a/opensora/models/causalvideovae/model/modules/conv.py +++ b/opensora/models/causalvideovae/model/modules/conv.py @@ -2,17 +2,9 @@ import torch.nn as nn from typing import Union, Tuple import torch.nn.functional as F import torch -from .block import Block from .ops import cast_tuple -from einops import rearrange from .ops import video_to_image from torch.utils.checkpoint import checkpoint -try: - import torch_npu - from opensora.npu_config import npu_config -except: - torch_npu = None - npu_config = None class Conv2d(nn.Conv2d): def __init__( @@ -99,20 +91,12 @@ class CausalConv3d(nn.Module): nn.init.constant_(self.conv.bias, 0) def forward(self, x): - if npu_config is not None and npu_config.on_npu: - x_dtype = x.dtype - first_frame_pad = x[:, :, :1, :, :].repeat( - (1, 1, self.time_kernel_size - 1, 1, 1) - ) # b c t h w - x = torch.concatenate((first_frame_pad, x), dim=2) # 3 + 16 - return npu_config.run_conv3d(self.conv, x, x_dtype) - else: - # 1 + 16 16 as video, 1 as image - first_frame_pad = x[:, :, :1, :, :].repeat( - (1, 1, self.time_kernel_size - 1, 1, 1) - ) # b c t h w - x = torch.concatenate((first_frame_pad, x), dim=2) # 3 + 16 - return self.conv(x) + # 1 + 16 16 as video, 1 as image + first_frame_pad = x[:, :, :1, :, :].repeat( + (1, 1, self.time_kernel_size - 1, 1, 1) + ) # b c t h w + x = torch.concatenate((first_frame_pad, x), dim=2) # 3 + 16 + return self.conv(x) class CausalConv3d_GC(CausalConv3d): diff --git a/opensora/models/causalvideovae/model/modules/resnet_block.py b/opensora/models/causalvideovae/model/modules/resnet_block.py index ce70334f1..d0e2551b5 100755 --- a/opensora/models/causalvideovae/model/modules/resnet_block.py +++ b/opensora/models/causalvideovae/model/modules/resnet_block.py @@ -6,12 +6,6 @@ from .ops import nonlinearity, video_to_image from .conv import CausalConv3d, CausalConv3d_GC from .block import Block from torch.utils.checkpoint import checkpoint -try: - import torch_npu - from opensora.npu_config import npu_config -except: - torch_npu = None - npu_config = None class ResnetBlock2D(Block): def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, @@ -78,16 +72,10 @@ class ResnetBlock3D(Block): def forward(self, x): h = x - if npu_config is None: - h = self.norm1(h) - else: - h = npu_config.run_group_norm(self.norm1, h) + h = self.norm1(h) h = nonlinearity(h) h = self.conv1(h) - if npu_config is None: - h = self.norm2(h) - else: - h = npu_config.run_group_norm(self.norm2, h) + h = self.norm2(h) h = nonlinearity(h) h = self.dropout(h) h = self.conv2(h) diff --git a/opensora/models/causalvideovae/model/modules/updownsample.py b/opensora/models/causalvideovae/model/modules/updownsample.py index 86a48cc38..b2592851e 100755 --- a/opensora/models/causalvideovae/model/modules/updownsample.py +++ b/opensora/models/causalvideovae/model/modules/updownsample.py @@ -9,12 +9,6 @@ from .ops import cast_tuple, video_to_image from .conv import CausalConv3d, CausalConv3d_GC from einops import rearrange from .block import Block -try: - import torch_npu - from opensora.npu_config import npu_config -except: - torch_npu = None - npu_config = None class Upsample(Block): @@ -58,22 +52,13 @@ class Downsample(Block): def forward(self, x): if self.with_conv: if self.undown: - if npu_config is not None and npu_config.on_npu: - x_dtype = x.dtype - x = x.to(npu_config.replaced_type) - x = npu_config.run_conv3d(self.conv, x, x_dtype) - else: - x = self.conv(x) + + x = self.conv(x) else: pad = (0, 1, 0, 1) - if npu_config is not None and npu_config.on_npu: - x_dtype = x.dtype - x = x.to(npu_config.replaced_type) - x = torch.nn.functional.pad(x, pad, mode="constant", value=0) - x = npu_config.run_conv3d(self.conv, x, x_dtype) - else: - x = torch.nn.functional.pad(x, pad, mode="constant", value=0) - x = self.conv(x) + + x = torch.nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) else: x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) return x @@ -179,26 +164,16 @@ class TimeDownsample2x(Block): ): super().__init__() self.kernel_size = kernel_size - if npu_config is not None and npu_config.on_npu: - self.avg_pool = nn.AvgPool2d((kernel_size, 1), stride=(2, 1)) - self.pad = nn.ReplicationPad3d((0, 0, 0, 0, self.kernel_size - 1, 0)) - else: - self.conv = nn.AvgPool3d((kernel_size, 1, 1), stride=(2, 1, 1)) + + self.conv = nn.AvgPool3d((kernel_size, 1, 1), stride=(2, 1, 1)) def forward(self, x): - if npu_config is not None and npu_config.on_npu: - n, c, d, h, w = x.shape - x = self.pad(x) - x = x.view(n * c, -1, h * w) - pooled = self.avg_pool(x) - output = pooled.view(n, c, -1, h, w) - return output - else: - first_frame_pad = x[:, :, :1, :, :].repeat( - (1, 1, self.kernel_size - 1, 1, 1) - ) - x = torch.concatenate((first_frame_pad, x), dim=2) - return self.conv(x) + + first_frame_pad = x[:, :, :1, :, :].repeat( + (1, 1, self.kernel_size - 1, 1, 1) + ) + x = torch.concatenate((first_frame_pad, x), dim=2) + return self.conv(x) class TimeUpsample2x(Block): def __init__( @@ -224,11 +199,8 @@ class TimeDownsampleRes2x(nn.Module): ): super().__init__() self.kernel_size = cast_tuple(kernel_size, 3) - if npu_config is not None and npu_config.on_npu: - self.avg_pool = nn.AvgPool2d((kernel_size, 1), stride=(2, 1)) - self.pad = nn.ReplicationPad3d((0, 0, 0, 0, kernel_size - 1, 0)) - else: - self.avg_pool = nn.AvgPool3d((kernel_size, 1, 1), stride=(2, 1, 1)) + + self.avg_pool = nn.AvgPool3d((kernel_size, 1, 1), stride=(2, 1, 1)) self.conv = nn.Conv3d( in_channels, out_channels, self.kernel_size, stride=(2,1,1), padding=(0,1,1) ) @@ -236,21 +208,12 @@ class TimeDownsampleRes2x(nn.Module): def forward(self, x): alpha = torch.sigmoid(self.mix_factor) - if npu_config is not None and npu_config.on_npu: - n, c, d, h, w = x.shape - x_dtype = x.dtype - x = x.to(npu_config.replaced_type) - x = self.pad(x) - pad_x = x.view(n, c, -1, h, w) - avg_x = self.avg_pool(x.view(n * c, -1, h * w)).view(n, c, -1, h, w).to(x_dtype) - conv_x = npu_config.run_conv3d(self.conv, pad_x, x_dtype) - return alpha * avg_x + (1 - alpha) * conv_x - else: - first_frame_pad = x[:, :, :1, :, :].repeat( - (1, 1, self.kernel_size[0] - 1, 1, 1) - ) - x = torch.concatenate((first_frame_pad, x), dim=2) - return alpha * self.avg_pool(x) + (1 - alpha) * self.conv(x) + + first_frame_pad = x[:, :, :1, :, :].repeat( + (1, 1, self.kernel_size[0] - 1, 1, 1) + ) + x = torch.concatenate((first_frame_pad, x), dim=2) + return alpha * self.avg_pool(x) + (1 - alpha) * self.conv(x) class TimeUpsampleRes2x(nn.Module): def __init__( @@ -270,13 +233,8 @@ class TimeUpsampleRes2x(nn.Module): alpha = torch.sigmoid(self.mix_factor) if x.size(2) > 1: x,x_= x[:,:,:1],x[:,:,1:] - if npu_config is not None and npu_config.on_npu: - x_dtype = x_.dtype - x_ = x_.to(npu_config.replaced_type) - x_ = F.interpolate(x_, scale_factor=(2, 1, 1), mode='trilinear') - x_ = x_.to(x_dtype) - else: - x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear') + + x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear') x = torch.concat([x, x_], dim=2) return alpha * x + (1-alpha) * self.conv(x) @@ -326,13 +284,8 @@ class TimeUpsampleResAdv2x(nn.Module): def forward(self, x): if x.size(2) > 1: x,x_= x[:,:,:1],x[:,:,1:] - if npu_config is not None and npu_config.on_npu: - x_dtype = x_.dtype - x_ = x_.to(npu_config.replaced_type) - x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear') - x_ = x_.to(x_dtype) - else: - x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear') + + x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear') x = torch.concat([x, x_], dim=2) alpha = torch.sigmoid(self.mix_factor) return alpha * x + (1 - alpha) * self.conv(self.attn(self.res(x))) diff --git a/opensora/models/diffusion/opensora/modeling_inpaint.py b/opensora/models/diffusion/opensora/modeling_inpaint.py index 4e2c17473..61e50edb3 100644 --- a/opensora/models/diffusion/opensora/modeling_inpaint.py +++ b/opensora/models/diffusion/opensora/modeling_inpaint.py @@ -14,15 +14,6 @@ from diffusers.models.embeddings import PixArtAlphaTextProjection from opensora.models.diffusion.opensora.modules import OverlapPatchEmbed3D, OverlapPatchEmbed2D, PatchEmbed2D, BasicTransformerBlock from opensora.utils.utils import to_2tuple -try: - import torch_npu - from opensora.npu_config import npu_config - from opensora.acceleration.parallel_states import get_sequence_parallel_state, hccl_info -except: - torch_npu = None - npu_config = None - from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info - from PIL import Image import numpy as np from enum import Enum, auto diff --git a/opensora/models/diffusion/opensora/modeling_opensora.py b/opensora/models/diffusion/opensora/modeling_opensora.py index adfcce5fe..0fab8db16 100644 --- a/opensora/models/diffusion/opensora/modeling_opensora.py +++ b/opensora/models/diffusion/opensora/modeling_opensora.py @@ -13,14 +13,8 @@ from diffusers.models.normalization import AdaLayerNormSingle from diffusers.models.embeddings import PixArtAlphaTextProjection from opensora.models.diffusion.opensora.modules import OverlapPatchEmbed3D, OverlapPatchEmbed2D, PatchEmbed2D, BasicTransformerBlock from opensora.utils.utils import to_2tuple -try: - import torch_npu - from opensora.npu_config import npu_config - from opensora.acceleration.parallel_states import get_sequence_parallel_state, hccl_info -except: - torch_npu = None - npu_config = None - from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info + +from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info class OpenSoraT2V(ModelMixin, ConfigMixin): """ @@ -353,14 +347,11 @@ class OpenSoraT2V(ModelMixin, ConfigMixin): # b, 1, h, w -> only images attention_mask = attention_mask.to(self.dtype) if get_sequence_parallel_state(): - if npu_config is not None: - attention_mask_vid = attention_mask[:, :frame * hccl_info.world_size] # b, frame, h, w - attention_mask_img = attention_mask[:, frame * hccl_info.world_size:] # b, use_image_num, h, w - else: - # print('before attention_mask.shape', attention_mask.shape) - attention_mask_vid = attention_mask[:, :frame * nccl_info.world_size] # b, frame, h, w - attention_mask_img = attention_mask[:, frame * nccl_info.world_size:] # b, use_image_num, h, w - # print('after attention_mask.shape', attention_mask_vid.shape) + + # print('before attention_mask.shape', attention_mask.shape) + attention_mask_vid = attention_mask[:, :frame * nccl_info.world_size] # b, frame, h, w + attention_mask_img = attention_mask[:, frame * nccl_info.world_size:] # b, use_image_num, h, w + # print('after attention_mask.shape', attention_mask_vid.shape) else: attention_mask_vid = attention_mask[:, :frame] # b, frame, h, w attention_mask_img = attention_mask[:, frame:] # b, use_image_num, h, w @@ -399,14 +390,7 @@ class OpenSoraT2V(ModelMixin, ConfigMixin): encoder_attention_mask_img = encoder_attention_mask_vid encoder_attention_mask_vid = None - if npu_config is not None and attention_mask_vid is not None: - attention_mask_vid = npu_config.get_attention_mask(attention_mask_vid, attention_mask_vid.shape[-1]) - encoder_attention_mask_vid = npu_config.get_attention_mask(encoder_attention_mask_vid, - attention_mask_vid.shape[-2]) - if npu_config is not None and attention_mask_img is not None: - attention_mask_img = npu_config.get_attention_mask(attention_mask_img, attention_mask_img.shape[-1]) - encoder_attention_mask_img = npu_config.get_attention_mask(encoder_attention_mask_img, - attention_mask_img.shape[-2]) + # 1. Input diff --git a/opensora/models/diffusion/opensora/modules.py b/opensora/models/diffusion/opensora/modules.py index 8904bc0c5..92cea6623 100644 --- a/opensora/models/diffusion/opensora/modules.py +++ b/opensora/models/diffusion/opensora/modules.py @@ -19,17 +19,8 @@ from diffusers.models.attention_processor import Attention as Attention_ from diffusers.models.embeddings import SinusoidalPositionalEmbedding from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm from .rope import PositionGetter3D, RoPE3D -try: - import torch_npu - from opensora.npu_config import npu_config, set_run_dtype - from opensora.acceleration.parallel_states import get_sequence_parallel_state, hccl_info - from opensora.acceleration.communications import all_to_all_SBH -except: - torch_npu = None - npu_config = None - set_run_dtype = None - from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info - from opensora.utils.communications import all_to_all_SBH +from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info +from opensora.utils.communications import all_to_all_SBH logger = logging.get_logger(__name__) def get_3d_sincos_pos_embed( @@ -227,30 +218,18 @@ class PatchEmbed2D(nn.Module): # import ipdb;ipdb.set_trace() # raise NotImplementedError if get_sequence_parallel_state(): - if npu_config is not None: - sp_size = hccl_info.world_size - temp_pos_embed = get_1d_sincos_pos_embed( - embed_dim=self.temp_pos_embed.shape[-1], - grid_size=num_frames * sp_size, - base_size=self.base_size_t, - interpolation_scale=self.interpolation_scale_t, - ) - rank = hccl_info.rank % sp_size - st_frame = rank * num_frames - ed_frame = st_frame + num_frames - temp_pos_embed = temp_pos_embed[st_frame: ed_frame] - else: - sp_size = nccl_info.world_size - temp_pos_embed = get_1d_sincos_pos_embed( - embed_dim=self.temp_pos_embed.shape[-1], - grid_size=num_frames * sp_size, - base_size=self.base_size_t, - interpolation_scale=self.interpolation_scale_t, - ) - rank = nccl_info.rank % sp_size - st_frame = rank * num_frames - ed_frame = st_frame + num_frames - temp_pos_embed = temp_pos_embed[st_frame: ed_frame] + + sp_size = nccl_info.world_size + temp_pos_embed = get_1d_sincos_pos_embed( + embed_dim=self.temp_pos_embed.shape[-1], + grid_size=num_frames * sp_size, + base_size=self.base_size_t, + interpolation_scale=self.interpolation_scale_t, + ) + rank = nccl_info.rank % sp_size + st_frame = rank * num_frames + ed_frame = st_frame + num_frames + temp_pos_embed = temp_pos_embed[st_frame: ed_frame] else: temp_pos_embed = get_1d_sincos_pos_embed( @@ -609,12 +588,7 @@ class DownSampler3d(nn.Module): def forward(self, x, attention_mask, t, h, w): b = x.shape[0] x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w) - if npu_config is None: - x = self.layer(x) + (x if self.down_shortcut else 0) - else: - x_dtype = x.dtype - x = npu_config.run_conv3d(self.layer, x, x_dtype) + (x if self.down_shortcut else 0) - + x = self.layer(x) + (x if self.down_shortcut else 0) self.t = t//self.down_factor[0] self.h = h//self.down_factor[1] self.w = w//self.down_factor[2] @@ -722,30 +696,17 @@ class AttnProcessor2_0: hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) if get_sequence_parallel_state(): - if npu_config is not None: - sequence_length, batch_size, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - else: - sequence_length, batch_size, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) + + sequence_length, batch_size, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) else: batch_size, sequence_length, _ = ( hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape ) if attention_mask is not None: - if npu_config is None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length * nccl_info.world_size, batch_size) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - if get_sequence_parallel_state(): - attention_mask = attention_mask.view(batch_size, attn.heads // nccl_info.world_size, -1, attention_mask.shape[-1]) - else: - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - else: - attention_mask = attention_mask.view(batch_size, 1, -1, attention_mask.shape[-1]) + attention_mask = attention_mask.view(batch_size, 1, -1, attention_mask.shape[-1]) if attn.group_norm is not None: hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) @@ -762,165 +723,109 @@ class AttnProcessor2_0: inner_dim = key.shape[-1] head_dim = inner_dim // attn.heads - if npu_config is not None and npu_config.on_npu: - if get_sequence_parallel_state(): - query = query.view(-1, attn.heads, head_dim) # [s // sp, b, h * d] -> [s // sp * b, h, d] - key = key.view(-1, attn.heads, head_dim) - value = value.view(-1, attn.heads, head_dim) - # query = attn.q_norm(query) - # key = attn.k_norm(key) - h_size = attn.heads * head_dim - sp_size = hccl_info.world_size - h_size_sp = h_size // sp_size - # apply all_to_all to gather sequence and split attention heads [s // sp * b, h, d] -> [s * b, h // sp, d] - query = all_to_all_SBH(query, scatter_dim=1, gather_dim=0).view(-1, batch_size, h_size_sp) - key = all_to_all_SBH(key, scatter_dim=1, gather_dim=0).view(-1, batch_size, h_size_sp) - value = all_to_all_SBH(value, scatter_dim=1, gather_dim=0).view(-1, batch_size, h_size_sp) - if self.use_rope: - query = query.view(-1, batch_size, attn.heads // sp_size, head_dim) - key = key.view(-1, batch_size, attn.heads // sp_size, head_dim) - # require the shape of (batch_size x nheads x ntokens x dim) - pos_thw = self.position_getter(batch_size, t=frame * sp_size, h=height, w=width, device=query.device) - query = self.rope(query, pos_thw) - key = self.rope(key, pos_thw) - query = query.view(-1, batch_size, h_size_sp) - key = key.view(-1, batch_size, h_size_sp) - value = value.view(-1, batch_size, h_size_sp) - hidden_states = npu_config.run_attention(query, key, value, attention_mask, "SBH", - head_dim, attn.heads // sp_size) + if get_sequence_parallel_state(): + query = query.reshape(-1, attn.heads, head_dim) # [s // sp, b, h * d] -> [s // sp * b, h, d] + key = key.reshape(-1, attn.heads, head_dim) + value = value.reshape(-1, attn.heads, head_dim) + # query = attn.q_norm(query) + # key = attn.k_norm(key) + h_size = attn.heads * head_dim + sp_size = nccl_info.world_size + h_size_sp = h_size // sp_size + # apply all_to_all to gather sequence and split attention heads [s // sp * b, h, d] -> [s * b, h // sp, d] + query = all_to_all_SBH(query, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp) + key = all_to_all_SBH(key, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp) + value = all_to_all_SBH(value, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp) + query = query.reshape(-1, batch_size, attn.heads // sp_size, head_dim) + key = key.reshape(-1, batch_size, attn.heads // sp_size, head_dim) + value = value.reshape(-1, batch_size, attn.heads // sp_size, head_dim) + # print('query', query.shape, 'key', key.shape, 'value', value.shape) + if self.use_rope: + # require the shape of (batch_size x nheads x ntokens x dim) + pos_thw = self.position_getter(batch_size, t=frame * sp_size, h=height, w=width, device=query.device) + query = self.rope(query, pos_thw) + key = self.rope(key, pos_thw) - hidden_states = hidden_states.view(-1, attn.heads // sp_size, head_dim) + # print('after rope query', query.shape, 'key', key.shape, 'value', value.shape) + query = rearrange(query, 's b h d -> b h s d') + key = rearrange(key, 's b h d -> b h s d') + value = rearrange(value, 's b h d -> b h s d') + # print('rearrange query', query.shape, 'key', key.shape, 'value', value.shape) - # [s * b, h // sp, d] -> [s // sp * b, h, d] -> [s // sp, b, h * d] - hidden_states = all_to_all_SBH(hidden_states, scatter_dim=0, gather_dim=1).view(-1, batch_size, h_size) + # 0, -10000 ->(bool) False, True ->(any) True ->(not) False + # 0, 0 ->(bool) False, False ->(any) False ->(not) True + if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible + attention_mask = None + # the output of sdp = (batch, num_heads, seq_len, head_dim) + # TODO: add support for attn.scale when we move to Torch 2.1 + # import ipdb;ipdb.set_trace() + # print(attention_mask) + if self.attention_mode == 'flash': + assert attention_mask is None, 'flash-attn do not support attention_mask' + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + hidden_states = F.scaled_dot_product_attention( + query, key, value, dropout_p=0.0, is_causal=False + ) + elif self.attention_mode == 'xformers': + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + elif self.attention_mode == 'math': + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) else: - if npu_config.enable_FA and query.dtype == torch.float32: - dtype = torch.bfloat16 - else: - dtype = None + raise NotImplementedError(f'Found attention_mode: {self.attention_mode}') + + hidden_states = rearrange(hidden_states, 'b h s d -> s b h d') - query = query.view(batch_size, -1, attn.heads, head_dim) - key = key.view(batch_size, -1, attn.heads, head_dim) - # query = attn.q_norm(query) - # key = attn.k_norm(key) - if self.use_rope: - # require the shape of (batch_size x nheads x ntokens x dim) - pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device) - query = self.rope(query, pos_thw) - key = self.rope(key, pos_thw) - query = query.view(batch_size, -1, attn.heads * head_dim) - key = key.view(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.reshape(-1, attn.heads // sp_size, head_dim) - with set_run_dtype(query, dtype): - query, key, value = npu_config.set_current_run_dtype([query, key, value]) - hidden_states = npu_config.run_attention(query, key, value, attention_mask, "BSH", - head_dim, attn.heads) - - hidden_states = npu_config.restore_dtype(hidden_states) + # [s * b, h // sp, d] -> [s // sp * b, h, d] -> [s // sp, b, h * d] + hidden_states = all_to_all_SBH(hidden_states, scatter_dim=0, gather_dim=1).reshape(-1, batch_size, h_size) else: - if get_sequence_parallel_state(): - query = query.reshape(-1, attn.heads, head_dim) # [s // sp, b, h * d] -> [s // sp * b, h, d] - key = key.reshape(-1, attn.heads, head_dim) - value = value.reshape(-1, attn.heads, head_dim) - # query = attn.q_norm(query) - # key = attn.k_norm(key) - h_size = attn.heads * head_dim - sp_size = nccl_info.world_size - h_size_sp = h_size // sp_size - # apply all_to_all to gather sequence and split attention heads [s // sp * b, h, d] -> [s * b, h // sp, d] - query = all_to_all_SBH(query, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp) - key = all_to_all_SBH(key, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp) - value = all_to_all_SBH(value, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp) - query = query.reshape(-1, batch_size, attn.heads // sp_size, head_dim) - key = key.reshape(-1, batch_size, attn.heads // sp_size, head_dim) - value = value.reshape(-1, batch_size, attn.heads // sp_size, head_dim) - # print('query', query.shape, 'key', key.shape, 'value', value.shape) - if self.use_rope: - # require the shape of (batch_size x nheads x ntokens x dim) - pos_thw = self.position_getter(batch_size, t=frame * sp_size, h=height, w=width, device=query.device) - query = self.rope(query, pos_thw) - key = self.rope(key, pos_thw) + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + # query = attn.q_norm(query) + # key = attn.k_norm(key) - # print('after rope query', query.shape, 'key', key.shape, 'value', value.shape) - query = rearrange(query, 's b h d -> b h s d') - key = rearrange(key, 's b h d -> b h s d') - value = rearrange(value, 's b h d -> b h s d') - # print('rearrange query', query.shape, 'key', key.shape, 'value', value.shape) + if self.use_rope: + # require the shape of (batch_size x nheads x ntokens x dim) + pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device) + query = self.rope(query, pos_thw) + key = self.rope(key, pos_thw) + + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - # 0, -10000 ->(bool) False, True ->(any) True ->(not) False - # 0, 0 ->(bool) False, False ->(any) False ->(not) True - if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible - attention_mask = None - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - # import ipdb;ipdb.set_trace() - # print(attention_mask) - if self.attention_mode == 'flash': - assert attention_mask is None, 'flash-attn do not support attention_mask' - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): - hidden_states = F.scaled_dot_product_attention( - query, key, value, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'xformers': - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'math': + # 0, -10000 ->(bool) False, True ->(any) True ->(not) False + # 0, 0 ->(bool) False, False ->(any) False ->(not) True + if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible + attention_mask = None + # the output of sdp = (batch, num_heads, seq_len, head_dim) + # TODO: add support for attn.scale when we move to Torch 2.1 + # import ipdb;ipdb.set_trace() + # print(attention_mask) + if self.attention_mode == 'flash': + assert attention_mask is None, 'flash-attn do not support attention_mask' + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + hidden_states = F.scaled_dot_product_attention( + query, key, value, dropout_p=0.0, is_causal=False + ) + elif self.attention_mode == 'xformers': + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): hidden_states = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) - else: - raise NotImplementedError(f'Found attention_mode: {self.attention_mode}') - - hidden_states = rearrange(hidden_states, 'b h s d -> s b h d') - - hidden_states = hidden_states.reshape(-1, attn.heads // sp_size, head_dim) - - # [s * b, h // sp, d] -> [s // sp * b, h, d] -> [s // sp, b, h * d] - hidden_states = all_to_all_SBH(hidden_states, scatter_dim=0, gather_dim=1).reshape(-1, batch_size, h_size) + elif self.attention_mode == 'math': + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) else: - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # qk norm - # query = attn.q_norm(query) - # key = attn.k_norm(key) - - if self.use_rope: - # require the shape of (batch_size x nheads x ntokens x dim) - pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device) - query = self.rope(query, pos_thw) - key = self.rope(key, pos_thw) - - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # 0, -10000 ->(bool) False, True ->(any) True ->(not) False - # 0, 0 ->(bool) False, False ->(any) False ->(not) True - if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible - attention_mask = None - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - # import ipdb;ipdb.set_trace() - # print(attention_mask) - if self.attention_mode == 'flash': - assert attention_mask is None, 'flash-attn do not support attention_mask' - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): - hidden_states = F.scaled_dot_product_attention( - query, key, value, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'xformers': - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'math': - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - else: - raise NotImplementedError(f'Found attention_mode: {self.attention_mode}') - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + raise NotImplementedError(f'Found attention_mode: {self.attention_mode}') + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states = hidden_states.to(query.dtype) # linear proj @@ -964,25 +869,14 @@ class FeedForward_Conv3d(nn.Module): def forward(self, x, t, h, w): # import ipdb;ipdb.set_trace() - if npu_config is None: - x = self.project_in(x) - x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w) - x = F.gelu(x) - out = x - for module in self.dwconv: - out = out + module(x) - out = rearrange(out, 'b d t h w -> b (t h w) d', t=t, h=h, w=w) - x = self.project_out(out) - else: - x_dtype = x.dtype - x = npu_config.run_conv3d(self.project_in, x, npu_config.replaced_type) - x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w) - x = F.gelu(x) - out = x - for module in self.dwconv: - out = out + npu_config.run_conv3d(module, x, npu_config.replaced_type) - out = rearrange(out, 'b d t h w -> b (t h w) d', t=t, h=h, w=w) - x = npu_config.run_conv3d(self.project_out, out, x_dtype) + x = self.project_in(x) + x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w) + x = F.gelu(x) + out = x + for module in self.dwconv: + out = out + module(x) + out = rearrange(out, 'b d t h w -> b (t h w) d', t=t, h=h, w=w) + x = self.project_out(out) return x diff --git a/opensora/models/diffusion/opensora/rope.py b/opensora/models/diffusion/opensora/rope.py index 8ead2b4b4..f43dcae04 100644 --- a/opensora/models/diffusion/opensora/rope.py +++ b/opensora/models/diffusion/opensora/rope.py @@ -1,12 +1,6 @@ import torch -try: - import torch_npu - from opensora.npu_config import npu_config, set_run_dtype - from opensora.acceleration.parallel_states import get_sequence_parallel_state -except: - torch_npu = None - npu_config = None - from opensora.utils.parallel_states import get_sequence_parallel_state + +from opensora.utils.parallel_states import get_sequence_parallel_state class PositionGetter3D(object): """ return positions of patches """ @@ -63,12 +57,11 @@ class RoPE3D(torch.nn.Module): def apply_rope1d(self, tokens, pos1d, cos, sin): assert pos1d.ndim == 2 - if npu_config is None and not get_sequence_parallel_state(): + if not get_sequence_parallel_state(): # for (batch_size x nheads x ntokens x dim) cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :] sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :] - else: - # for (batch_size x ntokens x nheads x dim) + else: cos = torch.nn.functional.embedding(pos1d, cos)[:, :, None, :] sin = torch.nn.functional.embedding(pos1d, sin)[:, :, None, :] diff --git a/opensora/models/diffusion/udit/modeling_udit.py b/opensora/models/diffusion/udit/modeling_udit.py index 8a22e618a..777a40749 100644 --- a/opensora/models/diffusion/udit/modeling_udit.py +++ b/opensora/models/diffusion/udit/modeling_udit.py @@ -15,15 +15,6 @@ from diffusers.models.embeddings import PixArtAlphaTextProjection from opensora.models.diffusion.udit.modules import Upsample2d, Downsample2d, PatchEmbed2D, BasicTransformerBlock, \ FP32_GELU, FP32_SiLU, FP32_Layernorm from opensora.utils.utils import to_2tuple -import math -import re -try: - import torch_npu - from opensora.npu_config import npu_config -except: - torch_npu = None - npu_config = None - class UDiTT2V(ModelMixin, ConfigMixin, PeftAdapterMixin): """ A 2D Transformer model for image-like data. diff --git a/opensora/models/diffusion/udit/modules.py b/opensora/models/diffusion/udit/modules.py index 631848900..1c0f9de29 100644 --- a/opensora/models/diffusion/udit/modules.py +++ b/opensora/models/diffusion/udit/modules.py @@ -20,13 +20,6 @@ from diffusers.models.embeddings import SinusoidalPositionalEmbedding from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm from .rope import PositionGetter3D, RoPE3D import re -try: - import torch_npu - from opensora.npu_config import npu_config, set_run_dtype -except: - torch_npu = None - npu_config = None - set_run_dtype = None logger = logging.get_logger(__name__) def get_3d_sincos_pos_embed( @@ -354,30 +347,18 @@ class PatchEmbed2D(nn.Module): # import ipdb;ipdb.set_trace() # raise NotImplementedError if get_sequence_parallel_state(): - if npu_config is not None: - sp_size = hccl_info.world_size - temp_pos_embed = get_1d_sincos_pos_embed( - embed_dim=self.temp_pos_embed.shape[-1], - grid_size=num_frames * sp_size, - base_size=self.base_size_t, - interpolation_scale=self.interpolation_scale_t, - ) - rank = hccl_info.rank % sp_size - st_frame = rank * num_frames - ed_frame = st_frame + num_frames - temp_pos_embed = temp_pos_embed[st_frame: ed_frame] - else: - sp_size = nccl_info.world_size - temp_pos_embed = get_1d_sincos_pos_embed( - embed_dim=self.temp_pos_embed.shape[-1], - grid_size=num_frames * sp_size, - base_size=self.base_size_t, - interpolation_scale=self.interpolation_scale_t, - ) - rank = nccl_info.rank % sp_size - st_frame = rank * num_frames - ed_frame = st_frame + num_frames - temp_pos_embed = temp_pos_embed[st_frame: ed_frame] + + sp_size = nccl_info.world_size + temp_pos_embed = get_1d_sincos_pos_embed( + embed_dim=self.temp_pos_embed.shape[-1], + grid_size=num_frames * sp_size, + base_size=self.base_size_t, + interpolation_scale=self.interpolation_scale_t, + ) + rank = nccl_info.rank % sp_size + st_frame = rank * num_frames + ed_frame = st_frame + num_frames + temp_pos_embed = temp_pos_embed[st_frame: ed_frame] else: temp_pos_embed = get_1d_sincos_pos_embed( @@ -442,11 +423,7 @@ class DownSampler3d(nn.Module): def forward(self, x, attention_mask, t, h, w): b = x.shape[0] x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w) - if npu_config is None: - x = self.layer(x) + (x if self.down_shortcut else 0) - else: - x_dtype = x.dtype - x = npu_config.run_conv3d(self.layer, x, x_dtype) + (x if self.down_shortcut else 0) + x = self.layer(x) + (x if self.down_shortcut else 0) self.t = t//self.down_factor[0] self.h = h//self.down_factor[1] @@ -561,16 +538,10 @@ class AttnProcessor2_0: ) if attention_mask is not None: - if npu_config is None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - else: - if npu_config.enable_FA: - attention_mask = attention_mask.to(torch.bool) - attention_mask = attention_mask.view(batch_size, 1, -1, attention_mask.shape[-1]) - attention_mask = attention_mask.repeat(1, 1, hidden_states.shape[1], 1) + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) if attn.group_norm is not None: hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) @@ -588,69 +559,46 @@ class AttnProcessor2_0: inner_dim = key.shape[-1] head_dim = inner_dim // attn.heads - if npu_config is not None and npu_config.on_npu: - if npu_config.enable_FA and query.dtype == torch.float32: - dtype = torch.bfloat16 - else: - dtype = None + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - if self.use_rope: - query = query.view(batch_size, -1, attn.heads, head_dim) - key = key.view(batch_size, -1, attn.heads, head_dim) - # require the shape of (batch_size x nheads x ntokens x dim) - pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device) - query = self.rope(query, pos_thw) - key = self.rope(key, pos_thw) - query = query.view(batch_size, -1, attn.heads * head_dim) - key = key.view(batch_size, -1, attn.heads * head_dim) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - with set_run_dtype(query, dtype): - query, key, value = npu_config.set_current_run_dtype([query, key, value]) - hidden_states = npu_config.run_attention(query, key, value, attention_mask, "BSH", - head_dim, attn.heads) + # qk norm + # query = attn.q_norm(query) + # key = attn.k_norm(key) - hidden_states = npu_config.restore_dtype(hidden_states) - else: - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + if self.use_rope: + # require the shape of (batch_size x nheads x ntokens x dim) + pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device) + query = self.rope(query, pos_thw) + key = self.rope(key, pos_thw) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # qk norm - # query = attn.q_norm(query) - # key = attn.k_norm(key) - - if self.use_rope: - # require the shape of (batch_size x nheads x ntokens x dim) - pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device) - query = self.rope(query, pos_thw) - key = self.rope(key, pos_thw) - - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - # 0, -10000 ->(bool) False, True ->(any) True ->(not) False - # 0, 0 ->(bool) False, False ->(any) False ->(not) True - if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible - attention_mask = None - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - - if self.attention_mode == 'flash': - assert attention_mask is None, 'flash-attn do not support attention_mask' - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): - hidden_states = F.scaled_dot_product_attention( - query, key, value, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'xformers': - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'math': + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + # 0, -10000 ->(bool) False, True ->(any) True ->(not) False + # 0, 0 ->(bool) False, False ->(any) False ->(not) True + if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible + attention_mask = None + # the output of sdp = (batch, num_heads, seq_len, head_dim) + # TODO: add support for attn.scale when we move to Torch 2.1 + + if self.attention_mode == 'flash': + assert attention_mask is None, 'flash-attn do not support attention_mask' + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + hidden_states = F.scaled_dot_product_attention( + query, key, value, dropout_p=0.0, is_causal=False + ) + elif self.attention_mode == 'xformers': + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): hidden_states = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) - else: - raise NotImplementedError(f'Found attention_mode: {self.attention_mode}') - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + elif self.attention_mode == 'math': + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + else: + raise NotImplementedError(f'Found attention_mode: {self.attention_mode}') + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states = hidden_states.to(query.dtype) # linear proj @@ -720,11 +668,7 @@ class Downsample3d(nn.Module): x = rearrange(x, 'b (t h w) d -> b d t h w', t=frames, h=height, w=width) # x = F.pad(x, (0, pad_w, 0, pad_h, 0, 0), mode='reflect') x = F.pad(x, (0, pad_w, 0, pad_h, 0, 0)) - if npu_config is None: - x = self.body(x) - else: - x_dtype = x.dtype - x = npu_config.run_conv3d(self.body, x, x_dtype) + x = self.body(x) x = rearrange(x, 'b d t h w -> b (t h w) d') attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> b 1 t h w', t=frames, h=height, w=width) @@ -744,11 +688,7 @@ class Upsample3d(nn.Module): def forward(self, x, attention_mask, frames, height, width, pad_h=0, pad_w=0): x = rearrange(x, 'b (t h w) d -> b d t h w', t=frames, h=height, w=width) - if npu_config is None: - x = self.body(x) - else: - x_dtype = x.dtype - x = npu_config.run_conv3d(self.body, x, x_dtype) + x = self.body(x) x = x[:, :, :, :height*2-pad_h, :width*2-pad_w] x = rearrange(x, 'b d t h w -> b (t h w) d') @@ -769,10 +709,7 @@ class Downsample2d(nn.Module): def forward(self, x, attention_mask, frames, height, width, pad_h=0, pad_w=0): # import ipdb;ipdb.set_trace() x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=frames, h=height, w=width) - if npu_config is None: - x = F.pad(x, (0, pad_w, 0, pad_h)) - else: - x = npu_config.run_pad_2d(F.pad, x, pad=(0, pad_w, 0, pad_h)) + x = F.pad(x, (0, pad_w, 0, pad_h)) x = self.body(x) x = rearrange(x, '(b t) d h w -> b (t h w) d', t=frames) diff --git a/opensora/sample/pipeline_opensora.py b/opensora/sample/pipeline_opensora.py new file mode 100644 index 000000000..a88370473 --- /dev/null +++ b/opensora/sample/pipeline_opensora.py @@ -0,0 +1,817 @@ +# Copyright 2024 PixArt-Sigma Authors and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import html +import inspect +import re +import urllib.parse as ul +from typing import Callable, List, Optional, Tuple, Union +import math +import torch +from transformers import T5EncoderModel, T5Tokenizer + +from diffusers.models import AutoencoderKL, Transformer2DModel +from diffusers.schedulers import DPMSolverMultistepScheduler +from diffusers.utils import ( + BACKENDS_MAPPING, + deprecate, + is_bs4_available, + is_ftfy_available, + logging, + replace_example_docstring, +) +from diffusers.utils.torch_utils import randn_tensor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput + + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +if is_bs4_available(): + from bs4 import BeautifulSoup + +if is_ftfy_available(): + import ftfy + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import PixArtSigmaPipeline + + >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too. + >>> pipe = PixArtSigmaPipeline.from_pretrained( + ... "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16 + ... ) + >>> # Enable memory optimizations. + >>> # pipe.enable_model_cpu_offload() + + >>> prompt = "A small cactus with a happy face in the Sahara desert." + >>> image = pipe(prompt).images[0] + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + **kwargs, +): + """ + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default + timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps` + must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class OpenSoraPipeline(DiffusionPipeline): + r""" + Pipeline for text-to-image generation using PixArt-Sigma. + """ + + bad_punct_regex = re.compile( + r"[" + + "#®•©™&@·º½¾¿¡§~" + + r"\)" + + r"\(" + + r"\]" + + r"\[" + + r"\}" + + r"\{" + + r"\|" + + "\\" + + r"\/" + + r"\*" + + r"]{1,}" + ) # noqa + + _optional_components = ["tokenizer", "text_encoder"] + model_cpu_offload_seq = "text_encoder->transformer->vae" + + def __init__( + self, + tokenizer: T5Tokenizer, + text_encoder: T5EncoderModel, + vae: AutoencoderKL, + transformer: Transformer2DModel, + scheduler: DPMSolverMultistepScheduler, + ): + super().__init__() + + self.register_modules( + tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler + ) + + # self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + + # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt + def encode_prompt( + self, + prompt: Union[str, List[str]], + do_classifier_free_guidance: bool = True, + negative_prompt: str = "", + num_images_per_prompt: int = 1, + device: Optional[torch.device] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_attention_mask: Optional[torch.FloatTensor] = None, + negative_prompt_attention_mask: Optional[torch.FloatTensor] = None, + clean_caption: bool = False, + max_sequence_length: int = 120, + **kwargs, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `List[str]`, *optional*): + The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` + instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For + PixArt-Alpha, this should be "". + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + whether to use classifier free guidance or not + num_images_per_prompt (`int`, *optional*, defaults to 1): + number of images that should be generated per prompt + device: (`torch.device`, *optional*): + torch device to place the resulting embeddings on + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the "" + string. + clean_caption (`bool`, defaults to `False`): + If `True`, the function will preprocess and clean the provided caption before encoding. + max_sequence_length (`int`, defaults to 120): Maximum sequence length to use for the prompt. + """ + + if "mask_feature" in kwargs: + deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version." + deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False) + + if device is None: + device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda') + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # See Section 3.1. of the paper. + max_length = max_sequence_length + + if prompt_embeds is None: + prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( + text_input_ids, untruncated_ids + ): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {max_length} tokens: {removed_text}" + ) + + prompt_attention_mask = text_inputs.attention_mask + prompt_attention_mask = prompt_attention_mask.to(device) + + prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) + prompt_embeds = prompt_embeds[0] + + if self.text_encoder is not None: + dtype = self.text_encoder.dtype + elif self.transformer is not None: + dtype = self.transformer.dtype + else: + dtype = None + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) + prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance and negative_prompt_embeds is None: + uncond_tokens = [negative_prompt] * batch_size + uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) + max_length = prompt_embeds.shape[1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_attention_mask=True, + add_special_tokens=True, + return_tensors="pt", + ) + negative_prompt_attention_mask = uncond_input.attention_mask + negative_prompt_attention_mask = negative_prompt_attention_mask.to(device) + + negative_prompt_embeds = self.text_encoder( + uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask + ) + negative_prompt_embeds = negative_prompt_embeds[0] + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device) + + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1) + negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1) + else: + negative_prompt_embeds = None + negative_prompt_attention_mask = None + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs + def check_inputs( + self, + prompt, + num_frames, + height, + width, + negative_prompt, + callback_steps, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + if num_frames <= 0: + raise ValueError(f"`num_frames` have to be positive but is {num_frames}.") + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}." + ) + + # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing + def _text_preprocessing(self, text, clean_caption=False): + if clean_caption and not is_bs4_available(): + logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) + logger.warning("Setting `clean_caption` to False...") + clean_caption = False + + if clean_caption and not is_ftfy_available(): + logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) + logger.warning("Setting `clean_caption` to False...") + clean_caption = False + + if not isinstance(text, (tuple, list)): + text = [text] + + def process(text: str): + if clean_caption: + text = self._clean_caption(text) + text = self._clean_caption(text) + else: + text = text.lower().strip() + return text + + return [process(t) for t in text] + + # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption + def _clean_caption(self, caption): + caption = str(caption) + caption = ul.unquote_plus(caption) + caption = caption.strip().lower() + caption = re.sub("", "person", caption) + # urls: + caption = re.sub( + r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa + "", + caption, + ) # regex for urls + caption = re.sub( + r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa + "", + caption, + ) # regex for urls + # html: + caption = BeautifulSoup(caption, features="html.parser").text + + # @ + caption = re.sub(r"@[\w\d]+\b", "", caption) + + # 31C0—31EF CJK Strokes + # 31F0—31FF Katakana Phonetic Extensions + # 3200—32FF Enclosed CJK Letters and Months + # 3300—33FF CJK Compatibility + # 3400—4DBF CJK Unified Ideographs Extension A + # 4DC0—4DFF Yijing Hexagram Symbols + # 4E00—9FFF CJK Unified Ideographs + caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) + caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) + caption = re.sub(r"[\u3200-\u32ff]+", "", caption) + caption = re.sub(r"[\u3300-\u33ff]+", "", caption) + caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) + caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) + # caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) + ####################################################### + + # все виды тире / all types of dash --> "-" + caption = re.sub( + r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa + "-", + caption, + ) + + # кавычки к одному стандарту + caption = re.sub(r"[`´«»“”¨]", '"', caption) + caption = re.sub(r"[‘’]", "'", caption) + + # " + caption = re.sub(r""?", "", caption) + # & + caption = re.sub(r"&", "", caption) + + # ip adresses: + caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) + + # article ids: + caption = re.sub(r"\d:\d\d\s+$", "", caption) + + # \n + caption = re.sub(r"\\n", " ", caption) + + # "#123" + caption = re.sub(r"#\d{1,3}\b", "", caption) + # "#12345.." + caption = re.sub(r"#\d{5,}\b", "", caption) + # "123456.." + caption = re.sub(r"\b\d{6,}\b", "", caption) + # filenames: + caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) + + # + caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" + caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" + + caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT + caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " + + # this-is-my-cute-cat / this_is_my_cute_cat + regex2 = re.compile(r"(?:\-|\_)") + if len(re.findall(regex2, caption)) > 3: + caption = re.sub(regex2, " ", caption) + + caption = ftfy.fix_text(caption) + caption = html.unescape(html.unescape(caption)) + + caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 + caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc + caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 + + caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) + caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) + caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) + caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) + caption = re.sub(r"\bpage\s+\d+\b", "", caption) + + caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... + + caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) + + caption = re.sub(r"\b\s+\:\s+", r": ", caption) + caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) + caption = re.sub(r"\s+", " ", caption) + + caption.strip() + + caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) + caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) + caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) + caption = re.sub(r"^\.\S+$", "", caption) + return caption.strip() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents + def prepare_latents(self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None): + shape = ( + batch_size, + num_channels_latents, + (math.ceil((int(num_frames) - 1) / self.vae.vae_scale_factor[0]) + 1) if int(num_frames) % 2 == 1 else math.ceil(int(num_frames) / self.vae.vae_scale_factor[0]), + math.ceil(int(height) / self.vae.vae_scale_factor[1]), + math.ceil(int(width) / self.vae.vae_scale_factor[2]), + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + + + return latents + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: str = "", + num_inference_steps: int = 20, + timesteps: List[int] = None, + guidance_scale: float = 4.5, + num_images_per_prompt: Optional[int] = 1, + num_frames: Optional[int] = None, + height: Optional[int] = None, + width: Optional[int] = None, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_attention_mask: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_attention_mask: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: int = 1, + clean_caption: bool = True, + use_resolution_binning: bool = True, + max_sequence_length: int = 300, + **kwargs, + ) -> Union[ImagePipelineOutput, Tuple]: + """ + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_inference_steps (`int`, *optional*, defaults to 100): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` + timesteps are used. Must be in descending order. + guidance_scale (`float`, *optional*, defaults to 4.5): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + height (`int`, *optional*, defaults to self.unet.config.sample_size): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to self.unet.config.sample_size): + The width in pixels of the generated image. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to + [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + negative_prompt_attention_mask (`torch.FloatTensor`, *optional*): + Pre-generated attention mask for negative text embeddings. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple. + callback (`Callable`, *optional*): + A function that will be called every `callback_steps` steps during inference. The function will be + called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function will be called. If not specified, the callback will be + called at every step. + clean_caption (`bool`, *optional*, defaults to `True`): + Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to + be installed. If the dependencies are not installed, the embeddings will be created from the raw + prompt. + use_resolution_binning (`bool` defaults to `True`): + If set to `True`, the requested height and width are first mapped to the closest resolutions using + `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to + the requested resolution. Useful for generating non-square images. + max_sequence_length (`int` defaults to 120): Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.ImagePipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is + returned where the first element is a list with the generated images + """ + # 1. Check inputs. Raise error if not correct + num_frames = num_frames or self.transformer.config.sample_size_t * self.vae.vae_scale_factor[0] + height = height or self.transformer.config.sample_size[0] * self.vae.vae_scale_factor[1] + width = width or self.transformer.config.sample_size[1] * self.vae.vae_scale_factor[2] + self.check_inputs( + prompt, + num_frames, + height, + width, + negative_prompt, + callback_steps, + prompt_embeds, + negative_prompt_embeds, + prompt_attention_mask, + negative_prompt_attention_mask, + ) + + # 2. Default height and width to transformer + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + # import ipdb;ipdb.set_trace() + device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda') + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt, + do_classifier_free_guidance, + negative_prompt=negative_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + ) + if do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + # 4. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + + # 5. Prepare latents. + latent_channels = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + latent_channels, + num_frames, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 6.1 Prepare micro-conditions. + added_cond_kwargs = {"resolution": None, "aspect_ratio": None} + + # 7. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + current_timestep = t + if not torch.is_tensor(current_timestep): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = latent_model_input.device.type == "mps" + if isinstance(current_timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device) + elif len(current_timestep.shape) == 0: + current_timestep = current_timestep[None].to(latent_model_input.device) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + current_timestep = current_timestep.expand(latent_model_input.shape[0]) + + # import ipdb;ipdb.set_trace() + if prompt_embeds.ndim == 3: + prompt_embeds = prompt_embeds.unsqueeze(1) # b l d -> b 1 l d + if prompt_attention_mask.ndim == 2: + prompt_attention_mask = prompt_attention_mask.unsqueeze(1) # b l -> b 1 l + # prepare attention_mask. + # b c t h w -> b t h w + attention_mask = torch.ones_like(latent_model_input)[:, 0] + # predict noise model_output + noise_pred = self.transformer( + latent_model_input, + attention_mask=attention_mask, + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + timestep=current_timestep, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + )[0] + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # learned sigma + if self.transformer.config.out_channels // 2 == latent_channels: + noise_pred = noise_pred.chunk(2, dim=1)[0] + else: + noise_pred = noise_pred + + # compute previous image: x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + # print(f'latents_{i}_{t}', torch.max(latents), torch.min(latents), torch.mean(latents), torch.std(latents)) + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + # import ipdb;ipdb.set_trace() + # latents = latents.squeeze(2) + if not output_type == "latent": + # b t h w c + image = self.decode_latents(latents) + image = image[:, :num_frames, :height, :width] + else: + image = latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return ImagePipelineOutput(images=image) + + + def decode_latents(self, latents): + # print(f'before vae decode', torch.max(latents).item(), torch.min(latents).item(), torch.mean(latents).item(), torch.std(latents).item()) + video = self.vae.decode(latents.to(self.vae.vae.dtype)) + # print(f'after vae decode', torch.max(video).item(), torch.min(video).item(), torch.mean(video).item(), torch.std(video).item()) + video = ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, 2).contiguous() # b t h w c + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 + return video \ No newline at end of file diff --git a/opensora/sample/sample_t2v_sp.py b/opensora/sample/sample_t2v_sp.py index 64cc10d74..5a593320d 100644 --- a/opensora/sample/sample_t2v_sp.py +++ b/opensora/sample/sample_t2v_sp.py @@ -28,16 +28,7 @@ from opensora.utils.utils import save_video_grid from opensora.sample.pipeline_opensora_sp import OpenSoraPipeline import imageio - -try: - import torch_npu - from opensora.npu_config import npu_config - from opensora.acceleration.parallel_states import initialize_sequence_parallel_state, hccl_info -except: - torch_npu = None - npu_config = None - from opensora.utils.parallel_states import initialize_sequence_parallel_state, nccl_info - pass +from opensora.utils.parallel_states import initialize_sequence_parallel_state, nccl_info import time @@ -194,17 +185,13 @@ if __name__ == "__main__": parser.add_argument('--save_memory', action='store_true') args = parser.parse_args() - if torch_npu is not None: - npu_config.print_msg(args) # 初始化分布式环境 local_rank = int(os.getenv('RANK', 0)) world_size = int(os.getenv('WORLD_SIZE', 1)) print('world_size', world_size) - if torch_npu is not None and npu_config.on_npu: - torch_npu.npu.set_device(local_rank) - else: - torch.cuda.set_device(local_rank) + + torch.cuda.set_device(local_rank) dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank) initialize_sequence_parallel_state(world_size) # torch.manual_seed(args.seed) @@ -275,19 +262,6 @@ if __name__ == "__main__": latest_path = None save_img_path = args.save_img_path - # while True: - # cur_path = get_latest_path() - # print(cur_path, latest_path) - # if cur_path == latest_path: - # time.sleep(5) - # continue - # time.sleep(1) - # latest_path = cur_path - - # if npu_config is not None: - # npu_config.print_msg(f"The latest_path is {latest_path}") - # else: - # print(f"The latest_path is {latest_path}") if latest_path is None: @@ -297,26 +271,5 @@ if __name__ == "__main__": # full_path = f"{args.model_path}/{latest_path}/model" pipeline = load_t2v_checkpoint(full_path) print('load model') - if npu_config is not None and npu_config.on_npu and npu_config.profiling: - experimental_config = torch_npu.profiler._ExperimentalConfig( - profiler_level=torch_npu.profiler.ProfilerLevel.Level1, - aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization - ) - profile_output_path = "/home/image_data/npu_profiling_t2v" - os.makedirs(profile_output_path, exist_ok=True) - - with torch_npu.profiler.profile( - activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU], - with_stack=True, - record_shapes=True, - profile_memory=True, - experimental_config=experimental_config, - schedule=torch_npu.profiler.schedule(wait=10000, warmup=0, active=1, repeat=1, - skip_first=0), - on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/") - ) as prof: - run_model_and_save_images(pipeline, latest_path) - prof.step() - else: - # print('gpu') - run_model_and_save_images(pipeline, latest_path) + + run_model_and_save_images(pipeline, latest_path) diff --git a/opensora/train/train_t2v_diffusers.py b/opensora/train/train_t2v_diffusers.py index 2c00977b8..7d076e726 100644 --- a/opensora/train/train_t2v_diffusers.py +++ b/opensora/train/train_t2v_diffusers.py @@ -13,7 +13,6 @@ import math import os import shutil from pathlib import Path -from typing import Optional import gc import numpy as np from einops import rearrange @@ -21,13 +20,10 @@ from tqdm import tqdm from opensora.adaptor.modules import replace_with_fp32_forwards - - from opensora.utils.parallel_states import initialize_sequence_parallel_state, \ destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state from opensora.utils.communications import prepare_parallel_data, broadcast import time -from dataclasses import field, dataclass from torch.utils.data import DataLoader from copy import deepcopy import accelerate @@ -678,7 +674,7 @@ def main(args): log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, weight_dtype, progress_info.global_step) - if args.use_ema and npu_config is None: + if args.use_ema: # Store the UNet parameters temporarily and load the EMA parameters to perform inference. ema_model.store(model.parameters()) ema_model.copy_to(model.parameters()) diff --git a/scripts/slurm/placeholder b/scripts/slurm/placeholder deleted file mode 100644 index e69de29bb..000000000 diff --git a/scripts/text_condition/sample_inpaint.sh b/scripts/text_condition/sample_inpaint.sh deleted file mode 100644 index 1a340aa3f..000000000 --- a/scripts/text_condition/sample_inpaint.sh +++ /dev/null @@ -1,18 +0,0 @@ -CUDA_VISIBLE_DEVICES=7 python -m opensora.sample.sample_inpaint \ - --model_path /path/to/checkpoint-xxx/model \ - --num_frames 93 \ - --height 480 \ - --width 640 \ - --cache_dir "../cache_dir" \ - --text_encoder_name google/mt5-xxl \ - --text_prompt examples/prompt_list_0.txt \ - --conditional_images_path examples/conditional_images_path.txt \ - --ae CausalVAEModel_D4_4x8x8 \ - --ae_path "/path/to/causalvideovae" \ - --save_img_path "./sample_video_93x480p_cfg7.5_step100" \ - --fps 24 \ - --guidance_scale 7.5 \ - --num_sampling_steps 100 \ - --enable_tiling \ - --max_sequence_length 512 \ - --sample_method PDNM \ diff --git a/scripts/text_condition/sample_inpaint_ddp.sh b/scripts/text_condition/sample_inpaint_ddp.sh deleted file mode 100644 index a2706e7b9..000000000 --- a/scripts/text_condition/sample_inpaint_ddp.sh +++ /dev/null @@ -1,21 +0,0 @@ - -torchrun --nnodes=1 --nproc_per_node 8 --master_port 29503 \ - -m opensora.sample.sample_inpaint_ddp \ - --model_path /path/to/checkpoint-xxx/model \ - --num_frames 93 \ - --height 480 \ - --width 640 \ - --cache_dir "../cache_dir" \ - --text_encoder_name google/mt5-xxl \ - --text_prompt examples/prompt_list_0.txt \ - --conditional_images_path examples/conditional_images_path.txt \ - --ae CausalVAEModel_D4_4x8x8 \ - --ae_path "/path/to/causalvideovae" \ - --save_img_path "./sample_video_93x480p_cfg7.5_step100" \ - --fps 24 \ - --guidance_scale 7.5 \ - --num_sampling_steps 100 \ - --enable_tiling \ - --tile_overlap_factor 0.125 \ - --max_sequence_length 512 \ - --sample_method PNDM \ \ No newline at end of file diff --git a/scripts/text_condition/sample_inpaint_sp.sh b/scripts/text_condition/sample_inpaint_sp.sh deleted file mode 100644 index 9aa9c1b4c..000000000 --- a/scripts/text_condition/sample_inpaint_sp.sh +++ /dev/null @@ -1,21 +0,0 @@ - -torchrun --nnodes=1 --nproc_per_node 8 --master_port 29503 \ - -m opensora.sample.sample_inpaint_sp \ - --model_path /path/to/checkpoint-xxx/model \ - --num_frames 93 \ - --height 480 \ - --width 640 \ - --cache_dir "../cache_dir" \ - --text_encoder_name google/mt5-xxl \ - --text_prompt examples/prompt_list_0.txt \ - --conditional_images_path examples/conditional_images_path.txt \ - --ae CausalVAEModel_D4_4x8x8 \ - --ae_path "/path/to/causalvideovae" \ - --save_img_path "./sample_video_93x480p_cfg7.5_step100" \ - --fps 24 \ - --guidance_scale 7.5 \ - --num_sampling_steps 100 \ - --enable_tiling \ - --tile_overlap_factor 0.125 \ - --max_sequence_length 512 \ - --sample_method PNDM \ \ No newline at end of file diff --git a/scripts/text_condition/train_inpaint.sh b/scripts/text_condition/train_inpaint.sh deleted file mode 100644 index 2276883d6..000000000 --- a/scripts/text_condition/train_inpaint.sh +++ /dev/null @@ -1,74 +0,0 @@ -PROJECT="your project name" -export WANDB_API_KEY="your wanb key" -export WANDB_MODE="online" -export ENTITY="your name" -export PROJECT=$PROJECT -export HF_DATASETS_OFFLINE=1 -export TRANSFORMERS_OFFLINE=1 -export TOKENIZERS_PARALLELISM=false -# NCCL setting -export GLOO_SOCKET_IFNAME=bond0 -export NCCL_SOCKET_IFNAME=bond0 -export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1 -export NCCL_IB_GID_INDEX=3 -export NCCL_IB_TC=162 -export NCCL_IB_TIMEOUT=22 -export NCCL_PXN_DISABLE=0 -export NCCL_IB_QPS_PER_CONNECTION=4 -export NCCL_ALGO=Ring -export OMP_NUM_THREADS=1 -export MKL_NUM_THREADS=1 - -# export NCCL_ALGO=Tree - -accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ - opensora/train/train_inpaint.py \ - --model OpenSoraInpaint-ROPE-L/122 \ - --text_encoder_name google/mt5-xxl \ - --cache_dir "../cache_dir" \ - --dataset inpaint \ - --data "scripts/train_data/merge_data.txt" \ - --ae CausalVAEModel_D4_4x8x8 \ - --ae_path "/path/to/causalvideovae" \ - --sample_rate 1 \ - --num_frames 93 \ - --use_image_num 0 \ - --max_height 480 \ - --max_width 640 \ - --interpolation_scale_t 1.0 \ - --interpolation_scale_h 1.0 \ - --interpolation_scale_w 1.0 \ - --attention_mode xformers \ - --gradient_checkpointing \ - --train_batch_size=1 \ - --dataloader_num_workers 10 \ - --gradient_accumulation_steps=1 \ - --max_train_steps=1000000 \ - --learning_rate=1e-5 \ - --lr_scheduler="constant" \ - --lr_warmup_steps=0 \ - --mixed_precision="bf16" \ - --report_to="wandb" \ - --checkpointing_steps=200 \ - --allow_tf32 \ - --model_max_length 512 \ - --enable_tiling \ - --tile_overlap_factor 0.125 \ - --snr_gamma 5.0 \ - --use_ema \ - --ema_start_step 0 \ - --cfg 0.1 \ - --noise_offset 0.02 \ - --i2v_ratio 0.4 \ - --transition_ratio 0.4 \ - --v2v_ratio 0.1 \ - --clear_video_ratio 0.0 \ - --default_text_ratio 0.5 \ - --use_rope \ - --ema_decay 0.999 \ - --speed_factor 1.0 \ - --group_frame \ - --output_dir runs/$PROJECT \ - --pretrained_transformer_model_path "/path/to/pretrained_model" \ - # --resume_from_checkpoint="latest" \ diff --git a/scripts/train_data/merge_data_on_npu.txt b/scripts/train_data/merge_data_on_npu.txt deleted file mode 100644 index d7e10fc05..000000000 --- a/scripts/train_data/merge_data_on_npu.txt +++ /dev/null @@ -1 +0,0 @@ -/home/local_dataset_6t/video_data_obs/movie_aes_motion,/home/opensora/captions/movie_aes_motion/bbc01_clips_aes_motion_final_250508.json \ No newline at end of file