



from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


__lowerCamelCase     :							List[Any]  							=				logging.get_logger(__name__)

# General docstring
__lowerCamelCase     :							Optional[int]  							=				"ResNetConfig"

# Base docstring
__lowerCamelCase     :							Union[str, Any]  							=				"microsoft/resnet-50"
__lowerCamelCase     :							Tuple  							=				[1, 2_048, 7, 7]

# Image classification docstring
__lowerCamelCase     :							str  							=				"microsoft/resnet-50"
__lowerCamelCase     :							Optional[Any]  							=				"tiger cat"

__lowerCamelCase     :							Optional[Any]  							=				[
    "microsoft/resnet-50",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]





class  __magic_name__   (							nn.Module						):






		def __init__(						self      :     Optional[Any]				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int = 3				,		UpperCamelCase__      :     int = 1				,		UpperCamelCase__      :     str = "relu"      )				->					Optional[int]:
								'''simple docstring'''







								super().__init__()
								UpperCAmelCase         =     nn.Convad(
								    UpperCamelCase__				,		UpperCamelCase__				,		kernel_size=UpperCamelCase__				,		stride=UpperCamelCase__				,		padding=kernel_size // 2				,		bias=UpperCamelCase__      )
								UpperCAmelCase         =     nn.BatchNormad(UpperCamelCase__      )
								UpperCAmelCase         =     ACTaFN[activation] if activation is not None else nn.Identity()







		def SCREAMING_SNAKE_CASE_			(						self      :     List[str]				,		UpperCamelCase__      :     Tensor      )				->					Tensor:
								'''simple docstring'''







								UpperCAmelCase         =     self.convolution(UpperCamelCase__      )
								UpperCAmelCase         =     self.normalization(UpperCamelCase__      )
								UpperCAmelCase         =     self.activation(UpperCamelCase__      )
								return hidden_state





class  __magic_name__   (							nn.Module						):






		def __init__(						self      :     Union[str, Any]				,		UpperCamelCase__      :     ResNetConfig      )				->					Dict:
								'''simple docstring'''







								super().__init__()
								UpperCAmelCase         =     ResNetConvLayer(
								    config.num_channels				,		config.embedding_size				,		kernel_size=7				,		stride=2				,		activation=config.hidden_act      )
								UpperCAmelCase         =     nn.MaxPoolad(kernel_size=3				,		stride=2				,		padding=1      )
								UpperCAmelCase         =     config.num_channels







		def SCREAMING_SNAKE_CASE_			(						self      :     int				,		UpperCamelCase__      :     Tensor      )				->					Tensor:
								'''simple docstring'''







								UpperCAmelCase         =     pixel_values.shape[1]
								if 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."      )
								UpperCAmelCase         =     self.embedder(UpperCamelCase__      )
								UpperCAmelCase         =     self.pooler(UpperCamelCase__      )
								return embedding





class  __magic_name__   (							nn.Module						):






		def __init__(						self      :     int				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int = 2      )				->					Tuple:
								'''simple docstring'''







								super().__init__()
								UpperCAmelCase         =     nn.Convad(UpperCamelCase__				,		UpperCamelCase__				,		kernel_size=1				,		stride=UpperCamelCase__				,		bias=UpperCamelCase__      )
								UpperCAmelCase         =     nn.BatchNormad(UpperCamelCase__      )







		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]				,		UpperCamelCase__      :     Tensor      )				->					Tensor:
								'''simple docstring'''







								UpperCAmelCase         =     self.convolution(UpperCamelCase__      )
								UpperCAmelCase         =     self.normalization(UpperCamelCase__      )
								return hidden_state





