




"""simple docstring"""


import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, 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 					lowercase(     unittest.TestCase ):







 '''simple docstring'''





 def  UpperCamelCase_   (				self:		List[str]		):




   '''simple docstring'''



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


 def  UpperCamelCase_   (				self:		List[Any]		):




   '''simple docstring'''



   _snake_case		,					_snake_case							:		Union[str, Any]								=  FlaxStableDiffusionPipeline.from_pretrained(
       """stabilityai/stable-diffusion-2""",      revision="""bf16""",      dtype=jnp.bfloataa,      )

   _snake_case							:		str								=  """A painting of a squirrel eating a burger"""
   _snake_case							:		Union[str, Any]								=  jax.device_count()
   _snake_case							:		str								=  num_samples * [prompt]
   _snake_case							:		Union[str, Any]								=  sd_pipe.prepare_inputs(a_		)

   _snake_case							:		List[Any]								=  replicate(a_		)
   _snake_case							:		str								=  shard(a_		)

   _snake_case							:		List[Any]								=  jax.random.PRNGKey(0		)
   _snake_case							:		Tuple								=  jax.random.split(a_,      jax.device_count()		)

   _snake_case							:		Union[str, Any]								=  sd_pipe(a_,      a_,      a_,      num_inference_steps=25,      jit=a_		)[0]
   assert images.shape == (jax.device_count(), 1, 768, 768, 3)

   _snake_case							:		Tuple								=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
   _snake_case							:		List[Any]								=  images[0, 253:256, 253:256, -1]

   _snake_case							:		Optional[Any]								=  jnp.asarray(jax.device_get(image_slice.flatten()		)		)
   _snake_case							:		Any								=  jnp.array([0.4_238, 0.4_414, 0.4_395, 0.4_453, 0.4_629, 0.4_590, 0.4_531, 0.45_508, 0.4_512]		)
   print(f"output_slice: {output_slice}"		)
   assert jnp.abs(output_slice - expected_slice		).max() < 1E-2


 def  UpperCamelCase_   (				self:		Tuple		):




   '''simple docstring'''



   _snake_case							:		str								=  """stabilityai/stable-diffusion-2"""
   _snake_case		,					_snake_case							:		Dict								=  FlaxDPMSolverMultistepScheduler.from_pretrained(a_,      subfolder="""scheduler"""		)
   _snake_case		,					_snake_case							:		Any								=  FlaxStableDiffusionPipeline.from_pretrained(
       a_,      scheduler=a_,      revision="""bf16""",      dtype=jnp.bfloataa,      )
   _snake_case							:		Tuple								=  scheduler_params

   _snake_case							:		Dict								=  """A painting of a squirrel eating a burger"""
   _snake_case							:		str								=  jax.device_count()
   _snake_case							:		List[Any]								=  num_samples * [prompt]
   _snake_case							:		str								=  sd_pipe.prepare_inputs(a_		)

   _snake_case							:		List[str]								=  replicate(a_		)
   _snake_case							:		Tuple								=  shard(a_		)

   _snake_case							:		Union[str, Any]								=  jax.random.PRNGKey(0		)
   _snake_case							:		Dict								=  jax.random.split(a_,      jax.device_count()		)

   _snake_case							:		str								=  sd_pipe(a_,      a_,      a_,      num_inference_steps=25,      jit=a_		)[0]
   assert images.shape == (jax.device_count(), 1, 768, 768, 3)

   _snake_case							:		Tuple								=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
   _snake_case							:		List[Any]								=  images[0, 253:256, 253:256, -1]

   _snake_case							:		Optional[int]								=  jnp.asarray(jax.device_get(image_slice.flatten()		)		)
   _snake_case							:		Dict								=  jnp.array([0.4_336, 0.42_969, 0.4_453, 0.4_199, 0.4_297, 0.4_531, 0.4_434, 0.4_434, 0.4_297]		)
   print(f"output_slice: {output_slice}"		)
   assert jnp.abs(output_slice - expected_slice		).max() < 1E-2









import argparse
from collections import defaultdict

import yaml


A    :					str =      '''docs/source/en/_toctree.yml'''







