





'''simple docstring'''






import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()
class 					__snake_case  (					__SCREAMING_SNAKE_CASE		,				__SCREAMING_SNAKE_CASE		,				__SCREAMING_SNAKE_CASE		,				unittest.TestCase							):






   '''simple docstring'''





   lowerCamelCase__    							=      StableUnCLIPPipeline
   lowerCamelCase__    							=      TEXT_TO_IMAGE_PARAMS
   lowerCamelCase__    							=      TEXT_TO_IMAGE_BATCH_PARAMS
   lowerCamelCase__    							=      TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCamelCase__    							=      TEXT_TO_IMAGE_IMAGE_PARAMS

   # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
   lowerCamelCase__    							=      False

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					List[str]    =  3_2
    snake_case__     :					List[Any]    =  embedder_hidden_size

    # prior components

    torch.manual_seed(0       )
    snake_case__     :					str    =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

    torch.manual_seed(0       )
    snake_case__     :					Optional[int]    =  CLIPTextModelWithProjection(
        CLIPTextConfig(
            bos_token_id=0	,		eos_token_id=2	,		hidden_size=__SCREAMING_SNAKE_CASE	,		projection_dim=__SCREAMING_SNAKE_CASE	,		intermediate_size=3_7	,		layer_norm_eps=1e-0_5	,		num_attention_heads=4	,		num_hidden_layers=5	,		pad_token_id=1	,		vocab_size=1_0_0_0	,		)       )

    torch.manual_seed(0       )
    snake_case__     :					int    =  PriorTransformer(
        num_attention_heads=2	,		attention_head_dim=1_2	,		embedding_dim=__SCREAMING_SNAKE_CASE	,		num_layers=1	,		)

    torch.manual_seed(0       )
    snake_case__     :					Tuple    =  DDPMScheduler(
        variance_type="""fixed_small_log"""	,		prediction_type="""sample"""	,		num_train_timesteps=1_0_0_0	,		clip_sample=__SCREAMING_SNAKE_CASE	,		clip_sample_range=5.0	,		beta_schedule="""squaredcos_cap_v2"""	,		)

    # regular denoising components

    torch.manual_seed(0       )
    snake_case__     :					List[str]    =  StableUnCLIPImageNormalizer(embedding_dim=__SCREAMING_SNAKE_CASE       )
    snake_case__     :					str    =  DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""       )

    torch.manual_seed(0       )
    snake_case__     :					List[Any]    =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

    torch.manual_seed(0       )
    snake_case__     :					Optional[Any]    =  CLIPTextModel(
        CLIPTextConfig(
            bos_token_id=0	,		eos_token_id=2	,		hidden_size=__SCREAMING_SNAKE_CASE	,		projection_dim=3_2	,		intermediate_size=3_7	,		layer_norm_eps=1e-0_5	,		num_attention_heads=4	,		num_hidden_layers=5	,		pad_token_id=1	,		vocab_size=1_0_0_0	,		)       )

    torch.manual_seed(0       )
    snake_case__     :					Dict    =  UNetaDConditionModel(
        sample_size=3_2	,		in_channels=4	,		out_channels=4	,		down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")	,		up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")	,		block_out_channels=(3_2, 6_4)	,		attention_head_dim=(2, 4)	,		class_embed_type="""projection"""	,		projection_class_embeddings_input_dim=embedder_projection_dim * 2	,		cross_attention_dim=__SCREAMING_SNAKE_CASE	,		layers_per_block=1	,		upcast_attention=__SCREAMING_SNAKE_CASE	,		use_linear_projection=__SCREAMING_SNAKE_CASE	,		)

    torch.manual_seed(0       )
    snake_case__     :					List[Any]    =  DDIMScheduler(
        beta_schedule="""scaled_linear"""	,		beta_start=0.0_0085	,		beta_end=0.012	,		prediction_type="""v_prediction"""	,		set_alpha_to_one=__SCREAMING_SNAKE_CASE	,		steps_offset=1	,		)

    torch.manual_seed(0       )
    snake_case__     :					Dict    =  AutoencoderKL()

    snake_case__     :					Optional[int]    =  {
        # prior components
        """prior_tokenizer""": prior_tokenizer,
        """prior_text_encoder""": prior_text_encoder,
        """prior""": prior,
        """prior_scheduler""": prior_scheduler,
        # image noising components
        """image_normalizer""": image_normalizer,
        """image_noising_scheduler""": image_noising_scheduler,
        # regular denoising components
        """tokenizer""": tokenizer,
        """text_encoder""": text_encoder,
        """unet""": unet,
        """scheduler""": scheduler,
        """vae""": vae,
    }

    return components

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE=0       ):
    if str(__SCREAMING_SNAKE_CASE       ).startswith("""mps"""       ):
     snake_case__     :					List[str]    =  torch.manual_seed(__SCREAMING_SNAKE_CASE       )
    else:
     snake_case__     :					str    =  torch.Generator(device=__SCREAMING_SNAKE_CASE       ).manual_seed(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Any    =  {
        """prompt""": """A painting of a squirrel eating a burger""",
        """generator""": generator,
        """num_inference_steps""": 2,
        """prior_num_inference_steps""": 2,
        """output_type""": """numpy""",
    }
    return inputs

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					List[str]    =  torch_device == """cpu"""

    self._test_attention_slicing_forward_pass(test_max_difference=__SCREAMING_SNAKE_CASE       )


   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Any    =  torch_device in ["""cpu""", """mps"""]

    self._test_inference_batch_single_identical(test_max_difference=__SCREAMING_SNAKE_CASE       )



@slow
@require_torch_gpu
class 					__snake_case  (					unittest.TestCase							):






   '''simple docstring'''





   def 		__UpperCamelCase						(							self       ):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Tuple    =  load_numpy(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""       )

    snake_case__     :					Optional[Any]    =  StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""	,		torch_dtype=torch.floataa       )
    pipe.to(__SCREAMING_SNAKE_CASE       )
    pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )
    # stable unclip will oom when integration tests are run on a V100,
    # so turn on memory savings
    pipe.enable_attention_slicing()
    pipe.enable_sequential_cpu_offload()

    snake_case__     :					Optional[Any]    =  torch.Generator(device="""cpu"""       ).manual_seed(0       )
    snake_case__     :					int    =  pipe("""anime turle"""	,		generator=__SCREAMING_SNAKE_CASE	,		output_type="""np"""       )

    snake_case__     :					Union[str, Any]    =  output.images[0]

    assert image.shape == (7_6_8, 7_6_8, 3)

    assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )


   def 		__UpperCamelCase						(							self       ):
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

    snake_case__     :					Union[str, Any]    =  StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""	,		torch_dtype=torch.floataa       )
    snake_case__     :					Any    =  pipe.to(__SCREAMING_SNAKE_CASE       )
    pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )
    pipe.enable_attention_slicing()
    pipe.enable_sequential_cpu_offload()

    snake_case__     :					int    =  pipe(
        """anime turtle"""	,		prior_num_inference_steps=2	,		num_inference_steps=2	,		output_type="""np"""	,		)

    snake_case__     :					Union[str, Any]    =  torch.cuda.max_memory_allocated()
    # make sure that less than 7 GB is allocated
    assert mem_bytes < 7 * 1_0**9







