





from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFImageClassifierOutputWithNoAttention,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_resnet import ResNetConfig


_UpperCamelCase       = logging.get_logger(__name__)

# General docstring
_UpperCamelCase       = '''ResNetConfig'''

# Base docstring
_UpperCamelCase       = '''microsoft/resnet-50'''
_UpperCamelCase       = [1, 2048, 7, 7]

# Image classification docstring
_UpperCamelCase       = '''microsoft/resnet-50'''
_UpperCamelCase       = '''tiger cat'''

_UpperCamelCase       = [
    '''microsoft/resnet-50''',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]




class 				_lowerCamelCase (    tf.keras.layers.Layer    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						UpperCAmelCase = 3				,						UpperCAmelCase = 1				,						UpperCAmelCase = "relu"				,						**UpperCAmelCase   )       ->   None:






        '''simple docstring'''






        super().__init__(**lowercase_   )
        __snake_case   :				List[Any]        =       kernel_size // 2
        __snake_case   :				Tuple        =       tf.keras.layers.ConvaD(
            lowercase_				,						kernel_size=lowercase_				,						strides=lowercase_				,						padding="valid"				,						use_bias=lowercase_				,						name="convolution"   )
        # Use same default momentum and epsilon as PyTorch equivalent
        __snake_case   :				Optional[int]        =       tf.keras.layers.BatchNormalization(epsilon=1E-5				,						momentum=0.9				,						name="normalization"   )
        __snake_case   :				List[Any]        =       ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("linear"   )







      def 						UpperCAmelCase		( self				,						UpperCAmelCase   )       ->   tf.Tensor:






        '''simple docstring'''






        __snake_case   :				List[Any]        =       (self.pad_value, self.pad_value)
        __snake_case   :				Tuple        =       tf.pad(lowercase_				,						[(0, 0), height_pad, width_pad, (0, 0)]   )
        __snake_case   :				Union[str, Any]        =       self.conv(lowercase_   )
        return hidden_state







      def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase = False   )       ->   tf.Tensor:






        '''simple docstring'''






        __snake_case   :				Optional[int]        =       self.convolution(lowercase_   )
        __snake_case   :				Dict        =       self.normalization(lowercase_				,						training=lowercase_   )
        __snake_case   :				Any        =       self.activation(lowercase_   )
        return hidden_state




class 				_lowerCamelCase (    tf.keras.layers.Layer    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						**UpperCAmelCase   )       ->   None:






        '''simple docstring'''






        super().__init__(**lowercase_   )
        __snake_case   :				List[Any]        =       TFResNetConvLayer(
            config.embedding_size				,						kernel_size=7				,						stride=2				,						activation=config.hidden_act				,						name="embedder"				,						)
        __snake_case   :				Any        =       tf.keras.layers.MaxPoolaD(pool_size=3				,						strides=2				,						padding="valid"				,						name="pooler"   )
        __snake_case   :				Optional[Any]        =       config.num_channels







      def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase = False   )       ->   tf.Tensor:






        '''simple docstring'''






        __snake_case   :				List[Any]        =       shape_list(lowercase_   )
        if tf.executing_eagerly() and num_channels != self.num_channels:
          raise ValueError(
              "Make sure that the channel dimension of the pixel values match with the one set in the configuration."   )
        __snake_case   :				Optional[Any]        =       pixel_values
        __snake_case   :				Union[str, Any]        =       self.embedder(lowercase_   )
        __snake_case   :				str        =       tf.pad(lowercase_				,						[[0, 0], [1, 1], [1, 1], [0, 0]]   )
        __snake_case   :				Dict        =       self.pooler(lowercase_   )
        return hidden_state




