



"""simple docstring"""
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    UNetaDConditionModel,
    VideoToVideoSDPipeline,
)
from diffusers.utils import floats_tensor, is_xformers_available, skip_mps
from diffusers.utils.testing_utils import enable_full_determinism, slow, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


@skip_mps
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  unittest.TestCase   ):
    lowercase__      =   VideoToVideoSDPipeline
    lowercase__      =   TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"video"}   ) - {"image", "width", "height"}
    lowercase__      =   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"video"}   ) - {"image"}
    lowercase__      =   PipelineTesterMixin.required_optional_params - {"latents"}
    lowercase__      =   False

    # No `output_type`.
    lowercase__      =   frozenset(
        [
            "num_inference_steps",
            "generator",
            "latents",
            "return_dict",
            "callback",
            "callback_steps",
        ]   )

    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      torch.manual_seed(0)
      lowercase_								= UNetaDConditionModel(
          block_out_channels=(3_2, 6_4, 6_4, 6_4) ,  layers_per_block=2 ,  sample_size=3_2 ,  in_channels=4 ,  out_channels=4 ,  down_block_types=("""CrossAttnDownBlock3D""", """CrossAttnDownBlock3D""", """CrossAttnDownBlock3D""", """DownBlock3D""") ,  up_block_types=("""UpBlock3D""", """CrossAttnUpBlock3D""", """CrossAttnUpBlock3D""", """CrossAttnUpBlock3D""") ,  cross_attention_dim=3_2 ,  attention_head_dim=4 ,  )
      lowercase_								= DDIMScheduler(
          beta_start=0.00_085 ,  beta_end=0.012 ,  beta_schedule="""scaled_linear""" ,  clip_sample=lowerCAmelCase_ ,  set_alpha_to_one=lowerCAmelCase_ ,  )
      torch.manual_seed(0)
      lowercase_								= AutoencoderKL(
          block_out_channels=[3_2, 6_4] ,  in_channels=3 ,  out_channels=3 ,  down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""] ,  up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""] ,  latent_channels=4 ,  sample_size=1_2_8 ,  )
      torch.manual_seed(0)
      lowercase_								= CLIPTextConfig(
          bos_token_id=0 ,  eos_token_id=2 ,  hidden_size=3_2 ,  intermediate_size=3_7 ,  layer_norm_eps=1E-05 ,  num_attention_heads=4 ,  num_hidden_layers=5 ,  pad_token_id=1 ,  vocab_size=1_0_0_0 ,  hidden_act="""gelu""" ,  projection_dim=5_1_2 ,  )
      lowercase_								= CLIPTextModel(lowerCAmelCase_)
      lowercase_								= CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""")

      lowercase_								= {
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
      }
      return components

    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Optional[Any]=0):







      """simple docstring"""







      lowercase_								= floats_tensor((1, 3, 3, 3_2, 3_2) ,  rng=random.Random(lowerCAmelCase_)).to(lowerCAmelCase_)

      if str(lowerCAmelCase_).startswith("""mps"""):
        lowercase_								= torch.manual_seed(lowerCAmelCase_)
      else:
        lowercase_								= torch.Generator(device=lowerCAmelCase_).manual_seed(lowerCAmelCase_)
      lowercase_								= {
          """prompt""": """A painting of a squirrel eating a burger""",
          """video""": video,
          """generator""": generator,
          """num_inference_steps""": 2,
          """guidance_scale""": 6.0,
          """output_type""": """pt""",
      }
      return inputs

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= """cpu"""  # ensure determinism for the device-dependent torch.Generator
      lowercase_								= self.get_dummy_components()
      lowercase_								= VideoToVideoSDPipeline(**lowerCAmelCase_)
      lowercase_								= sd_pipe.to(lowerCAmelCase_)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_)

      lowercase_								= self.get_dummy_inputs(lowerCAmelCase_)
      lowercase_								= """np"""
      lowercase_								= sd_pipe(**lowerCAmelCase_).frames
      lowercase_								= frames[0][-3:, -3:, -1]

      assert frames[0].shape == (3_2, 3_2, 3)
      lowercase_								= np.array([1_0_6, 1_1_7, 1_1_3, 1_7_4, 1_3_7, 1_1_2, 1_4_8, 1_5_1, 1_3_1])

      assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2

    @unittest.skipIf(
        torch_device != """cuda""" or not is_xformers_available() ,  reason="""XFormers attention is only available with CUDA and `xformers` installed""" ,  )
    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowerCAmelCase_ ,  expected_max_diff=5E-3)

    @unittest.skip(reason="""Batching needs to be properly figured out first for this pipeline.""")
    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      pass

    @unittest.skip(reason="""Batching needs to be properly figured out first for this pipeline.""")
    def 					_UpperCAmelCase	(				self			:				Optional[int]):







      """simple docstring"""







      pass

    @unittest.skip(reason="""`num_images_per_prompt` argument is not supported for this pipeline.""")
    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      pass




    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      return super().test_progress_bar()







@slow
@skip_mps
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):

    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      lowercase_								= VideoToVideoSDPipeline.from_pretrained("""cerspense/zeroscope_v2_XL""" ,  torch_dtype=torch.floataa)
      pipe.enable_model_cpu_offload()

      # 10 frames
      lowercase_								= torch.Generator(device="""cpu""").manual_seed(0)
      lowercase_								= torch.randn((1, 1_0, 3, 1_0_2_4, 5_7_6) ,  generator=lowerCAmelCase_)
      lowercase_								= video.to("""cuda""")

      lowercase_								= """Spiderman is surfing"""

      lowercase_								= pipe(lowerCAmelCase_ ,  video=lowerCAmelCase_ ,  generator=lowerCAmelCase_ ,  num_inference_steps=3 ,  output_type="""pt""").frames

      lowercase_								= np.array([-1.0_458_984, -1.1_279_297, -0.9_663_086, -0.91_503_906, -0.75_097_656])
      assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array).sum() < 1E-2






"""simple docstring"""
from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin


class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  __UpperCAmelCase      ,  __UpperCAmelCase   ):
    lowercase__      =   [r"h\.\d+\.attn\.bias", r"h\.\d+\.attn\.masked_bias"]

    @register_to_config
    def __init__(				self			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Optional[int] = None ,  lowerCAmelCase_			:				int = 5_0_2_5_7 ,  lowerCAmelCase_			:				int = 1_0_2_4 ,  lowerCAmelCase_			:				int = 7_6_8 ,  lowerCAmelCase_			:				int = 1_2 ,  lowerCAmelCase_			:				int = 1_2 ,  lowerCAmelCase_			:				Optional[int] = None ,  lowerCAmelCase_			:				str = "gelu_new" ,  lowerCAmelCase_			:				float = 0.1 ,  lowerCAmelCase_			:				float = 0.1 ,  lowerCAmelCase_			:				float = 0.1 ,  lowerCAmelCase_			:				float = 1E-5 ,  lowerCAmelCase_			:				float = 0.02 ,  lowerCAmelCase_			:				bool = True ,  lowerCAmelCase_			:				bool = True ,  lowerCAmelCase_			:				bool = False ,  lowerCAmelCase_			:				bool = False ,  ):







      """simple docstring"""







      super().__init__()

      lowercase_								= prefix_length

      if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
        raise ValueError(
            F'''`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'''
            F''' `n_embd`: {n_embd} are not equal.''')

      lowercase_								= prefix_inner_dim
      lowercase_								= prefix_hidden_dim

      lowercase_								= (
          nn.Linear(self.prefix_inner_dim ,  self.prefix_hidden_dim)
          if self.prefix_hidden_dim is not None
          else nn.Identity()
      )
      lowercase_								= (
          nn.Linear(self.prefix_hidden_dim ,  lowerCAmelCase_) if self.prefix_hidden_dim is not None else nn.Identity()
      )

      lowercase_								= GPTaConfig(
          vocab_size=lowerCAmelCase_ ,  n_positions=lowerCAmelCase_ ,  n_embd=lowerCAmelCase_ ,  n_layer=lowerCAmelCase_ ,  n_head=lowerCAmelCase_ ,  n_inner=lowerCAmelCase_ ,  activation_function=lowerCAmelCase_ ,  resid_pdrop=lowerCAmelCase_ ,  embd_pdrop=lowerCAmelCase_ ,  attn_pdrop=lowerCAmelCase_ ,  layer_norm_epsilon=lowerCAmelCase_ ,  initializer_range=lowerCAmelCase_ ,  scale_attn_weights=lowerCAmelCase_ ,  use_cache=lowerCAmelCase_ ,  scale_attn_by_inverse_layer_idx=lowerCAmelCase_ ,  reorder_and_upcast_attn=lowerCAmelCase_ ,  )
      lowercase_								= GPTaLMHeadModel(lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				torch.Tensor ,  lowerCAmelCase_			:				torch.Tensor ,  lowerCAmelCase_			:				Optional[torch.Tensor] = None ,  lowerCAmelCase_			:				Optional[torch.Tensor] = None ,  ):







      """simple docstring"""







      lowercase_								= self.transformer.transformer.wte(lowerCAmelCase_)
      lowercase_								= self.encode_prefix(lowerCAmelCase_)
      lowercase_								= self.decode_prefix(lowerCAmelCase_)
      lowercase_								= torch.cat((prefix_embeds, embedding_text) ,  dim=1)

      if labels is not None:
        lowercase_								= self.get_dummy_token(input_ids.shape[0] ,  input_ids.device)
        lowercase_								= torch.cat((dummy_token, input_ids) ,  dim=1)
      lowercase_								= self.transformer(inputs_embeds=lowerCAmelCase_ ,  labels=lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_)
      if self.prefix_hidden_dim is not None:
        return out, hidden
      else:
        return out

    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				torch.device):







      """simple docstring"""







      return torch.zeros(lowerCAmelCase_ ,  self.prefix_length ,  dtype=torch.intaa ,  device=lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				List[str] ,  lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      return self.encode_prefix(lowerCAmelCase_)

    @torch.no_grad()
    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				str ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				Union[str, Any]):







      """simple docstring"""







      lowercase_								= torch.split(lowerCAmelCase_ ,  1 ,  dim=0)
      lowercase_								= []
      lowercase_								= []
      for feature in features:
        lowercase_								= self.decode_prefix(feature.to(lowerCAmelCase_))  # back to the clip feature
        # Only support beam search for now
        lowercase_					,    lowercase_								= self.generate_beam(
            input_embeds=lowerCAmelCase_ ,  device=lowerCAmelCase_ ,  eos_token_id=lowerCAmelCase_)
        generated_tokens.append(output_tokens[0])
        generated_seq_lengths.append(seq_lengths[0])
      lowercase_								= torch.stack(lowerCAmelCase_)
      lowercase_								= torch.stack(lowerCAmelCase_)
      return generated_tokens, generated_seq_lengths




    @torch.no_grad()
    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				List[str]=None ,  lowerCAmelCase_			:				Tuple=None ,  lowerCAmelCase_			:				str=None ,  lowerCAmelCase_			:				int = 5 ,  lowerCAmelCase_			:				int = 6_7 ,  lowerCAmelCase_			:				float = 1.0 ,  lowerCAmelCase_			:				Optional[int] = None ,  ):







      """simple docstring"""







      lowercase_								= eos_token_id
      lowercase_								= None
      lowercase_								= None
      lowercase_								= torch.ones(lowerCAmelCase_ ,  device=lowerCAmelCase_ ,  dtype=torch.int)
      lowercase_								= torch.zeros(lowerCAmelCase_ ,  device=lowerCAmelCase_ ,  dtype=torch.bool)

      if input_embeds is not None:
        lowercase_								= input_embeds
      else:
        lowercase_								= self.transformer.transformer.wte(lowerCAmelCase_)

      for i in range(lowerCAmelCase_):
        lowercase_								= self.transformer(inputs_embeds=lowerCAmelCase_)
        lowercase_								= outputs.logits
        lowercase_								= logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
        lowercase_								= logits.softmax(-1).log()

        if scores is None:
          lowercase_					,    lowercase_								= logits.topk(lowerCAmelCase_ ,  -1)
          lowercase_								= generated.expand(lowerCAmelCase_ ,  *generated.shape[1:])
          lowercase_					,    lowercase_								= next_tokens.permute(1 ,  0), scores.squeeze(0)
          if tokens is None:
            lowercase_								= next_tokens
          else:
            lowercase_								= tokens.expand(lowerCAmelCase_ ,  *tokens.shape[1:])
            lowercase_								= torch.cat((tokens, next_tokens) ,  dim=1)
        else:
          lowercase_								= -float(np.inf)
          lowercase_								= 0
          lowercase_								= scores[:, None] + logits
          seq_lengths[~is_stopped] += 1
          lowercase_								= scores_sum / seq_lengths[:, None]
          lowercase_					,    lowercase_								= scores_sum_average.view(-1).topk(lowerCAmelCase_ ,  -1)
          lowercase_								= next_tokens // scores_sum.shape[1]
          lowercase_								= seq_lengths[next_tokens_source]
          lowercase_								= next_tokens % scores_sum.shape[1]
          lowercase_								= next_tokens.unsqueeze(1)
          lowercase_								= tokens[next_tokens_source]
          lowercase_								= torch.cat((tokens, next_tokens) ,  dim=1)
          lowercase_								= generated[next_tokens_source]
          lowercase_								= scores_sum_average * seq_lengths
          lowercase_								= is_stopped[next_tokens_source]

        lowercase_								= self.transformer.transformer.wte(next_tokens.squeeze()).view(generated.shape[0] ,  1 ,  -1)
        lowercase_								= torch.cat((generated, next_token_embed) ,  dim=1)
        lowercase_								= is_stopped + next_tokens.eq(lowerCAmelCase_).squeeze()
        if is_stopped.all():
          break

      lowercase_								= scores / seq_lengths
      lowercase_								= scores.argsort(descending=lowerCAmelCase_)
      # tokens tensors are already padded to max_seq_length
      lowercase_								= [tokens[i] for i in order]
      lowercase_								= torch.stack(lowerCAmelCase_ ,  dim=0)
      lowercase_								= torch.tensor([seq_lengths[i] for i in order] ,  dtype=seq_lengths.dtype)
      return output_texts, seq_lengths

