diff --git "a/easyanimate/pipeline/pipeline_easyanimate_inpaint.py" "b/easyanimate/pipeline/pipeline_easyanimate_inpaint.py"
--- "a/easyanimate/pipeline/pipeline_easyanimate_inpaint.py"
+++ "b/easyanimate/pipeline/pipeline_easyanimate_inpaint.py"
@@ -1,4 +1,4 @@
-# Copyright 2023 PixArt-Alpha Authors and The HuggingFace Team. All rights reserved.
+# Copyright 2024 EasyAnimate 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.
@@ -12,149 +12,336 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-import copy
-import gc
-import html
import inspect
-import re
-import urllib.parse as ul
-from dataclasses import dataclass
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
import numpy as np
import torch
import torch.nn.functional as F
-from diffusers import DiffusionPipeline, ImagePipelineOutput
+from dataclasses import dataclass
+from diffusers import DiffusionPipeline
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
from diffusers.image_processor import VaeImageProcessor
-from diffusers.models import AutoencoderKL
-from diffusers.schedulers import DPMSolverMultistepScheduler
+from diffusers.models import AutoencoderKL, HunyuanDiT2DModel
+from diffusers.models.embeddings import (get_2d_rotary_pos_embed,
+ get_3d_rotary_pos_embed)
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion.safety_checker import \
+ StableDiffusionSafetyChecker
+from diffusers.schedulers import DDIMScheduler, FlowMatchEulerDiscreteScheduler
from diffusers.utils import (BACKENDS_MAPPING, BaseOutput, deprecate,
- is_bs4_available, is_ftfy_available, logging,
+ is_bs4_available, is_ftfy_available,
+ is_torch_xla_available, logging,
replace_example_docstring)
from diffusers.utils.torch_utils import randn_tensor
from einops import rearrange
from PIL import Image
from tqdm import tqdm
-from transformers import (CLIPImageProcessor, CLIPVisionModelWithProjection,
+from transformers import (BertModel, BertTokenizer, CLIPImageProcessor,
+ Qwen2Tokenizer, Qwen2VLForConditionalGeneration, CLIPVisionModelWithProjection,
T5EncoderModel, T5Tokenizer)
-from ..models.transformer3d import Transformer3DModel
+from ..models import AutoencoderKLMagvit, EasyAnimateTransformer3DModel
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+if is_torch_xla_available():
+ import torch_xla.core.xla_model as xm
-if is_bs4_available():
- from bs4 import BeautifulSoup
+ XLA_AVAILABLE = True
+else:
+ XLA_AVAILABLE = False
-if is_ftfy_available():
- import ftfy
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> import torch
- >>> from diffusers import EasyAnimatePipeline
-
- >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-XL-2-512x512" too.
- >>> pipe = EasyAnimatePipeline.from_pretrained("PixArt-alpha/PixArt-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]
+ >>> from diffusers import EasyAnimateInpaintPipeline
+ >>> from diffusers.easyanimate.pipeline_easyanimate_inpaint import get_image_to_video_latent
+ >>> from diffusers.utils import export_to_video, load_image
+
+ >>> pipe = EasyAnimateInpaintPipeline.from_pretrained("alibaba-pai/EasyAnimateV5.1-12b-zh-InP", torch_dtype=torch.bfloat16)
+ >>> pipe.to("cuda")
+
+ >>> prompt = "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+ >>> validation_image_start = load_image(
+ ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+ ... )
+ >>> validation_image_end = None
+ >>> sample_size = (576, 448)
+ >>> video_length = 49
+ >>> input_video, input_video_mask, _ = get_image_to_video_latent(validation_image_start, validation_image_end, video_length, sample_size)
+ >>> video = pipe(prompt, video_length=video_length, negative_prompt="bad detailed", height=sample_size[0], width=sample_size[1], input_video=input_video, mask_video=input_video_mask)
+ >>> export_to_video(video.sample[0], "output.mp4", fps=8)
```
"""
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
-def retrieve_latents(encoder_output, generator):
- if hasattr(encoder_output, "latent_dist"):
- return encoder_output.latent_dist.sample(generator)
- elif hasattr(encoder_output, "latents"):
- return encoder_output.latents
+
+
+# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid
+def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
+ tw = tgt_width
+ th = tgt_height
+ h, w = src
+ r = h / w
+ if r > (th / tw):
+ resize_height = th
+ resize_width = int(round(th / h * w))
+ else:
+ resize_width = tw
+ resize_height = int(round(tw / w * h))
+
+ crop_top = int(round((th - resize_height) / 2.0))
+ crop_left = int(round((tw - resize_width) / 2.0))
+
+ return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width)
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+ """
+ Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+ Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+ """
+ std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+ std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+ # rescale the results from guidance (fixes overexposure)
+ noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+ # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+ noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+ return noise_cfg
+
+
+# Resize mask information in magvit
+def resize_mask(mask, latent, process_first_frame_only=True):
+ latent_size = latent.size()
+
+ if process_first_frame_only:
+ target_size = list(latent_size[2:])
+ target_size[0] = 1
+ first_frame_resized = F.interpolate(
+ mask[:, :, 0:1, :, :],
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+
+ target_size = list(latent_size[2:])
+ target_size[0] = target_size[0] - 1
+ if target_size[0] != 0:
+ remaining_frames_resized = F.interpolate(
+ mask[:, :, 1:, :, :],
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+ resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2)
+ else:
+ resized_mask = first_frame_resized
else:
- raise AttributeError("Could not access latents of provided encoder_output")
+ target_size = list(latent_size[2:])
+ resized_mask = F.interpolate(
+ mask,
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+ return resized_mask
+
+
+## Add noise to reference video
+def add_noise_to_reference_video(image, ratio=None):
+ if ratio is None:
+ sigma = torch.normal(mean=-3.0, std=0.5, size=(image.shape[0],)).to(image.device)
+ sigma = torch.exp(sigma).to(image.dtype)
+ else:
+ sigma = torch.ones((image.shape[0],)).to(image.device, image.dtype) * ratio
+ image_noise = torch.randn_like(image) * sigma[:, None, None, None, None]
+ image_noise = torch.where(image==-1, torch.zeros_like(image), image_noise)
+ image = image + image_noise
+ return image
+
+
+# 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,
+ sigmas: Optional[List[float]] = 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 override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` 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 and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ 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)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, 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
+
+
@dataclass
class EasyAnimatePipelineOutput(BaseOutput):
- videos: Union[torch.Tensor, np.ndarray]
+ r"""
+ Output class for EasyAnimate pipelines.
+
+ Args:
+ frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
+ List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
+ `(batch_size, num_frames, channels, height, width)`.
+ """
+
+ frames: torch.Tensor
+
class EasyAnimateInpaintPipeline(DiffusionPipeline):
r"""
- Pipeline for text-to-image generation using PixArt-Alpha.
+ Pipeline for text-to-video generation using EasyAnimate.
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+ EasyAnimate uses one text encoder [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1.
+ EasyAnimate uses two text encoders: [mT5](https://huggingface.co/google/mt5-base) and [bilingual CLIP](fine-tuned by
+ HunyuanDiT team) in V5.
+
Args:
- vae ([`AutoencoderKL`]):
- Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
- text_encoder ([`T5EncoderModel`]):
- Frozen text-encoder. PixArt-Alpha uses
- [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
- [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant.
- tokenizer (`T5Tokenizer`):
- Tokenizer of class
- [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
- transformer ([`Transformer3DModel`]):
- A text conditioned `Transformer3DModel` to denoise the encoded image latents.
- scheduler ([`SchedulerMixin`]):
- A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+ vae ([`AutoencoderKLMagvit`]):
+ Variational Auto-Encoder (VAE) Model to encode and decode video to and from latent representations.
+ text_encoder (Optional[`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel`]):
+ EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1.
+ EasyAnimate uses [bilingual CLIP](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers) in V5.
+ tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]):
+ A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text.
+ transformer ([`EasyAnimateTransformer3DModel`]):
+ The EasyAnimate model designed by EasyAnimate Team.
+ text_encoder_2 (`T5EncoderModel`):
+ EasyAnimate does not use text_encoder_2 in V5.1.
+ EasyAnimate uses [mT5](https://huggingface.co/google/mt5-base) embedder in V5.
+ tokenizer_2 (`T5Tokenizer`):
+ The tokenizer for the mT5 embedder.
+ scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+ A scheduler to be used in combination with EasyAnimate to denoise the encoded image latents.
+ clip_image_processor (`CLIPImageProcessor`):
+ The CLIP image embedder.
+ clip_image_encoder (`CLIPVisionModelWithProjection`):
+ The image processor for the CLIP image embedder.
"""
- bad_punct_regex = re.compile(
- r"[" + "#®•©™&@·º½¾¿¡§~" + "\)" + "\(" + "\]" + "\[" + "\}" + "\{" + "\|" + "\\" + "\/" + "\*" + r"]{1,}"
- ) # noqa
- _optional_components = ["tokenizer", "text_encoder"]
- model_cpu_offload_seq = "text_encoder->transformer->vae"
+ model_cpu_offload_seq = "text_encoder->text_encoder_2->clip_image_encoder->transformer->vae"
+ _optional_components = [
+ "text_encoder_2",
+ "tokenizer_2",
+ "text_encoder",
+ "tokenizer",
+ "clip_image_encoder",
+ ]
+ _callback_tensor_inputs = [
+ "latents",
+ "prompt_embeds",
+ "negative_prompt_embeds",
+ "prompt_embeds_2",
+ "negative_prompt_embeds_2",
+ ]
def __init__(
self,
- tokenizer: T5Tokenizer,
- text_encoder: T5EncoderModel,
- vae: AutoencoderKL,
- transformer: Transformer3DModel,
- scheduler: DPMSolverMultistepScheduler,
- clip_image_processor:CLIPImageProcessor = None,
- clip_image_encoder:CLIPVisionModelWithProjection = None,
+ vae: AutoencoderKLMagvit,
+ text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel],
+ tokenizer: Union[Qwen2Tokenizer, BertTokenizer],
+ text_encoder_2: Optional[Union[T5EncoderModel, Qwen2VLForConditionalGeneration]],
+ tokenizer_2: Optional[Union[T5Tokenizer, Qwen2Tokenizer]],
+ transformer: EasyAnimateTransformer3DModel,
+ scheduler: FlowMatchEulerDiscreteScheduler,
+ clip_image_processor: CLIPImageProcessor = None,
+ clip_image_encoder: CLIPVisionModelWithProjection = None,
):
super().__init__()
self.register_modules(
- tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer,
+ vae=vae,
+ text_encoder=text_encoder,
+ tokenizer=tokenizer,
+ tokenizer_2=tokenizer_2,
+ transformer=transformer,
scheduler=scheduler,
- clip_image_processor=clip_image_processor, clip_image_encoder=clip_image_encoder,
+ text_encoder_2=text_encoder_2,
+ clip_image_processor=clip_image_processor,
+ clip_image_encoder=clip_image_encoder,
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
- self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_normalize=True)
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
self.mask_processor = VaeImageProcessor(
vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True
)
- self.enable_autocast_float8_transformer_flag = False
- # Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/utils.py
- def mask_text_embeddings(self, emb, mask):
- if emb.shape[0] == 1:
- keep_index = mask.sum().item()
- return emb[:, :, :keep_index, :], keep_index
- else:
- masked_feature = emb * mask[:, None, :, None]
- return masked_feature, emb.shape[2]
+ def enable_sequential_cpu_offload(self, *args, **kwargs):
+ super().enable_sequential_cpu_offload(*args, **kwargs)
+ if hasattr(self.transformer, "clip_projection") and self.transformer.clip_projection is not None:
+ import accelerate
+ accelerate.hooks.remove_hook_from_module(self.transformer.clip_projection, recurse=True)
+ self.transformer.clip_projection = self.transformer.clip_projection.to("cuda")
- # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt
def encode_prompt(
self,
- prompt: Union[str, List[str]],
- do_classifier_free_guidance: bool = True,
- negative_prompt: str = "",
+ prompt: str,
+ device: torch.device,
+ dtype: torch.dtype,
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,
+ do_classifier_free_guidance: bool = True,
+ negative_prompt: Optional[str] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_attention_mask: Optional[torch.Tensor] = None,
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ max_sequence_length: Optional[int] = None,
+ text_encoder_index: int = 0,
+ actual_max_sequence_length: int = 256
):
r"""
Encodes the prompt into text encoder hidden states.
@@ -162,33 +349,46 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
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):
+ device: (`torch.device`):
+ torch device
+ dtype (`torch.dtype`):
+ torch dtype
+ num_images_per_prompt (`int`):
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*):
+ do_classifier_free_guidance (`bool`):
+ whether to use classifier free guidance or not
+ 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`).
+ prompt_embeds (`torch.Tensor`, *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.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ prompt_attention_mask (`torch.Tensor`, *optional*):
+ Attention mask for the prompt. Required when `prompt_embeds` is passed directly.
+ negative_prompt_attention_mask (`torch.Tensor`, *optional*):
+ Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly.
+ max_sequence_length (`int`, *optional*): maximum sequence length to use for the prompt.
+ text_encoder_index (`int`, *optional*):
+ Index of the text encoder to use. `0` for clip and `1` for T5.
"""
+ tokenizers = [self.tokenizer, self.tokenizer_2]
+ text_encoders = [self.text_encoder, self.text_encoder_2]
- 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)
+ tokenizer = tokenizers[text_encoder_index]
+ text_encoder = text_encoders[text_encoder_index]
- if device is None:
- device = self._execution_device
+ if max_sequence_length is None:
+ if text_encoder_index == 0:
+ max_length = min(self.tokenizer.model_max_length, actual_max_sequence_length)
+ if text_encoder_index == 1:
+ max_length = min(self.tokenizer_2.model_max_length, actual_max_sequence_length)
+ else:
+ max_length = max_sequence_length
if prompt is not None and isinstance(prompt, str):
batch_size = 1
@@ -197,74 +397,199 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
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}"
+ if type(tokenizer) in [BertTokenizer, T5Tokenizer]:
+ text_inputs = tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_attention_mask=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ if text_input_ids.shape[-1] > actual_max_sequence_length:
+ reprompt = tokenizer.batch_decode(text_input_ids[:, :actual_max_sequence_length], skip_special_tokens=True)
+ text_inputs = tokenizer(
+ reprompt,
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_attention_mask=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ untruncated_ids = 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
+ ):
+ _actual_max_sequence_length = min(tokenizer.model_max_length, actual_max_sequence_length)
+ removed_text = tokenizer.batch_decode(untruncated_ids[:, _actual_max_sequence_length - 1 : -1])
+ logger.warning(
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
+ f" {_actual_max_sequence_length} tokens: {removed_text}"
+ )
+
+ prompt_attention_mask = text_inputs.attention_mask.to(device)
+
+ if self.transformer.config.enable_text_attention_mask:
+ prompt_embeds = text_encoder(
+ text_input_ids.to(device),
+ attention_mask=prompt_attention_mask,
+ )
+ else:
+ prompt_embeds = text_encoder(
+ text_input_ids.to(device)
+ )
+ prompt_embeds = prompt_embeds[0]
+ prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
+ else:
+ if prompt is not None and isinstance(prompt, str):
+ messages = [
+ {
+ "role": "user",
+ "content": [{"type": "text", "text": prompt}],
+ }
+ ]
+ else:
+ messages = [
+ {
+ "role": "user",
+ "content": [{"type": "text", "text": _prompt}],
+ } for _prompt in prompt
+ ]
+ text = tokenizer.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True
)
- 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
-
+ text_inputs = tokenizer(
+ text=[text],
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_attention_mask=True,
+ padding_side="right",
+ return_tensors="pt",
+ )
+ text_inputs = text_inputs.to(text_encoder.device)
+
+ text_input_ids = text_inputs.input_ids
+ prompt_attention_mask = text_inputs.attention_mask
+ if self.transformer.config.enable_text_attention_mask:
+ # Inference: Generation of the output
+ prompt_embeds = text_encoder(
+ input_ids=text_input_ids,
+ attention_mask=prompt_attention_mask,
+ output_hidden_states=True).hidden_states[-2]
+ else:
+ raise ValueError("LLM needs attention_mask")
+ prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
+
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
+ # duplicate text embeddings 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)
+ prompt_attention_mask = prompt_attention_mask.to(device=device)
# 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)
+ if type(tokenizer) in [BertTokenizer, T5Tokenizer]:
+ uncond_tokens: List[str]
+ if negative_prompt is None:
+ uncond_tokens = [""] * batch_size
+ elif prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif isinstance(negative_prompt, str):
+ uncond_tokens = [negative_prompt]
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+ else:
+ uncond_tokens = negative_prompt
+
+ max_length = prompt_embeds.shape[1]
+ uncond_input = tokenizer(
+ uncond_tokens,
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+ uncond_input_ids = uncond_input.input_ids
+ if uncond_input_ids.shape[-1] > actual_max_sequence_length:
+ reuncond_tokens = tokenizer.batch_decode(uncond_input_ids[:, :actual_max_sequence_length], skip_special_tokens=True)
+ uncond_input = tokenizer(
+ reuncond_tokens,
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_attention_mask=True,
+ return_tensors="pt",
+ )
+ uncond_input_ids = uncond_input.input_ids
+
+ negative_prompt_attention_mask = uncond_input.attention_mask.to(device)
+ if self.transformer.config.enable_text_attention_mask:
+ negative_prompt_embeds = text_encoder(
+ uncond_input.input_ids.to(device),
+ attention_mask=negative_prompt_attention_mask,
+ )
+ else:
+ negative_prompt_embeds = text_encoder(
+ uncond_input.input_ids.to(device)
+ )
+ negative_prompt_embeds = negative_prompt_embeds[0]
+ negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
+ else:
+ if negative_prompt is not None and isinstance(negative_prompt, str):
+ messages = [
+ {
+ "role": "user",
+ "content": [{"type": "text", "text": negative_prompt}],
+ }
+ ]
+ else:
+ messages = [
+ {
+ "role": "user",
+ "content": [{"type": "text", "text": _negative_prompt}],
+ } for _negative_prompt in negative_prompt
+ ]
+ text = tokenizer.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True
+ )
- 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]
+ text_inputs = tokenizer(
+ text=[text],
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_attention_mask=True,
+ padding_side="right",
+ return_tensors="pt",
+ )
+ text_inputs = text_inputs.to(text_encoder.device)
+
+ text_input_ids = text_inputs.input_ids
+ negative_prompt_attention_mask = text_inputs.attention_mask
+ if self.transformer.config.enable_text_attention_mask:
+ # Inference: Generation of the output
+ negative_prompt_embeds = text_encoder(
+ input_ids=text_input_ids,
+ attention_mask=negative_prompt_attention_mask,
+ output_hidden_states=True).hidden_states[-2]
+ else:
+ raise ValueError("LLM needs attention_mask")
+ negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
if do_classifier_free_guidance:
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
@@ -274,14 +599,9 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
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
+ negative_prompt_attention_mask = negative_prompt_attention_mask.to(device=device)
+
+ return prompt_embeds, negative_prompt_embeds, prompt_attention_mask, 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):
@@ -306,20 +626,25 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
prompt,
height,
width,
- negative_prompt,
- callback_steps,
+ negative_prompt=None,
prompt_embeds=None,
negative_prompt_embeds=None,
+ prompt_attention_mask=None,
+ negative_prompt_attention_mask=None,
+ prompt_embeds_2=None,
+ negative_prompt_embeds_2=None,
+ prompt_attention_mask_2=None,
+ negative_prompt_attention_mask_2=None,
+ callback_on_step_end_tensor_inputs=None,
):
- 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 height % 16 != 0 or width % 16 != 0:
+ raise ValueError(f"`height` and `width` have to be divisible by 16 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)
+ if callback_on_step_end_tensor_inputs is not None and not all(
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
):
raise ValueError(
- f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
- f" {type(callback_steps)}."
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
)
if prompt is not None and prompt_embeds is not None:
@@ -331,14 +656,18 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
+ elif prompt is None and prompt_embeds_2 is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds_2`. Cannot leave both `prompt` and `prompt_embeds_2` 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 prompt_embeds is not None and prompt_attention_mask is None:
+ raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
+
+ if prompt_embeds_2 is not None and prompt_attention_mask_2 is None:
+ raise ValueError("Must provide `prompt_attention_mask_2` when specifying `prompt_embeds_2`.")
if negative_prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
@@ -346,6 +675,13 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
+ 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 negative_prompt_embeds_2 is not None and negative_prompt_attention_mask_2 is None:
+ raise ValueError(
+ "Must provide `negative_prompt_attention_mask_2` when specifying `negative_prompt_embeds_2`."
+ )
if prompt_embeds is not None and negative_prompt_embeds is not None:
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
@@ -353,201 +689,83 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
+ if prompt_embeds_2 is not None and negative_prompt_embeds_2 is not None:
+ if prompt_embeds_2.shape != negative_prompt_embeds_2.shape:
+ raise ValueError(
+ "`prompt_embeds_2` and `negative_prompt_embeds_2` must have the same shape when passed directly, but"
+ f" got: `prompt_embeds_2` {prompt_embeds_2.shape} != `negative_prompt_embeds_2`"
+ f" {negative_prompt_embeds_2.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.warn(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
- logger.warn("Setting `clean_caption` to False...")
- clean_caption = False
-
- if clean_caption and not is_ftfy_available():
- logger.warn(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
- logger.warn("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>", "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
-
- # @<nickname>
- 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()
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps
+ def get_timesteps(self, num_inference_steps, strength, device):
+ # get the original timestep using init_timestep
+ init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
- 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)
+ t_start = max(num_inference_steps - init_timestep, 0)
+ timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
- return caption.strip()
+ return timesteps, num_inference_steps - t_start
def prepare_mask_latents(
- self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance
+ self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance, noise_aug_strength
):
# resize the mask to latents shape as we concatenate the mask to the latents
# we do that before converting to dtype to avoid breaking in case we're using cpu_offload
# and half precision
- video_length = mask.shape[2]
-
- mask = mask.to(device=device, dtype=self.vae.dtype)
- if self.vae.quant_conv.weight.ndim==5:
- bs = 1
- new_mask = []
- for i in range(0, mask.shape[0], bs):
- mask_bs = mask[i : i + bs]
- mask_bs = self.vae.encode(mask_bs)[0]
- mask_bs = mask_bs.sample()
- new_mask.append(mask_bs)
- mask = torch.cat(new_mask, dim = 0)
- mask = mask * self.vae.config.scaling_factor
+ if mask is not None:
+ mask = mask.to(device=device, dtype=dtype)
+ if self.vae.quant_conv is None or self.vae.quant_conv.weight.ndim==5:
+ bs = 1
+ new_mask = []
+ for i in range(0, mask.shape[0], bs):
+ mask_bs = mask[i : i + bs]
+ mask_bs = self.vae.encode(mask_bs)[0]
+ mask_bs = mask_bs.mode()
+ new_mask.append(mask_bs)
+ mask = torch.cat(new_mask, dim = 0)
+ mask = mask * self.vae.config.scaling_factor
- else:
- if mask.shape[1] == 4:
- mask = mask
else:
- video_length = mask.shape[2]
- mask = rearrange(mask, "b c f h w -> (b f) c h w")
- mask = self._encode_vae_image(mask, generator=generator)
- mask = rearrange(mask, "(b f) c h w -> b c f h w", f=video_length)
-
- masked_image = masked_image.to(device=device, dtype=self.vae.dtype)
- if self.vae.quant_conv.weight.ndim==5:
- bs = 1
- new_mask_pixel_values = []
- for i in range(0, masked_image.shape[0], bs):
- mask_pixel_values_bs = masked_image[i : i + bs]
- mask_pixel_values_bs = self.vae.encode(mask_pixel_values_bs)[0]
- mask_pixel_values_bs = mask_pixel_values_bs.sample()
- new_mask_pixel_values.append(mask_pixel_values_bs)
- masked_image_latents = torch.cat(new_mask_pixel_values, dim = 0)
- masked_image_latents = masked_image_latents * self.vae.config.scaling_factor
+ if mask.shape[1] == 4:
+ mask = mask
+ else:
+ video_length = mask.shape[2]
+ mask = rearrange(mask, "b c f h w -> (b f) c h w")
+ mask = self._encode_vae_image(mask, generator=generator)
+ mask = rearrange(mask, "(b f) c h w -> b c f h w", f=video_length)
+
+ if masked_image is not None:
+ masked_image = masked_image.to(device=device, dtype=dtype)
+ if self.transformer.config.add_noise_in_inpaint_model:
+ masked_image = add_noise_to_reference_video(masked_image, ratio=noise_aug_strength)
+ if self.vae.quant_conv is None or self.vae.quant_conv.weight.ndim==5:
+ bs = 1
+ new_mask_pixel_values = []
+ for i in range(0, masked_image.shape[0], bs):
+ mask_pixel_values_bs = masked_image[i : i + bs]
+ mask_pixel_values_bs = self.vae.encode(mask_pixel_values_bs)[0]
+ mask_pixel_values_bs = mask_pixel_values_bs.mode()
+ new_mask_pixel_values.append(mask_pixel_values_bs)
+ masked_image_latents = torch.cat(new_mask_pixel_values, dim = 0)
+ masked_image_latents = masked_image_latents * self.vae.config.scaling_factor
- else:
- if masked_image.shape[1] == 4:
- masked_image_latents = masked_image
else:
- video_length = mask.shape[2]
- masked_image = rearrange(masked_image, "b c f h w -> (b f) c h w")
- masked_image_latents = self._encode_vae_image(masked_image, generator=generator)
- masked_image_latents = rearrange(masked_image_latents, "(b f) c h w -> b c f h w", f=video_length)
+ if masked_image.shape[1] == 4:
+ masked_image_latents = masked_image
+ else:
+ video_length = masked_image.shape[2]
+ masked_image = rearrange(masked_image, "b c f h w -> (b f) c h w")
+ masked_image_latents = self._encode_vae_image(masked_image, generator=generator)
+ masked_image_latents = rearrange(masked_image_latents, "(b f) c h w -> b c f h w", f=video_length)
+
+ # aligning device to prevent device errors when concating it with the latent model input
+ masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
+ else:
+ masked_image_latents = None
- # aligning device to prevent device errors when concating it with the latent model input
- masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
return mask, masked_image_latents
-
+
def prepare_latents(
self,
batch_size,
@@ -565,10 +783,15 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
return_noise=False,
return_video_latents=False,
):
- if self.vae.quant_conv.weight.ndim==5:
- mini_batch_encoder = self.vae.mini_batch_encoder
- mini_batch_decoder = self.vae.mini_batch_decoder
- shape = (batch_size, num_channels_latents, int(video_length // mini_batch_encoder * mini_batch_decoder) if video_length != 1 else 1, height // self.vae_scale_factor, width // self.vae_scale_factor)
+ if self.vae.quant_conv is None or self.vae.quant_conv.weight.ndim==5:
+ if self.vae.cache_mag_vae:
+ mini_batch_encoder = self.vae.mini_batch_encoder
+ mini_batch_decoder = self.vae.mini_batch_decoder
+ shape = (batch_size, num_channels_latents, int((video_length - 1) // mini_batch_encoder * mini_batch_decoder + 1) if video_length != 1 else 1, height // self.vae_scale_factor, width // self.vae_scale_factor)
+ else:
+ mini_batch_encoder = self.vae.mini_batch_encoder
+ mini_batch_decoder = self.vae.mini_batch_decoder
+ shape = (batch_size, num_channels_latents, int(video_length // mini_batch_encoder * mini_batch_decoder) if video_length != 1 else 1, height // self.vae_scale_factor, width // self.vae_scale_factor)
else:
shape = (batch_size, num_channels_latents, video_length, height // self.vae_scale_factor, width // self.vae_scale_factor)
@@ -579,10 +802,9 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
)
if return_video_latents or (latents is None and not is_strength_max):
- video = video.to(device=device, dtype=self.vae.dtype)
- if self.vae.quant_conv.weight.ndim==5:
+ video = video.to(device=device, dtype=dtype)
+ if self.vae.quant_conv is None or self.vae.quant_conv.weight.ndim==5:
bs = 1
- mini_batch_encoder = self.vae.mini_batch_encoder
new_video = []
for i in range(0, video.shape[0], bs):
video_bs = video[i : i + bs]
@@ -601,16 +823,24 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
video = self._encode_vae_image(video, generator=generator)
video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
video_latents = video.repeat(batch_size // video.shape[0], 1, 1, 1, 1)
+ video_latents = video_latents.to(device=device, dtype=dtype)
if latents is None:
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# if strength is 1. then initialise the latents to noise, else initial to image + noise
- latents = noise if is_strength_max else self.scheduler.add_noise(video_latents, noise, timestep)
+ if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+ latents = noise if is_strength_max else self.scheduler.scale_noise(video_latents, timestep, noise)
+ else:
+ latents = noise if is_strength_max else self.scheduler.add_noise(video_latents, noise, timestep)
# if pure noise then scale the initial latents by the Scheduler's init sigma
- latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents
+ if hasattr(self.scheduler, "init_noise_sigma"):
+ latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents
else:
- noise = latents.to(device)
- latents = noise * self.scheduler.init_noise_sigma
+ if hasattr(self.scheduler, "init_noise_sigma"):
+ noise = latents.to(device)
+ latents = noise * self.scheduler.init_noise_sigma
+ else:
+ latents = latents.to(device)
# scale the initial noise by the standard deviation required by the scheduler
outputs = (latents,)
@@ -632,22 +862,23 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
# Encode middle videos
latents = self.vae.encode(pixel_values)[0]
- latents = latents.sample()
+ latents = latents.mode()
# Decode middle videos
middle_video = self.vae.decode(latents)[0]
video[:, :, prefix_index_before:-prefix_index_after] = (video[:, :, prefix_index_before:-prefix_index_after] + middle_video) / 2
return video
-
+
def decode_latents(self, latents):
video_length = latents.shape[2]
latents = 1 / self.vae.config.scaling_factor * latents
- if self.vae.quant_conv.weight.ndim==5:
+ if self.vae.quant_conv is None or self.vae.quant_conv.weight.ndim==5:
mini_batch_encoder = self.vae.mini_batch_encoder
mini_batch_decoder = self.vae.mini_batch_decoder
video = self.vae.decode(latents)[0]
video = video.clamp(-1, 1)
- video = self.smooth_output(video, mini_batch_encoder, mini_batch_decoder).cpu().clamp(-1, 1)
+ if not self.vae.cache_compression_vae and not self.vae.cache_mag_vae:
+ video = self.smooth_output(video, mini_batch_encoder, mini_batch_decoder).cpu().clamp(-1, 1)
else:
latents = rearrange(latents, "b c f h w -> (b f) c h w")
video = []
@@ -660,32 +891,28 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
video = video.cpu().float().numpy()
return video
- def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
- if isinstance(generator, list):
- image_latents = [
- retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
- for i in range(image.shape[0])
- ]
- image_latents = torch.cat(image_latents, dim=0)
- else:
- image_latents = retrieve_latents(self.vae.encode(image), generator=generator)
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
- image_latents = self.vae.config.scaling_factor * image_latents
+ @property
+ def guidance_rescale(self):
+ return self._guidance_rescale
- return image_latents
+ # 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.
+ @property
+ def do_classifier_free_guidance(self):
+ return self._guidance_scale > 1
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps
- def get_timesteps(self, num_inference_steps, strength, device):
- # get the original timestep using init_timestep
- init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
- t_start = max(num_inference_steps - init_timestep, 0)
- timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
-
- return timesteps, num_inference_steps - t_start
-
- def enable_autocast_float8_transformer(self):
- self.enable_autocast_float8_transformer_flag = True
+ @property
+ def interrupt(self):
+ return self._interrupt
@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
@@ -696,109 +923,167 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
video: Union[torch.FloatTensor] = None,
mask_video: Union[torch.FloatTensor] = None,
masked_video_latents: Union[torch.FloatTensor] = 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,
height: Optional[int] = None,
width: Optional[int] = None,
- strength: float = 1.0,
- eta: float = 0.0,
+ num_inference_steps: Optional[int] = 50,
+ guidance_scale: Optional[float] = 5.0,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ num_images_per_prompt: Optional[int] = 1,
+ eta: Optional[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,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds_2: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds_2: Optional[torch.Tensor] = None,
+ prompt_attention_mask: Optional[torch.Tensor] = None,
+ prompt_attention_mask_2: Optional[torch.Tensor] = None,
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ negative_prompt_attention_mask_2: Optional[torch.Tensor] = None,
output_type: Optional[str] = "latent",
return_dict: bool = True,
- callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
- callback_steps: int = 1,
- clean_caption: bool = True,
- mask_feature: bool = True,
- max_sequence_length: int = 120,
+ callback_on_step_end: Optional[
+ Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+ ] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ guidance_rescale: float = 0.0,
+ original_size: Optional[Tuple[int, int]] = (1024, 1024),
+ target_size: Optional[Tuple[int, int]] = None,
+ crops_coords_top_left: Tuple[int, int] = (0, 0),
clip_image: Image = None,
- clip_apply_ratio: float = 0.50,
+ clip_apply_ratio: float = 0.40,
+ strength: float = 1.0,
+ noise_aug_strength: float = 0.0563,
comfyui_progressbar: bool = False,
- **kwargs,
- ) -> Union[EasyAnimatePipelineOutput, Tuple]:
- """
- Function invoked when calling the pipeline for generation.
+ timesteps: Optional[List[int]] = None,
+ ):
+ r"""
+ The call function to the pipeline for generation with HunyuanDiT.
- Args:
+ Examples:
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.
+ The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+ video_length (`int`, *optional*):
+ Length of the video to be generated in seconds. This parameter influences the number of frames and
+ continuity of generated content.
+ video (`torch.FloatTensor`, *optional*):
+ A tensor representing an input video, which can be modified depending on the prompts provided.
+ mask_video (`torch.FloatTensor`, *optional*):
+ A tensor to specify areas of the video to be masked (omitted from generation).
+ masked_video_latents (`torch.FloatTensor`, *optional*):
+ Latents from masked portions of the video, utilized during image generation.
+ height (`int`, *optional*):
+ The height in pixels of the generated image or video frames.
+ width (`int`, *optional*):
+ The width in pixels of the generated image or video frames.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps. More denoising steps usually lead to a higher quality image but slower
+ inference time. This parameter is modulated by `strength`.
+ guidance_scale (`float`, *optional*, defaults to 5.0):
+ A higher guidance scale value encourages the model to generate images closely linked to the text
+ `prompt` at the expense of lower image quality. Guidance scale is effective when `guidance_scale > 1`.
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 7.0):
- 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.
+ The prompt or prompts to guide what to exclude in image generation. If not defined, you need to
+ provide `negative_prompt_embeds`. This parameter is ignored when not using guidance (`guidance_scale < 1`).
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.
+ A parameter defined in the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies to the
+ [`~schedulers.DDIMScheduler`] and is ignored in other schedulers. It adjusts noise level during the
+ inference process.
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.
- negative_prompt_embeds (`torch.FloatTensor`, *optional*):
- Pre-generated negative text embeddings. For PixArt-Alpha this negative prompt should be "". If not
- provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
- 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`.
+ A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) for setting
+ random seeds which helps in making generation deterministic.
+ latents (`torch.Tensor`, *optional*):
+ A pre-computed latent representation which can be used to guide the generation process.
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+ provided, embeddings are generated from the `prompt` input argument.
+ prompt_embeds_2 (`torch.Tensor`, *optional*):
+ Secondary set of pre-generated text embeddings, useful for advanced prompt weighting.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings, aiding in fine-tuning what should not be represented in the outputs.
+ If not provided, embeddings are generated from the `negative_prompt` argument.
+ negative_prompt_embeds_2 (`torch.Tensor`, *optional*):
+ Secondary set of pre-generated negative text embeddings for further control.
+ prompt_attention_mask (`torch.Tensor`, *optional*):
+ Attention mask guiding the focus of the model on specific parts of the prompt text. Required when using
+ `prompt_embeds`.
+ prompt_attention_mask_2 (`torch.Tensor`, *optional*):
+ Attention mask for the secondary prompt embedding.
+ negative_prompt_attention_mask (`torch.Tensor`, *optional*):
+ Attention mask for the negative prompt, needed when `negative_prompt_embeds` are used.
+ negative_prompt_attention_mask_2 (`torch.Tensor`, *optional*):
+ Attention mask for the secondary negative prompt embedding.
+ output_type (`str`, *optional*, defaults to `"latent"`):
+ The output format of the generated image. Choose between `PIL.Image` and `np.array` to define
+ how you want the results to be formatted.
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.
- mask_feature (`bool` defaults to `True`): If set to `True`, the text embeddings will be masked.
+ If set to `True`, a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] will be returned;
+ otherwise, a tuple containing the generated images and safety flags will be returned.
+ callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+ A callback function (or a list of them) that will be executed at the end of each denoising step,
+ allowing for custom processing during generation.
+ callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
+ Specifies which tensor inputs should be included in the callback function. If not defined, all tensor
+ inputs will be passed, facilitating enhanced logging or monitoring of the generation process.
+ guidance_rescale (`float`, *optional*, defaults to 0.0):
+ Rescale parameter for adjusting noise configuration based on guidance rescale. Based on findings from
+ [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+ original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`):
+ The original dimensions of the image. Used to compute time ids during the generation process.
+ target_size (`Tuple[int, int]`, *optional*):
+ The targeted dimensions of the generated image, also utilized in the time id calculations.
+ crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`):
+ Coordinates defining the top left corner of any cropping, utilized while calculating the time ids.
+ clip_image (`Image`, *optional*):
+ An optional image to assist in the generation process. It may be used as an additional visual cue.
+ clip_apply_ratio (`float`, *optional*, defaults to 0.40):
+ Ratio indicating how much influence the clip image should exert over the generated content.
+ strength (`float`, *optional*, defaults to 1.0):
+ Affects the overall styling or quality of the generated output. Values closer to 1 usually provide direct
+ adherence to prompts.
+ comfyui_progressbar (`bool`, *optional*, defaults to `False`):
+ Enables a progress bar in ComfyUI, providing visual feedback during the generation process.
Examples:
-
+ # Example usage of the function for generating images based on prompts.
+
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
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+ Returns either a structured output containing generated images and their metadata when `return_dict` is
+ `True`, or a simpler tuple, where the first element is a list of generated images and the second
+ element indicates if any of them contain "not-safe-for-work" (NSFW) content.
"""
- # 1. Check inputs. Raise error if not correct
- height = height or self.transformer.config.sample_size * self.vae_scale_factor
- width = width or self.transformer.config.sample_size * self.vae_scale_factor
+
+ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+ callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+ # 0. default height and width
height = int(height // 16 * 16)
width = int(width // 16 * 16)
- # 2. Default height and width to transformer
+ # 1. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ prompt_embeds,
+ negative_prompt_embeds,
+ prompt_attention_mask,
+ negative_prompt_attention_mask,
+ prompt_embeds_2,
+ negative_prompt_embeds_2,
+ prompt_attention_mask_2,
+ negative_prompt_attention_mask_2,
+ callback_on_step_end_tensor_inputs,
+ )
+ self._guidance_scale = guidance_scale
+ self._guidance_rescale = guidance_rescale
+ self._interrupt = False
+
+ # 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
@@ -807,40 +1092,68 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
batch_size = prompt_embeds.shape[0]
device = self._execution_device
-
- # 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
-
+ if self.text_encoder is not None:
+ dtype = self.text_encoder.dtype
+ elif self.text_encoder_2 is not None:
+ dtype = self.text_encoder_2.dtype
+ else:
+ dtype = self.transformer.dtype
+
# 3. Encode input prompt
(
prompt_embeds,
- prompt_attention_mask,
negative_prompt_embeds,
+ prompt_attention_mask,
negative_prompt_attention_mask,
) = self.encode_prompt(
- prompt,
- do_classifier_free_guidance,
- negative_prompt=negative_prompt,
- num_images_per_prompt=num_images_per_prompt,
+ prompt=prompt,
device=device,
+ dtype=dtype,
+ num_images_per_prompt=num_images_per_prompt,
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
+ negative_prompt=negative_prompt,
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,
+ text_encoder_index=0,
)
- 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)
+ if self.tokenizer_2 is not None:
+ (
+ prompt_embeds_2,
+ negative_prompt_embeds_2,
+ prompt_attention_mask_2,
+ negative_prompt_attention_mask_2,
+ ) = self.encode_prompt(
+ prompt=prompt,
+ device=device,
+ dtype=dtype,
+ num_images_per_prompt=num_images_per_prompt,
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
+ negative_prompt=negative_prompt,
+ prompt_embeds=prompt_embeds_2,
+ negative_prompt_embeds=negative_prompt_embeds_2,
+ prompt_attention_mask=prompt_attention_mask_2,
+ negative_prompt_attention_mask=negative_prompt_attention_mask_2,
+ text_encoder_index=1,
+ )
+ else:
+ prompt_embeds_2 = None
+ negative_prompt_embeds_2 = None
+ prompt_attention_mask_2 = None
+ negative_prompt_attention_mask_2 = None
# 4. set timesteps
- self.scheduler.set_timesteps(num_inference_steps, device=device)
+ if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, mu=1)
+ else:
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
timesteps, num_inference_steps = self.get_timesteps(
num_inference_steps=num_inference_steps, strength=strength, device=device
)
+ if comfyui_progressbar:
+ from comfy.utils import ProgressBar
+ pbar = ProgressBar(num_inference_steps + 3)
# at which timestep to set the initial noise (n.b. 50% if strength is 0.5)
latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
# create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise
@@ -857,7 +1170,7 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
# Prepare latent variables
num_channels_latents = self.vae.config.latent_channels
num_channels_transformer = self.transformer.config.in_channels
- return_image_latents = True # num_channels_transformer == 4
+ return_image_latents = num_channels_transformer == num_channels_latents
# 5. Prepare latents.
latents_outputs = self.prepare_latents(
@@ -866,7 +1179,7 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
height,
width,
video_length,
- prompt_embeds.dtype,
+ dtype,
device,
generator,
latents,
@@ -880,91 +1193,153 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
latents, noise, image_latents = latents_outputs
else:
latents, noise = latents_outputs
- latents_dtype = latents.dtype
+ if comfyui_progressbar:
+ pbar.update(1)
+
+ # 6. Prepare clip latents if it needs.
+ if clip_image is not None and self.transformer.enable_clip_in_inpaint:
+ inputs = self.clip_image_processor(images=clip_image, return_tensors="pt")
+ inputs["pixel_values"] = inputs["pixel_values"].to(device, dtype=dtype)
+ if self.transformer.config.get("position_of_clip_embedding", "full") == "full":
+ clip_encoder_hidden_states = self.clip_image_encoder(**inputs).last_hidden_state[:, 1:]
+ clip_encoder_hidden_states_neg = torch.zeros(
+ [
+ batch_size,
+ int(self.clip_image_encoder.config.image_size / self.clip_image_encoder.config.patch_size) ** 2,
+ int(self.clip_image_encoder.config.hidden_size)
+ ]
+ ).to(device, dtype=dtype)
+
+ else:
+ clip_encoder_hidden_states = self.clip_image_encoder(**inputs).image_embeds
+ clip_encoder_hidden_states_neg = torch.zeros([batch_size, 768]).to(device, dtype=dtype)
+
+ clip_attention_mask = torch.ones([batch_size, self.transformer.n_query]).to(device, dtype=dtype)
+ clip_attention_mask_neg = torch.zeros([batch_size, self.transformer.n_query]).to(device, dtype=dtype)
+
+ clip_encoder_hidden_states_input = torch.cat([clip_encoder_hidden_states_neg, clip_encoder_hidden_states]) if self.do_classifier_free_guidance else clip_encoder_hidden_states
+ clip_attention_mask_input = torch.cat([clip_attention_mask_neg, clip_attention_mask]) if self.do_classifier_free_guidance else clip_attention_mask
+
+ elif clip_image is None and num_channels_transformer != num_channels_latents and self.transformer.enable_clip_in_inpaint:
+ if self.transformer.config.get("position_of_clip_embedding", "full") == "full":
+ clip_encoder_hidden_states = torch.zeros(
+ [
+ batch_size,
+ int(self.clip_image_encoder.config.image_size / self.clip_image_encoder.config.patch_size) ** 2,
+ int(self.clip_image_encoder.config.hidden_size)
+ ]
+ ).to(device, dtype=dtype)
+ else:
+ clip_encoder_hidden_states = torch.zeros([batch_size, 768]).to(device, dtype=dtype)
+
+ clip_attention_mask = torch.zeros([batch_size, self.transformer.n_query])
+ clip_attention_mask = clip_attention_mask.to(device, dtype=dtype)
+
+ clip_encoder_hidden_states_input = torch.cat([clip_encoder_hidden_states] * 2) if self.do_classifier_free_guidance else clip_encoder_hidden_states
+ clip_attention_mask_input = torch.cat([clip_attention_mask] * 2) if self.do_classifier_free_guidance else clip_attention_mask
+
+ else:
+ clip_encoder_hidden_states_input = None
+ clip_attention_mask_input = None
+ if comfyui_progressbar:
+ pbar.update(1)
+
+ # 7. Prepare inpaint latents if it needs.
if mask_video is not None:
- # Prepare mask latent variables
- video_length = video.shape[2]
- mask_condition = self.mask_processor.preprocess(rearrange(mask_video, "b c f h w -> (b f) c h w"), height=height, width=width)
- mask_condition = mask_condition.to(dtype=torch.float32)
- mask_condition = rearrange(mask_condition, "(b f) c h w -> b c f h w", f=video_length)
-
- if num_channels_transformer == 12:
- mask_condition_tile = torch.tile(mask_condition, [1, 3, 1, 1, 1])
- if masked_video_latents is None:
- masked_video = init_video * (mask_condition_tile < 0.5) + torch.ones_like(init_video) * (mask_condition_tile > 0.5) * -1
+ if self.transformer.config.get("enable_zero_in_inpaint", True) and (mask_video == 255).all():
+ # Use zero latents if we want to t2v.
+ mask = torch.zeros_like(latents).to(device, dtype)
+ if self.transformer.resize_inpaint_mask_directly:
+ mask_latents = torch.zeros_like(latents)[:, :1].to(device, dtype)
else:
- masked_video = masked_video_latents
-
- mask_latents, masked_video_latents = self.prepare_mask_latents(
- mask_condition_tile,
- masked_video,
- batch_size,
- height,
- width,
- prompt_embeds.dtype,
- device,
- generator,
- do_classifier_free_guidance,
- )
- mask = torch.tile(mask_condition, [1, num_channels_transformer // 3, 1, 1, 1])
- mask = F.interpolate(mask, size=latents.size()[-3:], mode='trilinear', align_corners=True).to(latents.device, latents.dtype)
-
- mask_input = torch.cat([mask_latents] * 2) if do_classifier_free_guidance else mask_latents
+ mask_latents = torch.zeros_like(latents).to(device, dtype)
+ masked_video_latents = torch.zeros_like(latents).to(device, dtype)
+
+ mask_input = torch.cat([mask_latents] * 2) if self.do_classifier_free_guidance else mask_latents
masked_video_latents_input = (
- torch.cat([masked_video_latents] * 2) if do_classifier_free_guidance else masked_video_latents
+ torch.cat([masked_video_latents] * 2) if self.do_classifier_free_guidance else masked_video_latents
)
- inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(latents.dtype)
+ inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(dtype)
else:
- mask = torch.tile(mask_condition, [1, num_channels_transformer, 1, 1, 1])
- mask = F.interpolate(mask, size=latents.size()[-3:], mode='trilinear', align_corners=True).to(latents.device, latents.dtype)
-
- inpaint_latents = None
+ # Prepare mask latent variables
+ video_length = video.shape[2]
+ mask_condition = self.mask_processor.preprocess(rearrange(mask_video, "b c f h w -> (b f) c h w"), height=height, width=width)
+ mask_condition = mask_condition.to(dtype=torch.float32)
+ mask_condition = rearrange(mask_condition, "(b f) c h w -> b c f h w", f=video_length)
+
+ if num_channels_transformer != num_channels_latents:
+ mask_condition_tile = torch.tile(mask_condition, [1, 3, 1, 1, 1])
+ if masked_video_latents is None:
+ masked_video = init_video * (mask_condition_tile < 0.5) + torch.ones_like(init_video) * (mask_condition_tile > 0.5) * -1
+ else:
+ masked_video = masked_video_latents
+
+ if self.transformer.resize_inpaint_mask_directly:
+ _, masked_video_latents = self.prepare_mask_latents(
+ None,
+ masked_video,
+ batch_size,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ self.do_classifier_free_guidance,
+ noise_aug_strength=noise_aug_strength,
+ )
+ mask_latents = resize_mask(1 - mask_condition, masked_video_latents, self.vae.cache_mag_vae)
+ mask_latents = mask_latents.to(device, dtype) * self.vae.config.scaling_factor
+ else:
+ mask_latents, masked_video_latents = self.prepare_mask_latents(
+ mask_condition_tile,
+ masked_video,
+ batch_size,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ self.do_classifier_free_guidance,
+ noise_aug_strength=noise_aug_strength,
+ )
+
+ mask_input = torch.cat([mask_latents] * 2) if self.do_classifier_free_guidance else mask_latents
+ masked_video_latents_input = (
+ torch.cat([masked_video_latents] * 2) if self.do_classifier_free_guidance else masked_video_latents
+ )
+ inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(dtype)
+ else:
+ inpaint_latents = None
+
+ mask = torch.tile(mask_condition, [1, num_channels_latents, 1, 1, 1])
+ mask = F.interpolate(mask, size=latents.size()[-3:], mode='trilinear', align_corners=True).to(device, dtype)
else:
- if num_channels_transformer == 12:
- mask = torch.zeros_like(latents).to(latents.device, latents.dtype)
- masked_video_latents = torch.zeros_like(latents).to(latents.device, latents.dtype)
+ if num_channels_transformer != num_channels_latents:
+ mask = torch.zeros_like(latents).to(device, dtype)
+ if self.transformer.resize_inpaint_mask_directly:
+ mask_latents = torch.zeros_like(latents)[:, :1].to(device, dtype)
+ else:
+ mask_latents = torch.zeros_like(latents).to(device, dtype)
+ masked_video_latents = torch.zeros_like(latents).to(device, dtype)
- mask_input = torch.cat([mask] * 2) if do_classifier_free_guidance else mask
+ mask_input = torch.cat([mask_latents] * 2) if self.do_classifier_free_guidance else mask_latents
masked_video_latents_input = (
- torch.cat([masked_video_latents] * 2) if do_classifier_free_guidance else masked_video_latents
+ torch.cat([masked_video_latents] * 2) if self.do_classifier_free_guidance else masked_video_latents
)
- inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(latents.dtype)
+ inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(dtype)
else:
mask = torch.zeros_like(init_video[:, :1])
- mask = torch.tile(mask, [1, num_channels_transformer, 1, 1, 1])
- mask = F.interpolate(mask, size=latents.size()[-3:], mode='trilinear', align_corners=True).to(latents.device, latents.dtype)
+ mask = torch.tile(mask, [1, num_channels_latents, 1, 1, 1])
+ mask = F.interpolate(mask, size=latents.size()[-3:], mode='trilinear', align_corners=True).to(device, dtype)
inpaint_latents = None
-
- if clip_image is not None:
- inputs = self.clip_image_processor(images=clip_image, return_tensors="pt")
- inputs["pixel_values"] = inputs["pixel_values"].to(latents.device, dtype=latents.dtype)
- clip_encoder_hidden_states = self.clip_image_encoder(**inputs).image_embeds
- clip_encoder_hidden_states_neg = torch.zeros([batch_size, 768]).to(latents.device, dtype=latents.dtype)
-
- clip_attention_mask = torch.ones([batch_size, 8]).to(latents.device, dtype=latents.dtype)
- clip_attention_mask_neg = torch.zeros([batch_size, 8]).to(latents.device, dtype=latents.dtype)
- clip_encoder_hidden_states_input = torch.cat([clip_encoder_hidden_states_neg, clip_encoder_hidden_states]) if do_classifier_free_guidance else clip_encoder_hidden_states
- clip_attention_mask_input = torch.cat([clip_attention_mask_neg, clip_attention_mask]) if do_classifier_free_guidance else clip_attention_mask
-
- elif clip_image is None and num_channels_transformer == 12:
- clip_encoder_hidden_states = torch.zeros([batch_size, 768]).to(latents.device, dtype=latents.dtype)
-
- clip_attention_mask = torch.zeros([batch_size, 8])
- clip_attention_mask = clip_attention_mask.to(latents.device, dtype=latents.dtype)
-
- clip_encoder_hidden_states_input = torch.cat([clip_encoder_hidden_states] * 2) if do_classifier_free_guidance else clip_encoder_hidden_states
- clip_attention_mask_input = torch.cat([clip_attention_mask] * 2) if do_classifier_free_guidance else clip_attention_mask
-
- else:
- clip_encoder_hidden_states_input = None
- clip_attention_mask_input = None
+ if comfyui_progressbar:
+ pbar.update(1)
# Check that sizes of mask, masked image and latents match
- if num_channels_transformer == 12:
- # default case for runwayml/stable-diffusion-inpainting
+ if num_channels_transformer != num_channels_latents:
num_channels_mask = mask_latents.shape[1]
num_channels_masked_image = masked_video_latents.shape[1]
if num_channels_latents + num_channels_mask + num_channels_masked_image != self.transformer.config.in_channels:
@@ -975,45 +1350,89 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of"
" `pipeline.transformer` or your `mask_image` or `image` input."
)
- elif num_channels_transformer != 4:
- raise ValueError(
- f"The transformer {self.transformer.__class__} should have 9 input channels, not {self.transformer.config.in_channels}."
- )
-
- # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+
+ # 8. 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.
+ # 9 create image_rotary_emb, style embedding & time ids
+ grid_height = height // 8 // self.transformer.config.patch_size
+ grid_width = width // 8 // self.transformer.config.patch_size
+ if self.transformer.config.get("time_position_encoding_type", "2d_rope") == "3d_rope":
+ base_size_width = 720 // 8 // self.transformer.config.patch_size
+ base_size_height = 480 // 8 // self.transformer.config.patch_size
+
+ grid_crops_coords = get_resize_crop_region_for_grid(
+ (grid_height, grid_width), base_size_width, base_size_height
+ )
+ image_rotary_emb = get_3d_rotary_pos_embed(
+ self.transformer.config.attention_head_dim, grid_crops_coords, grid_size=(grid_height, grid_width),
+ temporal_size=latents.size(2), use_real=True,
+ )
+ else:
+ base_size = 512 // 8 // self.transformer.config.patch_size
+ grid_crops_coords = get_resize_crop_region_for_grid(
+ (grid_height, grid_width), base_size, base_size
+ )
+ image_rotary_emb = get_2d_rotary_pos_embed(
+ self.transformer.config.attention_head_dim, grid_crops_coords, (grid_height, grid_width)
+ )
+
+ # Get other hunyuan params
+ target_size = target_size or (height, width)
+ add_time_ids = list(original_size + target_size + crops_coords_top_left)
+ add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
+ style = torch.tensor([0], device=device)
+
+ if self.do_classifier_free_guidance:
+ add_time_ids = torch.cat([add_time_ids] * 2, dim=0)
+ style = torch.cat([style] * 2, dim=0)
+
+ # To latents.device
+ add_time_ids = add_time_ids.to(dtype=dtype, device=device).repeat(
+ batch_size * num_images_per_prompt, 1
+ )
+ style = style.to(device=device).repeat(batch_size * num_images_per_prompt)
+
+ # Get other pixart params
added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
- if self.transformer.config.sample_size == 128:
+ if self.transformer.config.get("sample_size", 64) == 128:
resolution = torch.tensor([height, width]).repeat(batch_size * num_images_per_prompt, 1)
aspect_ratio = torch.tensor([float(height / width)]).repeat(batch_size * num_images_per_prompt, 1)
- resolution = resolution.to(dtype=prompt_embeds.dtype, device=device)
- aspect_ratio = aspect_ratio.to(dtype=prompt_embeds.dtype, device=device)
+ resolution = resolution.to(dtype=dtype, device=device)
+ aspect_ratio = aspect_ratio.to(dtype=dtype, device=device)
- if do_classifier_free_guidance:
+ if self.do_classifier_free_guidance:
resolution = torch.cat([resolution, resolution], dim=0)
aspect_ratio = torch.cat([aspect_ratio, aspect_ratio], dim=0)
-
+
added_cond_kwargs = {"resolution": resolution, "aspect_ratio": aspect_ratio}
- gc.collect()
- torch.cuda.empty_cache()
- torch.cuda.ipc_collect()
- if self.enable_autocast_float8_transformer_flag:
- origin_weight_dtype = self.transformer.dtype
- self.transformer = self.transformer.to(torch.float8_e4m3fn)
+ if self.do_classifier_free_guidance:
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+ prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask])
+ if prompt_embeds_2 is not None:
+ prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2])
+ prompt_attention_mask_2 = torch.cat([negative_prompt_attention_mask_2, prompt_attention_mask_2])
+
+ # To latents.device
+ prompt_embeds = prompt_embeds.to(device=device)
+ prompt_attention_mask = prompt_attention_mask.to(device=device)
+ if prompt_embeds_2 is not None:
+ prompt_embeds_2 = prompt_embeds_2.to(device=device)
+ prompt_attention_mask_2 = prompt_attention_mask_2.to(device=device)
# 10. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
- if comfyui_progressbar:
- from comfy.utils import ProgressBar
- pbar = ProgressBar(num_inference_steps)
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)
+ if self.interrupt:
+ continue
+
+ # expand the latents if we are doing classifier free guidance
+ latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+ if hasattr(self.scheduler, "scale_model_input"):
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
if i < len(timesteps) * (1 - clip_apply_ratio) and clip_encoder_hidden_states_input is not None:
clip_encoder_hidden_states_actual_input = torch.zeros_like(clip_encoder_hidden_states_input)
@@ -1021,74 +1440,83 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
else:
clip_encoder_hidden_states_actual_input = clip_encoder_hidden_states_input
clip_attention_mask_actual_input = clip_attention_mask_input
+
+ # expand scalar t to 1-D tensor to match the 1st dim of latent_model_input
+ t_expand = torch.tensor([t] * latent_model_input.shape[0], device=device).to(
+ dtype=latent_model_input.dtype
+ )
- 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])
-
- # predict noise model_output
+ # predict the noise residual
noise_pred = self.transformer(
latent_model_input,
+ t_expand,
encoder_hidden_states=prompt_embeds,
- encoder_attention_mask=prompt_attention_mask,
- timestep=current_timestep,
- added_cond_kwargs=added_cond_kwargs,
+ text_embedding_mask=prompt_attention_mask,
+ encoder_hidden_states_t5=prompt_embeds_2,
+ text_embedding_mask_t5=prompt_attention_mask_2,
+ image_meta_size=add_time_ids,
+ style=style,
+ image_rotary_emb=image_rotary_emb,
inpaint_latents=inpaint_latents,
clip_encoder_hidden_states=clip_encoder_hidden_states_actual_input,
clip_attention_mask=clip_attention_mask_actual_input,
+ added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)[0]
+ if noise_pred.size()[1] != self.vae.config.latent_channels:
+ noise_pred, _ = noise_pred.chunk(2, dim=1)
# perform guidance
- if do_classifier_free_guidance:
+ if self.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
- noise_pred = noise_pred.chunk(2, dim=1)[0]
+ if self.do_classifier_free_guidance and guidance_rescale > 0.0:
+ # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+ noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
- # compute previous image: x_t -> x_t-1
+ # compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
- if num_channels_transformer == 4:
+ if num_channels_transformer == num_channels_latents:
init_latents_proper = image_latents
init_mask = mask
if i < len(timesteps) - 1:
noise_timestep = timesteps[i + 1]
- init_latents_proper = self.scheduler.add_noise(
- init_latents_proper, noise, torch.tensor([noise_timestep])
- )
-
+ if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+ init_latents_proper = self.scheduler.scale_noise(
+ init_latents_proper, torch.tensor([noise_timestep], noise)
+ )
+ else:
+ init_latents_proper = self.scheduler.add_noise(
+ init_latents_proper, noise, torch.tensor([noise_timestep])
+ )
+
latents = (1 - init_mask) * init_latents_proper + init_mask * latents
- # call the callback, if provided
+ if callback_on_step_end is not None:
+ callback_kwargs = {}
+ for k in callback_on_step_end_tensor_inputs:
+ callback_kwargs[k] = locals()[k]
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+ latents = callback_outputs.pop("latents", latents)
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+ negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+ prompt_embeds_2 = callback_outputs.pop("prompt_embeds_2", prompt_embeds_2)
+ negative_prompt_embeds_2 = callback_outputs.pop(
+ "negative_prompt_embeds_2", negative_prompt_embeds_2
+ )
+
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)
+
+ if XLA_AVAILABLE:
+ xm.mark_step()
if comfyui_progressbar:
pbar.update(1)
- if self.enable_autocast_float8_transformer_flag:
- self.transformer = self.transformer.to("cpu", origin_weight_dtype)
-
- gc.collect()
- torch.cuda.empty_cache()
- torch.cuda.ipc_collect()
-
# Post-processing
video = self.decode_latents(latents)
@@ -1096,7 +1524,10 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
if output_type == "latent":
video = torch.from_numpy(video)
+ # Offload all models
+ self.maybe_free_model_hooks()
+
if not return_dict:
return video
- return EasyAnimatePipelineOutput(videos=video)
\ No newline at end of file
+ return EasyAnimatePipelineOutput(frames=video)
\ No newline at end of file