class  __magic_name__   (							nn.Module						):






		def __init__(						self      :     int				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int = 1				,		UpperCamelCase__      :     str = "relu"      )				->					Tuple:
								'''simple docstring'''







								super().__init__()
								UpperCAmelCase         =     in_channels != out_channels or stride != 1
								UpperCAmelCase         =     (
								    ResNetShortCut(UpperCamelCase__				,		UpperCamelCase__				,		stride=UpperCamelCase__      ) if should_apply_shortcut else nn.Identity()
								)
								UpperCAmelCase         =     nn.Sequential(
								    ResNetConvLayer(UpperCamelCase__				,		UpperCamelCase__				,		stride=UpperCamelCase__      )				,		ResNetConvLayer(UpperCamelCase__				,		UpperCamelCase__				,		activation=UpperCamelCase__      )				,		)
								UpperCAmelCase         =     ACTaFN[activation]







		def SCREAMING_SNAKE_CASE_			(						self      :     List[str]				,		UpperCamelCase__      :     Tuple      )				->					str:
								'''simple docstring'''







								UpperCAmelCase         =     hidden_state
								UpperCAmelCase         =     self.layer(UpperCamelCase__      )
								UpperCAmelCase         =     self.shortcut(UpperCamelCase__      )
								hidden_state += residual
								UpperCAmelCase         =     self.activation(UpperCamelCase__      )
								return hidden_state





class  __magic_name__   (							nn.Module						):






		def __init__(						self      :     Tuple				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int = 1				,		UpperCamelCase__      :     str = "relu"				,		UpperCamelCase__      :     int = 4      )				->					List[Any]:
								'''simple docstring'''







								super().__init__()
								UpperCAmelCase         =     in_channels != out_channels or stride != 1
								UpperCAmelCase         =     out_channels // reduction
								UpperCAmelCase         =     (
								    ResNetShortCut(UpperCamelCase__				,		UpperCamelCase__				,		stride=UpperCamelCase__      ) if should_apply_shortcut else nn.Identity()
								)
								UpperCAmelCase         =     nn.Sequential(
								    ResNetConvLayer(UpperCamelCase__				,		UpperCamelCase__				,		kernel_size=1      )				,		ResNetConvLayer(UpperCamelCase__				,		UpperCamelCase__				,		stride=UpperCamelCase__      )				,		ResNetConvLayer(UpperCamelCase__				,		UpperCamelCase__				,		kernel_size=1				,		activation=UpperCamelCase__      )				,		)
								UpperCAmelCase         =     ACTaFN[activation]







		def SCREAMING_SNAKE_CASE_			(						self      :     List[Any]				,		UpperCamelCase__      :     Union[str, Any]      )				->					Dict:
								'''simple docstring'''







								UpperCAmelCase         =     hidden_state
								UpperCAmelCase         =     self.layer(UpperCamelCase__      )
								UpperCAmelCase         =     self.shortcut(UpperCamelCase__      )
								hidden_state += residual
								UpperCAmelCase         =     self.activation(UpperCamelCase__      )
								return hidden_state





class  __magic_name__   (							nn.Module						):






		def __init__(						self      :     Union[str, Any]				,		UpperCamelCase__      :     ResNetConfig				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     int = 2				,		UpperCamelCase__      :     int = 2				,		)				->					Optional[Any]:
								'''simple docstring'''







								super().__init__()

								UpperCAmelCase         =     ResNetBottleNeckLayer if config.layer_type == "bottleneck" else ResNetBasicLayer

								UpperCAmelCase         =     nn.Sequential(
								    # downsampling is done in the first layer with stride of 2
								    layer(UpperCamelCase__				,		UpperCamelCase__				,		stride=UpperCamelCase__				,		activation=config.hidden_act      )				,		*[layer(UpperCamelCase__				,		UpperCamelCase__				,		activation=config.hidden_act      ) for _ in range(depth - 1      )]				,		)







		def SCREAMING_SNAKE_CASE_			(						self      :     List[Any]				,		UpperCamelCase__      :     Tensor      )				->					Tensor:
								'''simple docstring'''







								UpperCAmelCase         =     input
								for layer in self.layers:
														UpperCAmelCase         =     layer(UpperCamelCase__      )
								return hidden_state





