content: # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
  
  import tqdm
  
  import inspect
  from typing import Callable, List, Optional, Union
  from dataclasses import dataclass
  
  import numpy as np
  import torch
  
  from diffusers.utils import is_accelerate_available
  from packaging import version
  from transformers import CLIPTextModel, CLIPTokenizer
  import torchvision.transforms.functional as TF
  
  from diffusers.configuration_utils import FrozenDict
  from diffusers.models import AutoencoderKL
  from diffusers import DiffusionPipeline
  from diffusers.schedulers import (
      DDIMScheduler,
      DPMSolverMultistepScheduler,
      EulerAncestralDiscreteScheduler,
      EulerDiscreteScheduler,
      LMSDiscreteScheduler,
      PNDMScheduler,
  )
  from diffusers.utils import deprecate, logging, BaseOutput
  
  from einops import rearrange
  
  from ..models.unet import UNet3DConditionModel
  from torchvision.transforms import InterpolationMode
  
  import comfy.utils
  
  logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  @dataclass
  class TuneAVideoPipelineOutput(BaseOutput):
      videos: Union[torch.Tensor, np.ndarray]
  
  
  class TuneAVideoPipeline(DiffusionPipeline):
      _optional_components = []
  
      def __init__(
          self,
          vae: AutoencoderKL,
          text_encoder: CLIPTextModel,
          tokenizer: CLIPTokenizer,
          unet: UNet3DConditionModel,
          
          scheduler: Union[
              DDIMScheduler,
              PNDMScheduler,
              LMSDiscreteScheduler,
              EulerDiscreteScheduler,
              EulerAncestralDiscreteScheduler,
              DPMSolverMultistepScheduler,
          ],
          ref_unet = None,
          feature_extractor=None,
          image_encoder=None
      ):
          super().__init__()
          self.ref_unet = ref_unet
          self.feature_extractor = feature_extractor
          self.image_encoder = image_encoder
          
          if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
              deprecation_message = (
                  f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
                  f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
                  "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
                  " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
                  " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
                  " file"
              )
              deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
              new_config = dict(scheduler.config)
              new_config["steps_offset"] = 1
              scheduler._internal_dict = FrozenDict(new_config)
  
          if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
              deprecation_message = (
                  f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
                  " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
                  " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
                  " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
                  " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
              )
              deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
              new_config = dict(scheduler.config)
              new_config["clip_sample"] = False
              scheduler._internal_dict = FrozenDict(new_config)
  
          is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
              version.parse(unet.config._diffusers_version).base_version
          ) < version.parse("0.9.0.dev0")
          is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
          if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
              deprecation_message = (
                  "The configuration file of the unet has set the default `sample_size` to smaller than"
                  " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
                  " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
                  " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
                  " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
                  " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
                  " in the config might lead to incorrect results in future versions. If you have downloaded this"
                  " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
                  " the `unet/config.json` file"
              )
              deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
              new_config = dict(unet.config)
              new_config["sample_size"] = 64
              unet._internal_dict = FrozenDict(new_config)
  
          self.register_modules(
              vae=vae,
              text_encoder=text_encoder,
              tokenizer=tokenizer,
              unet=unet,
              scheduler=scheduler,
          )
          self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
  
      def enable_vae_slicing(self):
          self.vae.enable_slicing()
  
      def disable_vae_slicing(self):
          self.vae.disable_slicing()
  
      def enable_sequential_cpu_offload(self, gpu_id=0):
          if is_accelerate_available():
              from accelerate import cpu_offload
          else:
              raise ImportError("Please install accelerate via `pip install accelerate`")
  
          device = torch.device(f"cuda:{gpu_id}")
  
          for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
              if cpu_offloaded_model is not None:
                  cpu_offload(cpu_offloaded_model, device)
  
  
      @property
      def _execution_device(self):
          if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
              return self.device
          for module in self.unet.modules():
              if (
                  hasattr(module, "_hf_hook")
                  and hasattr(module._hf_hook, "execution_device")
                  and module._hf_hook.execution_device is not None
              ):
                  return torch.device(module._hf_hook.execution_device)
          return self.device
  
      def _encode_image(self, image_pil, device, num_images_per_prompt, do_classifier_free_guidance, img_proj=None):
          dtype = next(self.image_encoder.parameters()).dtype
  
          # image_pt = self.feature_extractor(images=image_pil, return_tensors="pt").pixel_values
          # image_pt = image_pt.to(device=device, dtype=dtype)
          # image_embeddings = self.image_encoder(image_pt).image_embeds
          # image_embeddings = image_embeddings.unsqueeze(1)
  
          # # image encoding
          clip_image_mean = torch.as_tensor(self.feature_extractor.image_mean)[:,None,None].to(device, dtype=torch.float32)
          clip_image_std = torch.as_tensor(self.feature_extractor.image_std)[:,None,None].to(device, dtype=torch.float32)
          imgs_in_proc = TF.resize(image_pil, (self.feature_extractor.crop_size['height'], self.feature_extractor.crop_size['width']), interpolation=InterpolationMode.BICUBIC)
          # do the normalization in float32 to preserve precision
          imgs_in_proc = ((imgs_in_proc.float() - clip_image_mean) / clip_image_std).to(dtype)        
          if img_proj is None:
              # (B*Nv, 1, 768)
              image_embeddings = self.image_encoder(imgs_in_proc).image_embeds.unsqueeze(1)
              # duplicate image embeddings for each generation per prompt, using mps friendly method
              # Note: repeat differently from official pipelines
              # B1B2B3B4 -> B1B2B3B4B1B2B3B4
              bs_embed, seq_len, _ = image_embeddings.shape
              image_embeddings = image_embeddings.repeat(num_images_per_prompt, 1, 1)
              if do_classifier_free_guidance:
                  negative_prompt_embeds = torch.zeros_like(image_embeddings)
  
                  # For classifier free guidance, we need to do two forward passes.
                  # Here we concatenate the unconditional and text embeddings into a single batch
                  # to avoid doing two forward passes
                  image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings])
          else:
              if do_classifier_free_guidance:
                  negative_image_proc = torch.zeros_like(imgs_in_proc)
  
                  # For classifier free guidance, we need to do two forward passes.
                  # Here we concatenate the unconditional and text embeddings into a single batch
                  # to avoid doing two forward passes
                  imgs_in_proc = torch.cat([negative_image_proc, imgs_in_proc])
              
              image_embeds = self.image_encoder(imgs_in_proc, output_hidden_states=True).hidden_states[-2]
              image_embeddings = img_proj(image_embeds)
  
              # image_embeddings_unet = rearrange(image_embeddings_unet, 'B Nv d c -> (B Nv) d c')
          
          # image_pt = torch.stack([TF.to_tensor(img) for img in image_pil], dim=0).to(device)
          # image_pil = image_pil * 2.0 - 1.0
          image_latents = self.vae.encode(image_pil* 2.0 - 1.0).latent_dist.mode() * self.vae.config.scaling_factor
  
          # Note: repeat differently from official pipelines
          # B1B2B3B4 -> B1B2B3B4B1B2B3B4        
          image_latents = image_latents.repeat(num_images_per_prompt, 1, 1, 1)
  
          # if do_classifier_free_guidance:
          #     image_latents = torch.cat([torch.zeros_like(image_latents), image_latents])
  
          return image_embeddings, image_latents
  
      def _encode_prompt(self, prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt):
          batch_size = len(prompt) if isinstance(prompt, list) else 1
  
          text_inputs = self.tokenizer(
              prompt,
              padding="max_length",
              max_length=self.tokenizer.model_max_length,
              truncation=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[:, self.tokenizer.model_max_length - 1 : -1])
              logger.warning(
                  "The following part of your input was truncated because CLIP can only handle sequences up to"
                  f" {self.tokenizer.model_max_length} tokens: {removed_text}"
              )
  
          if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
              attention_mask = text_inputs.attention_mask.to(device)
          else:
              attention_mask = None
  
          text_embeddings = self.text_encoder(
              text_input_ids.to(device),
              attention_mask=attention_mask,
          )
          text_embeddings = text_embeddings[0]
  
          # duplicate text embeddings for each generation per prompt, using mps friendly method
          bs_embed, seq_len, _ = text_embeddings.shape
          text_embeddings = text_embeddings.repeat(1, num_videos_per_prompt, 1)
          text_embeddings = text_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1)
  
          # get unconditional embeddings for classifier free guidance
          if do_classifier_free_guidance:
              uncond_tokens: List[str]
              if negative_prompt is None:
                  uncond_tokens = [""] * batch_size
              elif 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 = text_input_ids.shape[-1]
              uncond_input = self.tokenizer(
                  uncond_tokens,
                  padding="max_length",
                  max_length=max_length,
                  truncation=True,
                  return_tensors="pt",
              )
  
              if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
                  attention_mask = uncond_input.attention_mask.to(device)
              else:
                  attention_mask = None
  
              uncond_embeddings = self.text_encoder(
                  uncond_input.input_ids.to(device),
                  attention_mask=attention_mask,
              )
              uncond_embeddings = uncond_embeddings[0]
  
              # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
              seq_len = uncond_embeddings.shape[1]
              uncond_embeddings = uncond_embeddings.repeat(1, num_videos_per_prompt, 1)
              uncond_embeddings = uncond_embeddings.view(batch_size * num_videos_per_prompt, seq_len, -1)
  
              # For classifier free guidance, we need to do two forward passes.
              # Here we concatenate the unconditional and text embeddings into a single batch
              # to avoid doing two forward passes
              text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
  
          return text_embeddings
  
      def decode_latents(self, latents):
          video_length = latents.shape[2]
          latents = 1 / 0.18215 * latents
          latents = rearrange(latents, "b c f h w -> (b f) c h w")
          video = self.vae.decode(latents).sample
          video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
          video = (video / 2 + 0.5).clamp(0, 1)
          # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
          video = video.cpu().float().numpy()
          return video
  
      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
  
      def check_inputs(self, prompt, height, width, callback_steps):
          if 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 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)}."
              )
  
      def prepare_latents(self, batch_size, num_channels_latents, video_length, height, width, dtype, device, generator, latents=None):
          shape = (batch_size, num_channels_latents, video_length, height // self.vae_scale_factor, width // self.vae_scale_factor)
          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:
              rand_device = "cpu" if device.type == "mps" else device
  
              if isinstance(generator, list):
                  shape = (1,) + shape[1:]
                  latents = [
                      torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype)
                      for i in range(batch_size)
                  ]
                  latents = torch.cat(latents, dim=0).to(device)
              else:
                  latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device)
          else:
              if latents.shape != shape:
                  raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
              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]],
          image: Union[str, List[str]],
          video_length: Optional[int],
          height: Optional[int] = None,
          width: Optional[int] = None,
          num_inference_steps: int = 50,
          guidance_scale: float = 7.5,
          negative_prompt: Optional[Union[str, List[str]]] = None,
          num_videos_per_prompt: Optional[int] = 1,
          eta: float = 0.0,
          generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
          latents: Optional[torch.FloatTensor] = None,
          output_type: Optional[str] = "tensor",
          return_dict: bool = True,
          callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
          callback_steps: Optional[int] = 1,
          camera_matrixs = None,
          class_labels = None,
          prompt_ids = None,
          unet_condition_type = None,
          pose_guider = None,
          pose_image = None,
          img_proj=None,
          use_noise=True,
          use_shifted_noise=False,
          rescale = 0.7,
          **kwargs,
      ):
          # Default height and width to unet
          height = height or self.unet.config.sample_size * self.vae_scale_factor
          width = width or self.unet.config.sample_size * self.vae_scale_factor
  
          # Check inputs. Raise error if not correct
          self.check_inputs(prompt, height, width, callback_steps)
          if isinstance(image, list):
              batch_size = len(image)
          else:
              batch_size = image.shape[0]
          # assert batch_size >= video_length and batch_size % video_length == 0
          # Define call parameters
          batch_size = 1 if isinstance(prompt, str) else len(prompt)
          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
          
          # 3. Encode input image
          # if isinstance(image, list):
          #     image_pil = image
          # elif isinstance(image, torch.Tensor):
          #     image_pil = [TF.to_pil_image(image[i]) for i in range(image.shape[0])]
          # encode input reference image
          image_embeddings, image_latents = self._encode_image(image, device, num_videos_per_prompt, do_classifier_free_guidance, img_proj=img_proj) #torch.Size([64, 1, 768]) torch.Size([64, 4, 32, 32])
          image_latents = rearrange(image_latents, "(b f) c h w -> b c f h w", f=1) #torch.Size([64, 4, 1, 32, 32])
  
          # Encode input prompt_id
          # encoder_hidden_states = self.text_encoder(prompt_ids)[0] #torch.Size([32, 77, 768])
  
          # Encode input prompt
          text_embeddings = self._encode_prompt( #torch.Size([64, 77, 768])
              prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt
          )
  
          # Prepare timesteps
          self.scheduler.set_timesteps(num_inference_steps, device=device)
          timesteps = self.scheduler.timesteps
  
          # Prepare latent variables
          num_channels_latents = self.unet.in_channels
          latents = self.prepare_latents( #torch.Size([32, 4, 4, 32, 32])
              batch_size * num_videos_per_prompt,
              num_channels_latents,
              video_length,
              height,
              width,
              text_embeddings.dtype,
              device,
              generator,
              latents,
          )
          latents_dtype = latents.dtype
          # import ipdb
          # ipdb.set_trace()
          # Prepare extra step kwargs.
          extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
          # prepare camera_matrix
          if camera_matrixs is not None:
              camera_matrixs = torch.cat([camera_matrixs] * 2) if do_classifier_free_guidance else camera_matrixs #(64, 4, 12)
          # Denoising loop
          comfy_pbar = comfy.utils.ProgressBar(num_inference_steps)
          
          num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
          if pose_guider is not None:
              if len(pose_image.shape) == 5:
                  pose_embeds = pose_guider(rearrange(pose_image, "b f c h w -> (b f) c h w"))
                  pose_embeds = rearrange(pose_embeds, "(b f) c h w-> b c f h w ", f=video_length)
              else:
                  pose_embeds = pose_guider(pose_image).unsqueeze(0)
              pose_embeds = torch.cat([pose_embeds]*2, dim=0)
              # import ipdb
              # ipdb.set_trace()
          with self.progress_bar(total=num_inference_steps) as progress_bar:
              for i, t in enumerate(tqdm.tqdm(timesteps)):
                  # expand the latents if we are doing classifier free guidance
                  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 pose_guider is not None:
                      latent_model_input = latent_model_input + pose_embeds
  
                  noise_cond = torch.randn_like(image_latents)
                  if use_noise:
                      cond_latents = self.scheduler.add_noise(image_latents, noise_cond, t)
                  else:
                      cond_latents = image_latents
                  cond_latent_model_input = torch.cat([cond_latents] * 2) if do_classifier_free_guidance else cond_latents
                  cond_latent_model_input = self.scheduler.scale_model_input(cond_latent_model_input, t)
  
                  # predict the noise residual
                  # ref text condition
                  ref_dict = {}
                  if self.ref_unet is not None:
                      noise_pred_cond = self.ref_unet(
                          cond_latent_model_input,    #torch.Size([64, 4, 1, 32, 32])
                          t,             #torch.Size([32])
                          encoder_hidden_states=text_embeddings.to(torch.float32), #torch.Size([64, 77, 768])
                          cross_attention_kwargs=dict(mode="w", ref_dict=ref_dict)
                      ).sample.to(dtype=latents_dtype)
                  
                  # if torch.isnan(noise_pred_cond).any():
                  #     ipdb.set_trace()
                  # Predict the noise residual and compute loss
                  # model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, camera_matrixs).sample
                  # unet
                  #text condition for unet 
                  text_embeddings_unet = text_embeddings.unsqueeze(1).repeat(1,latents.shape[2],1,1)
                  text_embeddings_unet = rearrange(text_embeddings_unet, 'B Nv d c -> (B Nv) d c')
                  #image condition for unet
                  image_embeddings_unet = image_embeddings.unsqueeze(1).repeat(1,latents.shape[2],1, 1)
                  image_embeddings_unet = rearrange(image_embeddings_unet, 'B Nv d c -> (B Nv) d c')
  
                  if unet_condition_type == 'text':
                      encoder_hidden_states_unet_cond = text_embeddings_unet  
                  elif unet_condition_type == 'image':
                      encoder_hidden_states_unet_cond = image_embeddings_unet
                  else:
                      raise('need unet_condition_type')
                  
                  if self.ref_unet is not None:
                      noise_pred = self.unet(
                          latent_model_input.to(torch.float32),         #torch.Size([64, 4, 4, 32, 32])
                          t,
                          encoder_hidden_states=encoder_hidden_states_unet_cond.to(torch.float32), 
                          camera_matrixs=camera_matrixs.to(torch.float32),      #torch.Size([64, 4, 12])
                          cross_attention_kwargs=dict(mode="r", ref_dict=ref_dict, is_cfg_guidance=do_classifier_free_guidance)
                          # cross_attention_kwargs=dict(mode="n", ref_dict=ref_dict, is_cfg_guidance=do_classifier_free_guidance)
                      ).sample.to(dtype=latents_dtype)
                  else:
                      noise_pred = self.unet(
                          latent_model_input.to(torch.float32),         #torch.Size([64, 4, 4, 32, 32])
                          t,
                          encoder_hidden_states=encoder_hidden_states_unet_cond.to(torch.float32), 
                          camera_matrixs=camera_matrixs.to(torch.float32),      #torch.Size([64, 4, 12])
                          # cross_attention_kwargs=dict(mode="r", ref_dict=ref_dict, is_cfg_guidance=do_classifier_free_guidance)
                          cross_attention_kwargs=dict(mode="n", ref_dict=ref_dict, is_cfg_guidance=do_classifier_free_guidance)
                      ).sample.to(dtype=latents_dtype)
                  # perform guidance
                  if do_classifier_free_guidance:
                      noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                      if use_shifted_noise:
                          # Apply regular classifier-free guidance.
                          cfg = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                          # Calculate standard deviations.
                          std_pos = noise_pred_text.std([1,2,3], keepdim=True)
                          std_cfg = cfg.std([1,2,3], keepdim=True)
                          # Apply guidance rescale with fused operations.
                          factor = std_pos / std_cfg
                          factor = rescale * factor + (1 - rescale)
                          noise_pred = cfg * factor
                      else:
                          noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                      # noise_pred_uncond_, noise_pred_text_ = noise_pred_cond.chunk(2)
                      # noise_pred_cond = noise_pred_uncond_ + guidance_scale * (noise_pred_text_ - noise_pred_uncond_)
  
                  # compute the previous noisy sample x_t -> x_t-1
                  noise_pred = rearrange(noise_pred, "(b f) c h w -> b c f h w", f=video_length)
                  latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
                  # noise_pred_cond = rearrange(noise_pred_cond, "(b f) c h w -> b c f h w", f=1)
                  # cond_latents = self.scheduler.step(noise_pred_cond, t, cond_latents, **extra_step_kwargs).prev_sample
  
                  # 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:
                          callback(i, t, latents)
                          
                  comfy_pbar.update_absolute(i + 1)
  
          # Post-processing
          video = self.decode_latents(latents)
  
          # Convert to tensor
          if output_type == "tensor":
              video = torch.from_numpy(video)
  
          if not return_dict:
              return video
  
          return TuneAVideoPipelineOutput(videos=video)
  