"""simple docstring"""





def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	if number < 0 or shift_amount < 0:
		raise ValueError('''both inputs must be positive integers'''					)

	lowerCAmelCase__      :		int    							=      str(bin(A_					)					)
	binary_number += "0" * shift_amount
	return binary_number

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	if number < 0 or shift_amount < 0:
		raise ValueError('''both inputs must be positive integers'''					)

	lowerCAmelCase__      :		Any    							=      str(bin(A_					)					)[2:]
	if shift_amount >= len(A_					):
		return "0b0"
	lowerCAmelCase__      :		Tuple    							=      binary_number[: len(A_					) - shift_amount]
	return "0b" + shifted_binary_number

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	if number >= 0:  # Get binary representation of positive number
		lowerCAmelCase__      :		Union[str, Any]    							=      '''0''' + str(bin(A_					)					).strip('''-'''					)[2:]
	else:  # Get binary (2's complement) representation of negative number
		lowerCAmelCase__      :		Optional[Any]    							=      len(bin(A_					)[3:]					)  # Find 2's complement of number
		lowerCAmelCase__      :		Optional[int]    							=      bin(abs(A_					) - (1 << binary_number_length)					)[3:]
		lowerCAmelCase__      :		int    							=      (
		    '''1''' + '''0''' * (binary_number_length - len(A_					)) + binary_number
		)

	if shift_amount >= len(A_					):
		return "0b" + binary_number[0] * len(A_					)
	return (
	    "0b"
	    + binary_number[0] * shift_amount
	    + binary_number[: len(A_					) - shift_amount]
	)


if __name__ == "__main__":
							import doctest

							doctest.testmod()