class  __magic_name__   (							nn.Module						):






		def __init__(						self      :     str				,		UpperCamelCase__      :     ResNetConfig      )				->					Optional[Any]:
								'''simple docstring'''







								super().__init__()
								UpperCAmelCase         =     nn.ModuleList([]      )
								# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
								self.stages.append(
								    ResNetStage(
								        UpperCamelCase__				,		config.embedding_size				,		config.hidden_sizes[0]				,		stride=2 if config.downsample_in_first_stage else 1				,		depth=config.depths[0]				,		)      )
								UpperCAmelCase         =     zip(config.hidden_sizes				,		config.hidden_sizes[1:]      )
								for (in_channels, out_channels), depth in zip(UpperCamelCase__				,		config.depths[1:]      ):
														self.stages.append(ResNetStage(UpperCamelCase__				,		UpperCamelCase__				,		UpperCamelCase__				,		depth=UpperCamelCase__      )      )







		def SCREAMING_SNAKE_CASE_			(						self      :     Dict				,		UpperCamelCase__      :     Tensor				,		UpperCamelCase__      :     bool = False				,		UpperCamelCase__      :     bool = True      )				->					BaseModelOutputWithNoAttention:
								'''simple docstring'''







								UpperCAmelCase         =     () if output_hidden_states else None

								for stage_module in self.stages:
														if output_hidden_states:
																				UpperCAmelCase         =     hidden_states + (hidden_state,)

														UpperCAmelCase         =     stage_module(UpperCamelCase__      )

								if output_hidden_states:
														UpperCAmelCase         =     hidden_states + (hidden_state,)

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

								return BaseModelOutputWithNoAttention(
								    last_hidden_state=UpperCamelCase__				,		hidden_states=UpperCamelCase__				,		)





class  __magic_name__   (							A__						):
		lowercase      :       int          =ResNetConfig
		lowercase      :       Dict          ='''resnet'''
		lowercase      :       Tuple          ='''pixel_values'''
		lowercase      :       str          =True






		def SCREAMING_SNAKE_CASE_			(						self      :     Union[str, Any]				,		UpperCamelCase__      :     Union[str, Any]      )				->					Optional[Any]:
								'''simple docstring'''







								if isinstance(UpperCamelCase__				,		nn.Convad      ):
														nn.init.kaiming_normal_(module.weight				,		mode="fan_out"				,		nonlinearity="relu"      )
								elif isinstance(UpperCamelCase__				,		(nn.BatchNormad, nn.GroupNorm)      ):
														nn.init.constant_(module.weight				,		1      )
														nn.init.constant_(module.bias				,		0      )







		def SCREAMING_SNAKE_CASE_			(						self      :     int				,		UpperCamelCase__      :     Dict				,		UpperCamelCase__      :     Optional[int]=False      )				->					Union[str, Any]:
								'''simple docstring'''







								if isinstance(UpperCamelCase__				,		UpperCamelCase__      ):
														UpperCAmelCase         =     value
__lowerCamelCase     :							Optional[int]  							=				r"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch 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 [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

__lowerCamelCase     :							List[Any]  							=				r"\n    Args:\n        pixel_values (`torch.FloatTensor` 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"





@add_start_docstrings(
    '''The bare ResNet model outputting raw features without any specific head on top.''',  A__,  )