class 				_lowerCamelCase (    tf.keras.layers.Layer    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						UpperCAmelCase = 2				,						**UpperCAmelCase   )       ->   None:






        '''simple docstring'''






        super().__init__(**lowercase_   )
        __snake_case   :				Optional[int]        =       tf.keras.layers.ConvaD(
            lowercase_				,						kernel_size=1				,						strides=lowercase_				,						use_bias=lowercase_				,						name="convolution"   )
        # Use same default momentum and epsilon as PyTorch equivalent
        __snake_case   :				Optional[Any]        =       tf.keras.layers.BatchNormalization(epsilon=1E-5				,						momentum=0.9				,						name="normalization"   )







      def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase = False   )       ->   tf.Tensor:






        '''simple docstring'''






        __snake_case   :				Optional[Any]        =       x
        __snake_case   :				Any        =       self.convolution(lowercase_   )
        __snake_case   :				Optional[int]        =       self.normalization(lowercase_				,						training=lowercase_   )
        return hidden_state




class 				_lowerCamelCase (    tf.keras.layers.Layer    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase = 1				,						UpperCAmelCase = "relu"				,						**UpperCAmelCase   )       ->   None:






        '''simple docstring'''






        super().__init__(**lowercase_   )
        __snake_case   :				Union[str, Any]        =       in_channels != out_channels or stride != 1
        __snake_case   :				List[str]        =       TFResNetConvLayer(lowercase_				,						stride=lowercase_				,						name="layer.0"   )
        __snake_case   :				str        =       TFResNetConvLayer(lowercase_				,						activation=lowercase_				,						name="layer.1"   )
        __snake_case   :				Tuple        =       (
            TFResNetShortCut(lowercase_				,						stride=lowercase_				,						name="shortcut"   )
            if should_apply_shortcut
            else tf.keras.layers.Activation("linear"				,						name="shortcut"   )
        )
        __snake_case   :				Tuple        =       ACTaFN[activation]







      def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase = False   )       ->   tf.Tensor:






        '''simple docstring'''






        __snake_case   :				List[str]        =       hidden_state
        __snake_case   :				Tuple        =       self.conva(lowercase_				,						training=lowercase_   )
        __snake_case   :				Optional[Any]        =       self.conva(lowercase_				,						training=lowercase_   )
        __snake_case   :				Dict        =       self.shortcut(lowercase_				,						training=lowercase_   )
        hidden_state += residual
        __snake_case   :				Dict        =       self.activation(lowercase_   )
        return hidden_state




class 				_lowerCamelCase (    tf.keras.layers.Layer    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase = 1				,						UpperCAmelCase = "relu"				,						UpperCAmelCase = 4				,						**UpperCAmelCase				,						)       ->   None:






        '''simple docstring'''






        super().__init__(**lowercase_   )
        __snake_case   :				List[str]        =       in_channels != out_channels or stride != 1
        __snake_case   :				int        =       out_channels // reduction
        __snake_case   :				List[Any]        =       TFResNetConvLayer(lowercase_				,						kernel_size=1				,						name="layer.0"   )
        __snake_case   :				str        =       TFResNetConvLayer(lowercase_				,						stride=lowercase_				,						name="layer.1"   )
        __snake_case   :				Optional[int]        =       TFResNetConvLayer(lowercase_				,						kernel_size=1				,						activation=lowercase_				,						name="layer.2"   )
        __snake_case   :				Optional[int]        =       (
            TFResNetShortCut(lowercase_				,						stride=lowercase_				,						name="shortcut"   )
            if should_apply_shortcut
            else tf.keras.layers.Activation("linear"				,						name="shortcut"   )
        )
        __snake_case   :				Union[str, Any]        =       ACTaFN[activation]







      def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase = False   )       ->   tf.Tensor:






        '''simple docstring'''






        __snake_case   :				Any        =       hidden_state
        __snake_case   :				Dict        =       self.conva(lowercase_				,						training=lowercase_   )
        __snake_case   :				List[str]        =       self.conva(lowercase_				,						training=lowercase_   )
        __snake_case   :				str        =       self.conva(lowercase_				,						training=lowercase_   )
        __snake_case   :				Dict        =       self.shortcut(lowercase_				,						training=lowercase_   )
        hidden_state += residual
        __snake_case   :				str        =       self.activation(lowercase_   )
        return hidden_state




