
'''simple docstring'''
from collections import defaultdict


class       _lowerCAmelCase   :







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  ,  UpperCAmelCase  )					->       Any:
														_snake_case       =  total  # total no of tasks (N)

														# DP table will have a dimension of (2^M)*N
														# initially all values are set to -1
														_snake_case       =  [
														    [-1 for i in range(total + 1  )] for j in range(2 ** len(UpperCAmelCase  )  )
														]

														_snake_case       =  defaultdict(UpperCAmelCase  )  # stores the list of persons for each task

														# final_mask is used to check if all persons are included by setting all bits
														# to 1
														_snake_case       =  (1 << len(UpperCAmelCase  )) - 1

							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase  )					->       List[Any]:
														# if mask == self.finalmask all persons are distributed tasks, return 1
														if mask == self.final_mask:
																					return 1

														# if not everyone gets the task and no more tasks are available, return 0
														if task_no > self.total_tasks:
																					return 0

														# if case already considered
														if self.dp[mask][task_no] != -1:
																					return self.dp[mask][task_no]

														# Number of ways when we don't this task in the arrangement
														_snake_case       =  self.count_ways_until(UpperCAmelCase  ,  task_no + 1  )

														# now assign the tasks one by one to all possible persons and recursively
														# assign for the remaining tasks.
														if task_no in self.task:
																					for p in self.task[task_no]:
																												# if p is already given a task
																												if mask & (1 << p):
																																			continue

																												# assign this task to p and change the mask value. And recursively
																												# assign tasks with the new mask value.
																												total_ways_util += self.count_ways_until(mask | (1 << p)  ,  task_no + 1  )

        # save the value.
														_snake_case       =  total_ways_util

														return self.dp[mask][task_no]





							def 	lowercase							(self  ,  UpperCAmelCase  )					->       Tuple:
														# Store the list of persons for each task
														for i in range(len(UpperCAmelCase  )  ):
																					for j in task_performed[i]:
																												self.task[j].append(UpperCAmelCase  )

        # call the function to fill the DP table, final answer is stored in dp[0][1]
														return self.count_ways_until(0  ,  1  )


if __name__ == "__main__":
				__lowerCAmelCase												=	5  # total no of tasks (the value of N)

				# the list of tasks that can be done by M persons.
				__lowerCAmelCase												=	[[1, 3, 4], [1, 2, 5], [3, 4]]
				print(
				    AssignmentUsingBitmask(task_performed, total_tasks).count_no_of_ways(
				        task_performed
				    )
				)

'''simple docstring'''
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       _lowerCAmelCase   (   unittest.TestCase				):







							'''simple docstring'''

							def 	lowercase							(self  )					->       Optional[int]:
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()

							def 	lowercase							(self  )					->       Dict:
														_snake_case,							_snake_case       =  FlaxControlNetModel.from_pretrained(
														    """lllyasviel/sd-controlnet-canny"""  ,  from_pt=UpperCAmelCase  ,  dtype=jnp.bfloataa  )
														_snake_case,							_snake_case       =  FlaxStableDiffusionControlNetPipeline.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""  ,  controlnet=UpperCAmelCase  ,  from_pt=UpperCAmelCase  ,  dtype=jnp.bfloataa  )
														_snake_case       =  controlnet_params

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

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

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

														_snake_case       =  replicate(UpperCAmelCase  )
														_snake_case       =  shard(UpperCAmelCase  )
														_snake_case       =  shard(UpperCAmelCase  )

														_snake_case       =  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, 768, 512, 3)

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

														_snake_case       =  jnp.asarray(jax.device_get(image_slice.flatten()  )  )
														_snake_case       =  jnp.array(
														    [0.16_7969, 0.11_6699, 0.08_1543, 0.15_4297, 0.13_2812, 0.10_8887, 0.16_9922, 0.16_9922, 0.20_5078]  )
														print(f"""output_slice: {output_slice}"""  )
														assert jnp.abs(output_slice - expected_slice  ).max() < 1e-2





							def 	lowercase							(self  )					->       Optional[int]:
														_snake_case,							_snake_case       =  FlaxControlNetModel.from_pretrained(
														    """lllyasviel/sd-controlnet-openpose"""  ,  from_pt=UpperCAmelCase  ,  dtype=jnp.bfloataa  )
														_snake_case,							_snake_case       =  FlaxStableDiffusionControlNetPipeline.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""  ,  controlnet=UpperCAmelCase  ,  from_pt=UpperCAmelCase  ,  dtype=jnp.bfloataa  )
														_snake_case       =  controlnet_params

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

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

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

														_snake_case       =  replicate(UpperCAmelCase  )
														_snake_case       =  shard(UpperCAmelCase  )
														_snake_case       =  shard(UpperCAmelCase  )

														_snake_case       =  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, 768, 512, 3)

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

														_snake_case       =  jnp.asarray(jax.device_get(image_slice.flatten()  )  )
														_snake_case       =  jnp.array(
														    [[0.27_1484, 0.26_1719, 0.27_5391, 0.27_7344, 0.27_9297, 0.29_1016, 0.29_4922, 0.30_2734, 0.30_2734]]  )
														print(f"""output_slice: {output_slice}"""  )
														assert jnp.abs(output_slice - expected_slice  ).max() < 1e-2