class  __magic_name__   (							A__						):






		def __init__(						self      :     Union[str, Any]				,		UpperCamelCase__      :     str      )				->					Tuple:
								'''simple docstring'''







								super().__init__(UpperCamelCase__      )
								UpperCAmelCase         =     config
								UpperCAmelCase         =     ResNetEmbeddings(UpperCamelCase__      )
								UpperCAmelCase         =     ResNetEncoder(UpperCamelCase__      )
								UpperCAmelCase         =     nn.AdaptiveAvgPoolad((1, 1)      )
								# Initialize weights and apply final processing
								self.post_init()







		@add_start_docstrings_to_model_forward(UpperCamelCase__      )
		@add_code_sample_docstrings(
		    checkpoint=_CHECKPOINT_FOR_DOC				,		output_type=UpperCamelCase__				,		config_class=_CONFIG_FOR_DOC				,		modality="vision"				,		expected_output=_EXPECTED_OUTPUT_SHAPE				,		)
		def SCREAMING_SNAKE_CASE_			(						self      :     int				,		UpperCamelCase__      :     Tensor				,		UpperCamelCase__      :     Optional[bool] = None				,		UpperCamelCase__      :     Optional[bool] = None      )				->					BaseModelOutputWithPoolingAndNoAttention:
								'''simple docstring'''







								UpperCAmelCase         =     (
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								UpperCAmelCase         =     return_dict if return_dict is not None else self.config.use_return_dict

								UpperCAmelCase         =     self.embedder(UpperCamelCase__      )

								UpperCAmelCase         =     self.encoder(
								    UpperCamelCase__				,		output_hidden_states=UpperCamelCase__				,		return_dict=UpperCamelCase__      )

								UpperCAmelCase         =     encoder_outputs[0]

								UpperCAmelCase         =     self.pooler(UpperCamelCase__      )

								if not return_dict:
														return (last_hidden_state, pooled_output) + encoder_outputs[1:]

								return BaseModelOutputWithPoolingAndNoAttention(
								    last_hidden_state=UpperCamelCase__				,		pooler_output=UpperCamelCase__				,		hidden_states=encoder_outputs.hidden_states				,		)





@add_start_docstrings(
    '''
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    ''',  A__,  )
class  __magic_name__   (							A__						):






		def __init__(						self      :     Any				,		UpperCamelCase__      :     Optional[Any]      )				->					int:
								'''simple docstring'''







								super().__init__(UpperCamelCase__      )
								UpperCAmelCase         =     config.num_labels
								UpperCAmelCase         =     ResNetModel(UpperCamelCase__      )
								# classification head
								UpperCAmelCase         =     nn.Sequential(
								    nn.Flatten()				,		nn.Linear(config.hidden_sizes[-1]				,		config.num_labels      ) if config.num_labels > 0 else nn.Identity()				,		)
								# initialize weights and apply final processing
								self.post_init()







		@add_start_docstrings_to_model_forward(UpperCamelCase__      )
		@add_code_sample_docstrings(
		    checkpoint=_IMAGE_CLASS_CHECKPOINT				,		output_type=UpperCamelCase__				,		config_class=_CONFIG_FOR_DOC				,		expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT				,		)
		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]				,		UpperCamelCase__      :     Optional[torch.FloatTensor] = None				,		UpperCamelCase__      :     Optional[torch.LongTensor] = None				,		UpperCamelCase__      :     Optional[bool] = None				,		UpperCamelCase__      :     Optional[bool] = None				,		)				->					ImageClassifierOutputWithNoAttention:
								'''simple docstring'''







								UpperCAmelCase         =     return_dict if return_dict is not None else self.config.use_return_dict

								UpperCAmelCase         =     self.resnet(UpperCamelCase__				,		output_hidden_states=UpperCamelCase__				,		return_dict=UpperCamelCase__      )

								UpperCAmelCase         =     outputs.pooler_output if return_dict else outputs[1]

								UpperCAmelCase         =     self.classifier(UpperCamelCase__      )

								UpperCAmelCase         =     None

								if labels is not None:
														if self.config.problem_type is None:
																				if self.num_labels == 1:
																										UpperCAmelCase         =     "regression"
																				elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																										UpperCAmelCase         =     "single_label_classification"
																				else:
																										UpperCAmelCase         =     "multi_label_classification"
														if self.config.problem_type == "regression":
																				UpperCAmelCase         =     MSELoss()
																				if self.num_labels == 1:
																										UpperCAmelCase         =     loss_fct(logits.squeeze()				,		labels.squeeze()      )
																				else:
																										UpperCAmelCase         =     loss_fct(UpperCamelCase__				,		UpperCamelCase__      )
														elif self.config.problem_type == "single_label_classification":
																				UpperCAmelCase         =     CrossEntropyLoss()
																				UpperCAmelCase         =     loss_fct(logits.view(-1				,		self.num_labels      )				,		labels.view(-1      )      )
														elif self.config.problem_type == "multi_label_classification":
																				UpperCAmelCase         =     BCEWithLogitsLoss()
																				UpperCAmelCase         =     loss_fct(UpperCamelCase__				,		UpperCamelCase__      )

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

								return ImageClassifierOutputWithNoAttention(loss=UpperCamelCase__				,		logits=UpperCamelCase__				,		hidden_states=outputs.hidden_states      )