class 				_lowerCamelCase (    tf.keras.layers.Layer    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase = 2				,						UpperCAmelCase = 2				,						**UpperCAmelCase   )       ->   None:






        '''simple docstring'''






        super().__init__(**lowercase_   )

        __snake_case   :				List[str]        =       TFResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else TFResNetBasicLayer

        __snake_case   :				int        =       [layer(lowercase_				,						lowercase_				,						stride=lowercase_				,						activation=config.hidden_act				,						name="layers.0"   )]
        layers += [
            layer(lowercase_				,						lowercase_				,						activation=config.hidden_act				,						name=F"""layers.{i + 1}"""   )
            for i in range(depth - 1   )
        ]
        __snake_case   :				Tuple        =       layers







      def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase = False   )       ->   tf.Tensor:






        '''simple docstring'''






        for layer in self.stage_layers:
          __snake_case   :				List[Any]        =       layer(lowercase_				,						training=lowercase_   )
        return hidden_state




class 				_lowerCamelCase (    tf.keras.layers.Layer    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						**UpperCAmelCase   )       ->   None:






        '''simple docstring'''






        super().__init__(**lowercase_   )
        # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
        __snake_case   :				Any        =       [
            TFResNetStage(
                lowercase_				,						config.embedding_size				,						config.hidden_sizes[0]				,						stride=2 if config.downsample_in_first_stage else 1				,						depth=config.depths[0]				,						name="stages.0"				,						)
        ]
        for i, (in_channels, out_channels, depth) in enumerate(
            zip(config.hidden_sizes				,						config.hidden_sizes[1:]				,						config.depths[1:]   )   ):
          self.stages.append(TFResNetStage(lowercase_				,						lowercase_				,						lowercase_				,						depth=lowercase_				,						name=F"""stages.{i + 1}"""   )   )







      def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase = False				,						UpperCAmelCase = True				,						UpperCAmelCase = False				,						)       ->   TFBaseModelOutputWithNoAttention:






        '''simple docstring'''






        __snake_case   :				Any        =       () if output_hidden_states else None

        for stage_module in self.stages:
          if output_hidden_states:
            __snake_case   :				Optional[int]        =       hidden_states + (hidden_state,)

          __snake_case   :				Dict        =       stage_module(lowercase_				,						training=lowercase_   )

        if output_hidden_states:
          __snake_case   :				Tuple        =       hidden_states + (hidden_state,)

        if not return_dict:
          return tuple(v for v in [hidden_state, hidden_states] if v is not None   )

        return TFBaseModelOutputWithNoAttention(last_hidden_state=lowercase_				,						hidden_states=lowercase_   )




class 				_lowerCamelCase (    __UpperCAmelCase    ):



      """simple docstring"""






      UpperCAmelCase_						:      Tuple       =ResNetConfig
      UpperCAmelCase_						:      Any       ="resnet"
      UpperCAmelCase_						:      Union[str, Any]       ="pixel_values"







      @property
      def 						UpperCAmelCase		( self   )       ->   str:






        '''simple docstring'''






        return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)				,						dtype=tf.floataa   )}




_UpperCamelCase       = R'''\n    This model is a TensorFlow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n'''


_UpperCamelCase       = R'''\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'''