def    __lowerCamelCase					(							__a   :str    )		->				List[Any]:






		"""simple docstring"""

		A__        =   defaultdict(__a    )
		A__        =   []
		A__        =   []
		for doc in doc_list:
				if "local" in doc:
						counts[doc["local"]] += 1

				if doc["title"].lower() == "overview":
						overview_doc.append({"""local""": doc["""local"""], """title""": doc["""title"""]}    )
				else:
						new_doc_list.append(__a    )

		A__        =   new_doc_list
		A__        =   [key for key, value in counts.items() if value > 1]

		A__        =   []
		for duplicate_key in duplicates:
				A__        =   list({doc["""title"""] for doc in doc_list if doc["""local"""] == duplicate_key}    )
				if len(__a    ) > 1:
						raise ValueError(
						    F'{duplicate_key} is present several times in the documentation table of content at '
						    """`docs/source/en/_toctree.yml` with different *Title* values. Choose one of those and remove the """
						    """others."""    )
				# Only add this once
				new_doc.append({"""local""": duplicate_key, """title""": titles[0]}    )

		# Add none duplicate-keys
		new_doc.extend([doc for doc in doc_list if """local""" not in counts or counts[doc["""local"""]] == 1]    )
		A__        =   sorted(__a    ,     key=lambda __a    : s["title"].lower()    )

		# "overview" gets special treatment and is always first
		if len(__a    ) > 1:
				raise ValueError("""{doc_list} has two 'overview' docs which is not allowed."""    )

		overview_doc.extend(__a    )

		# Sort
		return overview_doc







def    __lowerCamelCase					(							__a   :Any=False    )		->				List[str]:






		"""simple docstring"""

		with open(__a    ,     encoding="""utf-8"""    ) as f:
				A__        =   yaml.safe_load(f.read()    )

		# Get to the API doc
		A__        =   0
		while content[api_idx]["title"] != "API":
				api_idx += 1
		A__        =   content[api_idx]["""sections"""]

		# Then to the model doc
		A__        =   0
		while api_doc[scheduler_idx]["title"] != "Schedulers":
				scheduler_idx += 1

		A__        =   api_doc[scheduler_idx]["""sections"""]
		A__        =   clean_doc_toc(__a    )

		A__        =   False
		if new_scheduler_doc != scheduler_doc:
				A__        =   True
				if overwrite:
						A__        =   new_scheduler_doc

		if diff:
				if overwrite:
						A__        =   api_doc
						with open(__a    ,     """w"""    ,     encoding="""utf-8"""    ) as f:
								f.write(yaml.dump(__a    ,     allow_unicode=__a    )    )
				else:
						raise ValueError(
						    """The model doc part of the table of content is not properly sorted, run `make style` to fix this."""    )







def    __lowerCamelCase					(							__a   :Optional[int]=False    )		->				Dict:






		"""simple docstring"""

		with open(__a    ,     encoding="""utf-8"""    ) as f:
				A__        =   yaml.safe_load(f.read()    )

		# Get to the API doc
		A__        =   0
		while content[api_idx]["title"] != "API":
				api_idx += 1
		A__        =   content[api_idx]["""sections"""]

		# Then to the model doc
		A__        =   0
		while api_doc[pipeline_idx]["title"] != "Pipelines":
				pipeline_idx += 1

		A__        =   False
		A__        =   api_doc[pipeline_idx]["""sections"""]
		A__        =   []

		# sort sub pipeline docs
		for pipeline_doc in pipeline_docs:
				if "section" in pipeline_doc:
						A__        =   pipeline_doc["""section"""]
						A__        =   clean_doc_toc(__a    )
						if overwrite:
								A__        =   new_sub_pipeline_doc
				new_pipeline_docs.append(__a    )

		# sort overall pipeline doc
		A__        =   clean_doc_toc(__a    )

		if new_pipeline_docs != pipeline_docs:
				A__        =   True
				if overwrite:
						A__        =   new_pipeline_docs

		if diff:
				if overwrite:
						A__        =   api_doc
						with open(__a    ,     """w"""    ,     encoding="""utf-8"""    ) as f:
								f.write(yaml.dump(__a    ,     allow_unicode=__a    )    )
				else:
						raise ValueError(
						    """The model doc part of the table of content is not properly sorted, run `make style` to fix this."""    )


if __name__ == "__main__":
					A    :					Tuple =      argparse.ArgumentParser()
					parser.add_argument('''--fix_and_overwrite''', action='''store_true''', help='''Whether to fix inconsistencies.''')
					A    :					Optional[Any] =      parser.parse_args()

					check_scheduler_doc(args.fix_and_overwrite)
					check_pipeline_doc(args.fix_and_overwrite)