@add_start_docstrings(
    '''
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    ''',  A__,  )
class  __magic_name__   (							A__,  A__						):






		def __init__(						self      :     List[Any]				,		UpperCamelCase__      :     Optional[int]      )				->					str:
								'''simple docstring'''







								super().__init__(UpperCamelCase__      )
								super()._init_backbone(UpperCamelCase__      )

								UpperCAmelCase         =     [config.embedding_size] + config.hidden_sizes
								UpperCAmelCase         =     ResNetEmbeddings(UpperCamelCase__      )
								UpperCAmelCase         =     ResNetEncoder(UpperCamelCase__      )

								# initialize weights and apply final processing
								self.post_init()







		@add_start_docstrings_to_model_forward(UpperCamelCase__      )
		@replace_return_docstrings(output_type=UpperCamelCase__				,		config_class=_CONFIG_FOR_DOC      )
		def SCREAMING_SNAKE_CASE_			(						self      :     Any				,		UpperCamelCase__      :     Tensor				,		UpperCamelCase__      :     Optional[bool] = None				,		UpperCamelCase__      :     Optional[bool] = None      )				->					BackboneOutput:
								'''simple docstring'''







								UpperCAmelCase         =     return_dict if return_dict is not None else self.config.use_return_dict
								UpperCAmelCase         =     (
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)

								UpperCAmelCase         =     self.embedder(UpperCamelCase__      )

								UpperCAmelCase         =     self.encoder(UpperCamelCase__				,		output_hidden_states=UpperCamelCase__				,		return_dict=UpperCamelCase__      )

								UpperCAmelCase         =     outputs.hidden_states

								UpperCAmelCase         =     ()
								for idx, stage in enumerate(self.stage_names      ):
														if stage in self.out_features:
																				feature_maps += (hidden_states[idx],)

								if not return_dict:
														UpperCAmelCase         =     (feature_maps,)
														if output_hidden_states:
																				output += (outputs.hidden_states,)
														return output

								return BackboneOutput(
								    feature_maps=UpperCamelCase__				,		hidden_states=outputs.hidden_states if output_hidden_states else None				,		attentions=UpperCamelCase__				,		)








import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
		import jax
		import jax.numpy as jnp
		from flax.jax_utils import replicate
		from flax.training.common_utils import shard





