



"""simple docstring"""

from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)




@flax.struct.dataclass
class SCREAMING_SNAKE_CASE__					(		__snake_case      ):
						_lowerCAmelCase							=					42







@flax_register_to_config
class SCREAMING_SNAKE_CASE__					(		nn.Module     ,					__snake_case     ,					__snake_case      ):
						_lowerCAmelCase							=					3_2
						_lowerCAmelCase							=					4
						_lowerCAmelCase							=					4
						_lowerCAmelCase							=					(
						    "CrossAttnDownBlock2D",
						    "CrossAttnDownBlock2D",
						    "CrossAttnDownBlock2D",
						    "DownBlock2D",
						)
						_lowerCAmelCase							=					("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
						_lowerCAmelCase							=					False
						_lowerCAmelCase							=					(3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
						_lowerCAmelCase							=					2
						_lowerCAmelCase							=					8
						_lowerCAmelCase							=					None
						_lowerCAmelCase							=					1_2_8_0
						_lowerCAmelCase							=					0.0
						_lowerCAmelCase							=					False
						_lowerCAmelCase							=					jnp.floataa
						_lowerCAmelCase							=					True
						_lowerCAmelCase							=					0
						_lowerCAmelCase							=					False



						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

													__a		:		Dict									=       (1, self.in_channels, self.sample_size, self.sample_size)
													__a		:		Optional[int]									=       jnp.zeros(_lowercase						,     dtype=jnp.floataa							)
													__a		:		List[str]									=       jnp.ones((1,)						,     dtype=jnp.intaa							)
													__a		:		str									=       jnp.zeros((1, 1, self.cross_attention_dim)						,     dtype=jnp.floataa							)

													__a							,     __a		:		List[str]									=       jax.random.split(_lowercase							)
													__a		:		Tuple									=       {"""params""": params_rng, """dropout""": dropout_rng}

													return self.init(_lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)["params"]



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		List[Any]									=       self.block_out_channels
													__a		:		List[Any]									=       block_out_channels[0] * 4

													if self.num_attention_heads is not None:
																				raise ValueError(
																				    """At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."""							)

													# If `num_attention_heads` is not defined (which is the case for most models)
													# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
													# The reason for this behavior is to correct for incorrectly named variables that were introduced
													# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
													# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
													# which is why we correct for the naming here.
													__a		:		Union[str, Any]									=       self.num_attention_heads or self.attention_head_dim

													# input
													__a		:		List[Any]									=       nn.Conv(
													    block_out_channels[0]						,     kernel_size=(3, 3)						,     strides=(1, 1)						,     padding=((1, 1), (1, 1))						,     dtype=self.dtype						,     )

													# time
													__a		:		Optional[Any]									=       FlaxTimesteps(
													    block_out_channels[0]						,     flip_sin_to_cos=self.flip_sin_to_cos						,     freq_shift=self.config.freq_shift							)
													__a		:		Optional[Any]									=       FlaxTimestepEmbedding(_lowercase						,     dtype=self.dtype							)

													__a		:		str									=       self.only_cross_attention
													if isinstance(_lowercase						,     _lowercase							):
																				__a		:		Dict									=       (only_cross_attention,) * len(self.down_block_types							)

													if isinstance(_lowercase						,     _lowercase							):
																				__a		:		int									=       (num_attention_heads,) * len(self.down_block_types							)

													# down
													__a		:		Union[str, Any]									=       []
													__a		:		Optional[Any]									=       block_out_channels[0]
													for i, down_block_type in enumerate(self.down_block_types							):
																				__a		:		List[Any]									=       output_channel
																				__a		:		Dict									=       block_out_channels[i]
																				__a		:		Dict									=       i == len(_lowercase							) - 1

																				if down_block_type == "CrossAttnDownBlock2D":
																											__a		:		Optional[Any]									=       FlaxCrossAttnDownBlockaD(
																											    in_channels=_lowercase						,     out_channels=_lowercase						,     dropout=self.dropout						,     num_layers=self.layers_per_block						,     num_attention_heads=num_attention_heads[i]						,     add_downsample=not is_final_block						,     use_linear_projection=self.use_linear_projection						,     only_cross_attention=only_cross_attention[i]						,     use_memory_efficient_attention=self.use_memory_efficient_attention						,     dtype=self.dtype						,     )
																				else:
																											__a		:		Optional[Any]									=       FlaxDownBlockaD(
																											    in_channels=_lowercase						,     out_channels=_lowercase						,     dropout=self.dropout						,     num_layers=self.layers_per_block						,     add_downsample=not is_final_block						,     dtype=self.dtype						,     )

																				down_blocks.append(_lowercase							)
													__a		:		Optional[int]									=       down_blocks

													# mid
													__a		:		Tuple									=       FlaxUNetMidBlockaDCrossAttn(
													    in_channels=block_out_channels[-1]						,     dropout=self.dropout						,     num_attention_heads=num_attention_heads[-1]						,     use_linear_projection=self.use_linear_projection						,     use_memory_efficient_attention=self.use_memory_efficient_attention						,     dtype=self.dtype						,     )

													# up
													__a		:		Optional[int]									=       []
													__a		:		Union[str, Any]									=       list(reversed(_lowercase							)							)
													__a		:		List[str]									=       list(reversed(_lowercase							)							)
													__a		:		Tuple									=       list(reversed(_lowercase							)							)
													__a		:		Dict									=       reversed_block_out_channels[0]
													for i, up_block_type in enumerate(self.up_block_types							):
																				__a		:		Union[str, Any]									=       output_channel
																				__a		:		Tuple									=       reversed_block_out_channels[i]
																				__a		:		Optional[int]									=       reversed_block_out_channels[min(i + 1						,     len(_lowercase							) - 1							)]

																				__a		:		List[str]									=       i == len(_lowercase							) - 1

																				if up_block_type == "CrossAttnUpBlock2D":
																											__a		:		List[Any]									=       FlaxCrossAttnUpBlockaD(
																											    in_channels=_lowercase						,     out_channels=_lowercase						,     prev_output_channel=_lowercase						,     num_layers=self.layers_per_block + 1						,     num_attention_heads=reversed_num_attention_heads[i]						,     add_upsample=not is_final_block						,     dropout=self.dropout						,     use_linear_projection=self.use_linear_projection						,     only_cross_attention=only_cross_attention[i]						,     use_memory_efficient_attention=self.use_memory_efficient_attention						,     dtype=self.dtype						,     )
																				else:
																											__a		:		Any									=       FlaxUpBlockaD(
																											    in_channels=_lowercase						,     out_channels=_lowercase						,     prev_output_channel=_lowercase						,     num_layers=self.layers_per_block + 1						,     add_upsample=not is_final_block						,     dropout=self.dropout						,     dtype=self.dtype						,     )

																				up_blocks.append(_lowercase							)
																				__a		:		Union[str, Any]									=       output_channel
													__a		:		int									=       up_blocks

													# out
													__a		:		Optional[Any]									=       nn.GroupNorm(num_groups=32						,     epsilon=1e-5							)
													__a		:		Optional[Any]									=       nn.Conv(
													    self.out_channels						,     kernel_size=(3, 3)						,     strides=(1, 1)						,     padding=((1, 1), (1, 1))						,     dtype=self.dtype						,     )





						def __call__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase=None						,     _lowercase=None						,     _lowercase = True						,     _lowercase = False						,     ):





													'''simple docstring'''

													if not isinstance(_lowercase						,     jnp.ndarray							):
																				__a		:		Union[str, Any]									=       jnp.array([timesteps]						,     dtype=jnp.intaa							)
													elif isinstance(_lowercase						,     jnp.ndarray							) and len(timesteps.shape							) == 0:
																				__a		:		Union[str, Any]									=       timesteps.astype(dtype=jnp.floataa							)
																				__a		:		List[str]									=       jnp.expand_dims(_lowercase						,     0							)

													__a		:		Any									=       self.time_proj(_lowercase							)
													__a		:		Tuple									=       self.time_embedding(_lowercase							)

													# 2. pre-process
													__a		:		Optional[Any]									=       jnp.transpose(_lowercase						,     (0, 2, 3, 1)							)
													__a		:		List[str]									=       self.conv_in(_lowercase							)

													# 3. down
													__a		:		Dict									=       (sample,)
													for down_block in self.down_blocks:
																				if isinstance(_lowercase						,     _lowercase							):
																											__a							,     __a		:		Tuple									=       down_block(_lowercase						,     _lowercase						,     _lowercase						,     deterministic=not train							)
																				else:
																											__a							,     __a		:		Tuple									=       down_block(_lowercase						,     _lowercase						,     deterministic=not train							)
																				down_block_res_samples += res_samples

													if down_block_additional_residuals is not None:
																				__a		:		Dict									=       ()

																				for down_block_res_sample, down_block_additional_residual in zip(
																				    _lowercase						,     _lowercase							):
																											down_block_res_sample += down_block_additional_residual
																											new_down_block_res_samples += (down_block_res_sample,)

																				__a		:		Dict									=       new_down_block_res_samples

													# 4. mid
													__a		:		Dict									=       self.mid_block(_lowercase						,     _lowercase						,     _lowercase						,     deterministic=not train							)

													if mid_block_additional_residual is not None:
																				sample += mid_block_additional_residual

													# 5. up
													for up_block in self.up_blocks:
																				__a		:		List[str]									=       down_block_res_samples[-(self.layers_per_block + 1) :]
																				__a		:		Optional[Any]									=       down_block_res_samples[: -(self.layers_per_block + 1)]
																				if isinstance(_lowercase						,     _lowercase							):
																											__a		:		List[str]									=       up_block(
																											    _lowercase						,     temb=_lowercase						,     encoder_hidden_states=_lowercase						,     res_hidden_states_tuple=_lowercase						,     deterministic=not train						,     )
																				else:
																											__a		:		Tuple									=       up_block(_lowercase						,     temb=_lowercase						,     res_hidden_states_tuple=_lowercase						,     deterministic=not train							)

        # 6. post-process
													__a		:		Optional[Any]									=       self.conv_norm_out(_lowercase							)
													__a		:		Optional[Any]									=       nn.silu(_lowercase							)
													__a		:		Any									=       self.conv_out(_lowercase							)
													__a		:		int									=       jnp.transpose(_lowercase						,     (0, 3, 1, 2)							)

													if not return_dict:
																				return (sample,)

													return FlaxUNetaDConditionOutput(sample=_lowercase							)






"""simple docstring"""

import importlib

import torch
import yaml
from omegaconf import OmegaConf
from taming.models.vqgan import VQModel





def       __magic_name__    (   _lowerCamelCase				: int  ,       _lowerCamelCase				: Optional[Any]=False					):
							__a		:		Dict									=       OmegaConf.load(_lowerCamelCase					)
							if display:
														print(yaml.dump(OmegaConf.to_container(_lowerCamelCase					)					)					)
							return config





def       __magic_name__    (   _lowerCamelCase				: str  ,       _lowerCamelCase				: Union[str, Any]=None  ,       _lowerCamelCase				: int=None					):
							if conf_path is None:
														__a		:		str									=       """./model_checkpoints/vqgan_only.yaml"""
							__a		:		List[Any]									=       load_config(_lowerCamelCase  ,       display=_lowerCamelCase					)
							__a		:		Dict									=       VQModel(**config.model.params					)
							if ckpt_path is None:
														__a		:		List[Any]									=       """./model_checkpoints/vqgan_only.pt"""
							__a		:		Tuple									=       torch.load(_lowerCamelCase  ,       map_location=_lowerCamelCase					)
							if ".ckpt" in ckpt_path:
														__a		:		List[str]									=       sd["""state_dict"""]
							model.load_state_dict(_lowerCamelCase  ,       strict=_lowerCamelCase					)
							model.to(_lowerCamelCase					)
							del sd
							return model





def       __magic_name__    (   _lowerCamelCase				: Optional[int]  ,       _lowerCamelCase				: List[str]					):
							__a							,     __a							,     __a		:		Tuple									=       model.encode(_lowerCamelCase					)
							print(F'''VQGAN --- {model.__class__.__name__}: latent shape: {z.shape[2:]}'''					)
							__a		:		Union[str, Any]									=       model.decode(_lowerCamelCase					)
							return xrec





def       __magic_name__    (   _lowerCamelCase				: Optional[int]  ,       _lowerCamelCase				: Union[str, Any]=False					):
							__a							,     __a		:		Optional[Any]									=       string.rsplit("""."""  ,       1					)
							if reload:
														__a		:		Optional[Any]									=       importlib.import_module(_lowerCamelCase					)
														importlib.reload(_lowerCamelCase					)
							return getattr(importlib.import_module(_lowerCamelCase  ,       package=_lowerCamelCase					)  ,       cls					)





def       __magic_name__    (   _lowerCamelCase				: Any					):
							if "target" not in config:
														raise KeyError("""Expected key `target` to instantiate."""					)
							return get_obj_from_str(config["""target"""]					)(**config.get("""params"""  ,       {}					)					)





def       __magic_name__    (   _lowerCamelCase				: Any  ,       _lowerCamelCase				: Dict  ,       _lowerCamelCase				: int=True  ,       _lowerCamelCase				: int=True					):
							__a		:		Union[str, Any]									=       instantiate_from_config(_lowerCamelCase					)
							if sd is not None:
														model.load_state_dict(_lowerCamelCase					)
							if gpu:
														model.cuda()
							if eval_mode:
														model.eval()
							return {"model": model}







def       __magic_name__    (   _lowerCamelCase				: Any  ,       _lowerCamelCase				: Union[str, Any]  ,       _lowerCamelCase				: str  ,       _lowerCamelCase				: int					):
							# load the specified checkpoint
							if ckpt:
														__a		:		List[str]									=       torch.load(_lowerCamelCase  ,       map_location="""cpu"""					)
														__a		:		Any									=       pl_sd["""global_step"""]
														print(F'''loaded model from global step {global_step}.'''					)
							else:
														__a		:		List[Any]									=       {"""state_dict""": None}
														__a		:		Any									=       None
							__a		:		Union[str, Any]									=       load_model_from_config(config.model  ,       pl_sd["""state_dict"""]  ,       gpu=_lowerCamelCase  ,       eval_mode=_lowerCamelCase					)["""model"""]
							return model, global_step

