


from typing import Optional

from torch import nn

from .transformer_ad import TransformeraDModel, TransformeraDModelOutput
class 	lowerCamelCase_       (				nn.Module				):



					'''simple docstring'''







					def __init__(				self ,       snake_case_ = 1_6 ,       snake_case_ = 8_8 ,       snake_case_ = None ,       snake_case_ = 1 ,       snake_case_ = 0.0 ,       snake_case_ = 3_2 ,       snake_case_ = None ,       snake_case_ = False ,       snake_case_ = None ,       snake_case_ = None ,       snake_case_ = "geglu" ,       snake_case_ = None ,       )   ->			Dict:







								'''simple docstring'''

								super().__init__()
								__lowercase 					=   nn.ModuleList(
								    [
								        TransformeraDModel(
								            num_attention_heads=snake_case_ ,       attention_head_dim=snake_case_ ,       in_channels=snake_case_ ,       num_layers=snake_case_ ,       dropout=snake_case_ ,       norm_num_groups=snake_case_ ,       cross_attention_dim=snake_case_ ,       attention_bias=snake_case_ ,       sample_size=snake_case_ ,       num_vector_embeds=snake_case_ ,       activation_fn=snake_case_ ,       num_embeds_ada_norm=snake_case_ ,       )
								        for _ in range(2						)
								    ]						)

								# Variables that can be set by a pipeline:

								# The ratio of transformer1 to transformer2's output states to be combined during inference
								__lowercase 					=   0.5

								# The shape of `encoder_hidden_states` is expected to be
								# `(batch_size, condition_lengths[0]+condition_lengths[1], num_features)`
								__lowercase 					=   [7_7, 2_5_7]

								# Which transformer to use to encode which condition.
								# E.g. `(1, 0)` means that we'll use `transformers[1](conditions[0])` and `transformers[0](conditions[1])`
								__lowercase 					=   [1, 0]






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_=None ,       snake_case_=None ,       snake_case_=None ,       snake_case_ = True ,       )   ->			Any:







								'''simple docstring'''

								__lowercase 					=   hidden_states

								__lowercase 					=   []
								__lowercase 					=   0
								# attention_mask is not used yet
								for i in range(2						):
											# for each of the two transformers, pass the corresponding condition tokens
											__lowercase 					=   encoder_hidden_states[:, tokens_start : tokens_start + self.condition_lengths[i]]
											__lowercase 					=   self.transformer_index_for_condition[i]
											__lowercase 					=   self.transformers[transformer_index](
											    snake_case_ ,       encoder_hidden_states=snake_case_ ,       timestep=snake_case_ ,       cross_attention_kwargs=snake_case_ ,       return_dict=snake_case_ ,       )[0]
											encoded_states.append(encoded_state - input_states						)
											tokens_start += self.condition_lengths[i]

								__lowercase 					=   encoded_states[0] * self.mix_ratio + encoded_states[1] * (1 - self.mix_ratio)
								__lowercase 					=   output_states + input_states

								if not return_dict:
											return (output_states,)

								return TransformeraDModelOutput(sample=snake_case_						)






from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL

from ...utils import BaseOutput, OptionalDependencyNotAvailable, is_torch_available, is_transformers_available
from .timesteps import (
    fastaa_timesteps,
    smartaa_timesteps,
    smartaa_timesteps,
    smartaaa_timesteps,
    smartaaa_timesteps,
    superaa_timesteps,
    superaa_timesteps,
    superaaa_timesteps,
)
@dataclass
class 	lowerCamelCase_       (				lowerCAmelCase__				):



					'''simple docstring'''







					__UpperCAmelCase            =				42
					__UpperCAmelCase            =				42
					__UpperCAmelCase            =				42


try:
	if not (is_transformers_available() and is_torch_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
	from .pipeline_if import IFPipeline
	from .pipeline_if_imgaimg import IFImgaImgPipeline
	from .pipeline_if_imgaimg_superresolution import IFImgaImgSuperResolutionPipeline
	from .pipeline_if_inpainting import IFInpaintingPipeline
	from .pipeline_if_inpainting_superresolution import IFInpaintingSuperResolutionPipeline
	from .pipeline_if_superresolution import IFSuperResolutionPipeline
	from .safety_checker import IFSafetyChecker
	from .watermark import IFWatermarker