@slow
@require_flax
class  __magic_name__   (							unittest.TestCase						):






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[int]      )				->					Optional[Any]:
								'''simple docstring'''







								super().tearDown()
								gc.collect()






		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple      )				->					Dict:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase         =     FlaxControlNetModel.from_pretrained(
								    "lllyasviel/sd-controlnet-canny"				,		from_pt=UpperCamelCase__				,		dtype=jnp.bfloataa      )
								UpperCAmelCase       ,		UpperCAmelCase         =     FlaxStableDiffusionControlNetPipeline.from_pretrained(
								    "runwayml/stable-diffusion-v1-5"				,		controlnet=UpperCamelCase__				,		from_pt=UpperCamelCase__				,		dtype=jnp.bfloataa      )
								UpperCAmelCase         =     controlnet_params

								UpperCAmelCase         =     "bird"
								UpperCAmelCase         =     jax.device_count()
								UpperCAmelCase         =     pipe.prepare_text_inputs([prompts] * num_samples      )

								UpperCAmelCase         =     load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"      )
								UpperCAmelCase         =     pipe.prepare_image_inputs([canny_image] * num_samples      )

								UpperCAmelCase         =     jax.random.PRNGKey(0      )
								UpperCAmelCase         =     jax.random.split(UpperCamelCase__				,		jax.device_count()      )

								UpperCAmelCase         =     replicate(UpperCamelCase__      )
								UpperCAmelCase         =     shard(UpperCamelCase__      )
								UpperCAmelCase         =     shard(UpperCamelCase__      )

								UpperCAmelCase         =     pipe(
								    prompt_ids=UpperCamelCase__				,		image=UpperCamelCase__				,		params=UpperCamelCase__				,		prng_seed=UpperCamelCase__				,		num_inference_steps=50				,		jit=UpperCamelCase__				,		).images
								assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

								UpperCAmelCase         =     images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
								UpperCAmelCase         =     images[0, 2_53:2_56, 2_53:2_56, -1]

								UpperCAmelCase         =     jnp.asarray(jax.device_get(image_slice.flatten()      )      )
								UpperCAmelCase         =     jnp.array(
								    [0.16_79_69, 0.11_66_99, 0.08_15_43, 0.15_42_97, 0.13_28_12, 0.10_88_87, 0.16_99_22, 0.16_99_22, 0.20_50_78]      )
								print(F'output_slice: {output_slice}'      )
								assert jnp.abs(output_slice - expected_slice      ).max() < 1e-2







		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple      )				->					Tuple:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase         =     FlaxControlNetModel.from_pretrained(
								    "lllyasviel/sd-controlnet-openpose"				,		from_pt=UpperCamelCase__				,		dtype=jnp.bfloataa      )
								UpperCAmelCase       ,		UpperCAmelCase         =     FlaxStableDiffusionControlNetPipeline.from_pretrained(
								    "runwayml/stable-diffusion-v1-5"				,		controlnet=UpperCamelCase__				,		from_pt=UpperCamelCase__				,		dtype=jnp.bfloataa      )
								UpperCAmelCase         =     controlnet_params

								UpperCAmelCase         =     "Chef in the kitchen"
								UpperCAmelCase         =     jax.device_count()
								UpperCAmelCase         =     pipe.prepare_text_inputs([prompts] * num_samples      )

								UpperCAmelCase         =     load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png"      )
								UpperCAmelCase         =     pipe.prepare_image_inputs([pose_image] * num_samples      )

								UpperCAmelCase         =     jax.random.PRNGKey(0      )
								UpperCAmelCase         =     jax.random.split(UpperCamelCase__				,		jax.device_count()      )

								UpperCAmelCase         =     replicate(UpperCamelCase__      )
								UpperCAmelCase         =     shard(UpperCamelCase__      )
								UpperCAmelCase         =     shard(UpperCamelCase__      )

								UpperCAmelCase         =     pipe(
								    prompt_ids=UpperCamelCase__				,		image=UpperCamelCase__				,		params=UpperCamelCase__				,		prng_seed=UpperCamelCase__				,		num_inference_steps=50				,		jit=UpperCamelCase__				,		).images
								assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

								UpperCAmelCase         =     images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
								UpperCAmelCase         =     images[0, 2_53:2_56, 2_53:2_56, -1]

								UpperCAmelCase         =     jnp.asarray(jax.device_get(image_slice.flatten()      )      )
								UpperCAmelCase         =     jnp.array(
								    [[0.27_14_84, 0.26_17_19, 0.27_53_91, 0.27_73_44, 0.27_92_97, 0.29_10_16, 0.29_49_22, 0.30_27_34, 0.30_27_34]]      )
								print(F'output_slice: {output_slice}'      )
								assert jnp.abs(output_slice - expected_slice      ).max() < 1e-2