@keras_serializable
class 				_lowerCamelCase (    tf.keras.layers.Layer    ):



      """simple docstring"""






      UpperCAmelCase_						:      str       =ResNetConfig







      def __init__( self				,						UpperCAmelCase				,						**UpperCAmelCase   )       ->   None:






        '''simple docstring'''






        super().__init__(**lowercase_   )
        __snake_case   :				Optional[int]        =       config
        __snake_case   :				int        =       TFResNetEmbeddings(lowercase_				,						name="embedder"   )
        __snake_case   :				Union[str, Any]        =       TFResNetEncoder(lowercase_				,						name="encoder"   )
        __snake_case   :				str        =       tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowercase_   )







      @unpack_inputs
      def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase = None				,						UpperCAmelCase = None				,						UpperCAmelCase = False				,						)       ->   Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:






        '''simple docstring'''






        __snake_case   :				Optional[Any]        =       (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        __snake_case   :				str        =       return_dict if return_dict is not None else self.config.use_return_dict

        # TF 2.0 image layers can't use NCHW format when running on CPU.
        # We transpose to NHWC format and then transpose back after the full forward pass.
        # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
        __snake_case   :				List[Any]        =       tf.transpose(lowercase_				,						perm=[0, 2, 3, 1]   )
        __snake_case   :				Optional[Any]        =       self.embedder(lowercase_				,						training=lowercase_   )

        __snake_case   :				int        =       self.encoder(
            lowercase_				,						output_hidden_states=lowercase_				,						return_dict=lowercase_				,						training=lowercase_   )

        __snake_case   :				Union[str, Any]        =       encoder_outputs[0]

        __snake_case   :				str        =       self.pooler(lowercase_   )

        # Transpose all the outputs to the NCHW format
        # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
        __snake_case   :				Optional[Any]        =       tf.transpose(lowercase_				,						(0, 3, 1, 2)   )
        __snake_case   :				Tuple        =       tf.transpose(lowercase_				,						(0, 3, 1, 2)   )
        __snake_case   :				Dict        =       ()
        for hidden_state in encoder_outputs[1:]:
          __snake_case   :				Any        =       hidden_states + tuple(tf.transpose(lowercase_				,						(0, 3, 1, 2)   ) for h in hidden_state   )

        if not return_dict:
          return (last_hidden_state, pooled_output) + hidden_states

        __snake_case   :				Tuple        =       hidden_states if output_hidden_states else None

        return TFBaseModelOutputWithPoolingAndNoAttention(
            last_hidden_state=lowercase_				,						pooler_output=lowercase_				,						hidden_states=lowercase_				,						)




@add_start_docstrings(
    "The bare ResNet model outputting raw features without any specific head on top." ,   __UpperCAmelCase ,   )
class 				_lowerCamelCase (    __UpperCAmelCase    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						**UpperCAmelCase   )       ->   None:






        '''simple docstring'''






        super().__init__(lowercase_				,						**lowercase_   )
        __snake_case   :				str        =       TFResNetMainLayer(config=lowercase_				,						name="resnet"   )







      @add_start_docstrings_to_model_forward(lowercase_   )
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC				,						output_type=lowercase_				,						config_class=_CONFIG_FOR_DOC				,						modality="vision"				,						expected_output=_EXPECTED_OUTPUT_SHAPE				,						)
      @unpack_inputs
      def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase = None				,						UpperCAmelCase = None				,						UpperCAmelCase = False				,						)       ->   Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:






        '''simple docstring'''






        __snake_case   :				Dict        =       (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        __snake_case   :				List[Any]        =       return_dict if return_dict is not None else self.config.use_return_dict

        __snake_case   :				Union[str, Any]        =       self.resnet(
            pixel_values=lowercase_				,						output_hidden_states=lowercase_				,						return_dict=lowercase_				,						training=lowercase_				,						)
        return resnet_outputs




@add_start_docstrings(
    "\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    " ,   __UpperCAmelCase ,   )
class 				_lowerCamelCase (    __UpperCAmelCase ,   __UpperCAmelCase    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						**UpperCAmelCase   )       ->   None:






        '''simple docstring'''






        super().__init__(lowercase_				,						**lowercase_   )
        __snake_case   :				Optional[int]        =       config.num_labels
        __snake_case   :				Optional[Any]        =       TFResNetMainLayer(lowercase_				,						name="resnet"   )
        # classification head
        __snake_case   :				Dict        =       (
            tf.keras.layers.Dense(config.num_labels				,						name="classifier.1"   )
            if config.num_labels > 0
            else tf.keras.layers.Activation("linear"				,						name="classifier.1"   )
        )







      def 						UpperCAmelCase		( self				,						UpperCAmelCase   )       ->   tf.Tensor:






        '''simple docstring'''






        __snake_case   :				Tuple        =       tf.keras.layers.Flatten()(lowercase_   )
        __snake_case   :				int        =       self.classifier_layer(lowercase_   )
        return logits







      @add_start_docstrings_to_model_forward(lowercase_   )
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT				,						output_type=lowercase_				,						config_class=_CONFIG_FOR_DOC				,						expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT				,						)
      @unpack_inputs
      def 						UpperCAmelCase		( self				,						UpperCAmelCase = None				,						UpperCAmelCase = None				,						UpperCAmelCase = None				,						UpperCAmelCase = None				,						UpperCAmelCase = False				,						)       ->   Union[Tuple[tf.Tensor], TFImageClassifierOutputWithNoAttention]:






        '''simple docstring'''






        __snake_case   :				Any        =       return_dict if return_dict is not None else self.config.use_return_dict

        __snake_case   :				Tuple        =       self.resnet(
            lowercase_				,						output_hidden_states=lowercase_				,						return_dict=lowercase_				,						training=lowercase_   )

        __snake_case   :				Tuple        =       outputs.pooler_output if return_dict else outputs[1]

        __snake_case   :				List[str]        =       self.classifier(lowercase_   )

        __snake_case   :				Tuple        =       None if labels is None else self.hf_compute_loss(lowercase_				,						lowercase_   )

        if not return_dict:
          __snake_case   :				Optional[int]        =       (logits,) + outputs[2:]
          return (loss,) + output if loss is not None else output

        return TFImageClassifierOutputWithNoAttention(loss=lowercase_				,						logits=lowercase_				,						hidden_states=outputs.hidden_states   )








from typing import Union

import torch
from PIL import Image
from torchvision import transforms as tfms
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)




class 				_lowerCamelCase (    a    ):



      """simple docstring"""






      def __init__( self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						)       ->   Optional[Any]:






        '''simple docstring'''






        super().__init__()

        self.register_modules(vae=UpperCAmelCase				,						text_encoder=UpperCAmelCase				,						tokenizer=UpperCAmelCase				,						unet=UpperCAmelCase				,						scheduler=UpperCAmelCase   )







      def 						UpperCAmelCase		( self				,						UpperCAmelCase   )       ->   str:






        '''simple docstring'''






        with torch.no_grad():
          __snake_case   :				Tuple        =       self.vae.encode(tfms.ToTensor()(UpperCAmelCase   ).unsqueeze(0   ).to(self.device   ) * 2 - 1   )
          __snake_case   :				Optional[int]        =       0.18_215 * latent.latent_dist.sample()
        return latent







      def 						UpperCAmelCase		( self				,						UpperCAmelCase   )       ->   str:






        '''simple docstring'''






        __snake_case   :				Optional[Any]        =       (1 / 0.18_215) * latent
        with torch.no_grad():
          __snake_case   :				Any        =       self.vae.decode(UpperCAmelCase   ).sample
        __snake_case   :				Optional[int]        =       (img / 2 + 0.5).clamp(0				,						1   )
        __snake_case   :				Tuple        =       img.detach().cpu().permute(0				,						2				,						3				,						1   ).numpy()
        __snake_case   :				Dict        =       (img * 255).round().astype("uint8"   )
        return Image.fromarray(img[0]   )







      def 						UpperCAmelCase		( self				,						UpperCAmelCase   )       ->   Dict:






        '''simple docstring'''






        __snake_case   :				int        =       self.tokenizer(
            UpperCAmelCase				,						padding="max_length"				,						max_length=self.tokenizer.model_max_length				,						truncation=UpperCAmelCase				,						return_tensors="pt"				,						)

        __snake_case   :				List[Any]        =       self.text_encoder(text_input.input_ids.to(self.device   )   )[0]

        __snake_case   :				Union[str, Any]        =       self.tokenizer(
            ""				,						padding="max_length"				,						max_length=self.tokenizer.model_max_length				,						truncation=UpperCAmelCase				,						return_tensors="pt"				,						)

        __snake_case   :				int        =       self.text_encoder(uncond_input.input_ids.to(self.device   )   )[0]

        return torch.cat([uncond_embedding, text_embedding]   )







      def __call__( self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase = 0.3				,						UpperCAmelCase = 0.6				,						UpperCAmelCase = 0.5				,						UpperCAmelCase = 42				,						UpperCAmelCase = 50				,						UpperCAmelCase = 7.5				,						)       ->   Image.Image:






        '''simple docstring'''






        __snake_case   :				Tuple        =       steps - int(kmin * steps   )
        __snake_case   :				Optional[int]        =       steps - int(kmax * steps   )

        __snake_case   :				Tuple        =       self.prep_text(UpperCAmelCase   )

        self.scheduler.set_timesteps(UpperCAmelCase   )

        __snake_case    ,   __snake_case   :				Optional[Any]        =       img.size
        __snake_case   :				List[Any]        =       self.encode(UpperCAmelCase   )

        torch.manual_seed(UpperCAmelCase   )
        __snake_case   :				str        =       torch.randn(
            (1, self.unet.config.in_channels, height // 8, width // 8)				,						).to(self.device   )

        __snake_case   :				int        =       self.scheduler.add_noise(
            UpperCAmelCase				,						UpperCAmelCase				,						timesteps=self.scheduler.timesteps[tmax]				,						)

        __snake_case   :				Optional[Any]        =       torch.cat([latents] * 2   )

        __snake_case   :				Union[str, Any]        =       self.scheduler.scale_model_input(UpperCAmelCase				,						self.scheduler.timesteps[tmax]   )

        with torch.no_grad():
          __snake_case   :				int        =       self.unet(
              UpperCAmelCase				,						self.scheduler.timesteps[tmax]				,						encoder_hidden_states=UpperCAmelCase				,						).sample

        __snake_case    ,   __snake_case   :				Union[str, Any]        =       pred.chunk(2   )
        __snake_case   :				Optional[Any]        =       pred_uncond + guidance_scale * (pred_text - pred_uncond)

        __snake_case   :				int        =       self.scheduler.step(UpperCAmelCase				,						self.scheduler.timesteps[tmax]				,						UpperCAmelCase   ).prev_sample

        for i, t in enumerate(tqdm(self.scheduler.timesteps   )   ):
          if i > tmax:
            if i < tmin:  # layout generation phase
              __snake_case   :				int        =       self.scheduler.add_noise(
                  UpperCAmelCase				,						UpperCAmelCase				,						timesteps=UpperCAmelCase				,						)

              __snake_case   :				str        =       (mix_factor * latents) + (
                  1 - mix_factor
              ) * orig_latents  # interpolating between layout noise and conditionally generated noise to preserve layout sematics
              __snake_case   :				List[Any]        =       torch.cat([input] * 2   )

            else:  # content generation phase
              __snake_case   :				str        =       torch.cat([latents] * 2   )

            __snake_case   :				Tuple        =       self.scheduler.scale_model_input(UpperCAmelCase				,						UpperCAmelCase   )

            with torch.no_grad():
              __snake_case   :				List[str]        =       self.unet(
                  UpperCAmelCase				,						UpperCAmelCase				,						encoder_hidden_states=UpperCAmelCase				,						).sample

            __snake_case    ,   __snake_case   :				str        =       pred.chunk(2   )
            __snake_case   :				List[Any]        =       pred_uncond + guidance_scale * (pred_text - pred_uncond)

            __snake_case   :				List[Any]        =       self.scheduler.step(UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   ).prev_sample

        return self.decode(UpperCAmelCase   )

