






"""simple docstring"""







import argparse

import torch

from diffusers import UNetaDConditionModel






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase=None		,					lowerCAmelCase=None		,					lowerCAmelCase=None						):







       '''simple docstring'''
       assert isinstance(lowerCAmelCase		,					lowerCAmelCase						), "Paths should be a list of dicts containing 'old' and 'new' keys."

       # Splits the attention layers into three variables.
       if attention_paths_to_split is not None:
              for path, path_map in attention_paths_to_split.items():
                     UpperCAmelCase					=	old_checkpoint[path]
                     UpperCAmelCase					=	old_tensor.shape[0] // 3

                     UpperCAmelCase					=	(-1, channels) if len(old_tensor.shape						) == 3 else (-1)

                     UpperCAmelCase					=	old_tensor.shape[0] // config["""num_head_channels"""] // 3

                     UpperCAmelCase					=	old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]						)
                     UpperCAmelCase			,				UpperCAmelCase			,				UpperCAmelCase					=	old_tensor.split(channels // num_heads		,					dim=1						)

                     UpperCAmelCase					=	query.reshape(lowerCAmelCase						)
                     UpperCAmelCase					=	key.reshape(lowerCAmelCase						)
                     UpperCAmelCase					=	value.reshape(lowerCAmelCase						)

       for path in paths:
              UpperCAmelCase					=	path["""new"""]

              # These have already been assigned
              if attention_paths_to_split is not None and new_path in attention_paths_to_split:
                     continue

              if additional_replacements is not None:
                     for replacement in additional_replacements:
                            UpperCAmelCase					=	new_path.replace(replacement["""old"""]		,					replacement["""new"""]						)

        # proj_attn.weight has to be converted from conv 1D to linear
              UpperCAmelCase					=	old_checkpoint[path["""old"""]]
              UpperCAmelCase					=	["""proj_attn.weight"""]
              UpperCAmelCase					=	["""proj_out.weight""", """proj_in.weight"""]
              if any(k in new_path for k in names						):
                     UpperCAmelCase					=	weight[:, :, 0]
              elif any(k in new_path for k in names_a						) and len(weight.shape						) > 2 and ".attentions." not in new_path:
                     UpperCAmelCase					=	weight[:, :, 0]
              else:
                     UpperCAmelCase					=	weight






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase=0						):







       '''simple docstring'''
       UpperCAmelCase					=	[]
       for old_item in old_list:
              UpperCAmelCase					=	old_item

              #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
              #         new_item = new_item.replace('norm.bias', 'group_norm.bias')

              #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
              #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')

              #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)

              mapping.append({"""old""": old_item, """new""": new_item}						)

       return mapping






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase=1						):







       '''simple docstring'''
       if n_shave_prefix_segments >= 0:
              return ".".join(path.split("""."""						)[n_shave_prefix_segments:]						)
       else:
              return ".".join(path.split("""."""						)[:n_shave_prefix_segments]						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase=0						):







       '''simple docstring'''
       UpperCAmelCase					=	[]
       for old_item in old_list:
              mapping.append({"""old""": old_item, """new""": old_item}						)

       return mapping






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase=0						):







       '''simple docstring'''
       UpperCAmelCase					=	[]
       for old_item in old_list:
              UpperCAmelCase					=	old_item.replace("""in_layers.0"""		,					"""norm1"""						)
              UpperCAmelCase					=	new_item.replace("""in_layers.2"""		,					"""conv1"""						)

              UpperCAmelCase					=	new_item.replace("""out_layers.0"""		,					"""norm2"""						)
              UpperCAmelCase					=	new_item.replace("""out_layers.3"""		,					"""conv2"""						)

              UpperCAmelCase					=	new_item.replace("""emb_layers.1"""		,					"""time_emb_proj"""						)
              UpperCAmelCase					=	new_item.replace("""skip_connection"""		,					"""conv_shortcut"""						)

              UpperCAmelCase					=	shave_segments(lowerCAmelCase		,					n_shave_prefix_segments=lowerCAmelCase						)

              if "temopral_conv" not in old_item:
                     mapping.append({"""old""": old_item, """new""": new_item}						)

       return mapping






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase=None		,					lowerCAmelCase=False						):







       '''simple docstring'''
       UpperCAmelCase					=	{}
       UpperCAmelCase					=	list(checkpoint.keys()						)

       UpperCAmelCase					=	"""model.diffusion_model."""

       # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA
       if sum(k.startswith("""model_ema"""						) for k in keys						) > 100 and extract_ema:
              print(F'''Checkpoint {path} has both EMA and non-EMA weights.'''						)
              print(
                  """In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA"""
                  """ weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag."""						)
              for key in keys:
                     if key.startswith("""model.diffusion_model"""						):
                            UpperCAmelCase					=	"""model_ema.""" + """""".join(key.split("""."""						)[1:]						)
                            UpperCAmelCase					=	checkpoint.pop(lowerCAmelCase						)
       else:
              if sum(k.startswith("""model_ema"""						) for k in keys						) > 100:
                     print(
                         """In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA"""
                         """ weights (usually better for inference), please make sure to add the `--extract_ema` flag."""						)

              for key in keys:
                     UpperCAmelCase					=	checkpoint.pop(lowerCAmelCase						)

       UpperCAmelCase					=	{}

       UpperCAmelCase					=	unet_state_dict["""time_embed.0.weight"""]
       UpperCAmelCase					=	unet_state_dict["""time_embed.0.bias"""]
       UpperCAmelCase					=	unet_state_dict["""time_embed.2.weight"""]
       UpperCAmelCase					=	unet_state_dict["""time_embed.2.bias"""]

       if config["class_embed_type"] is None:
              # No parameters to port
              ...
       elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection":
              UpperCAmelCase					=	unet_state_dict["""label_emb.0.0.weight"""]
              UpperCAmelCase					=	unet_state_dict["""label_emb.0.0.bias"""]
              UpperCAmelCase					=	unet_state_dict["""label_emb.0.2.weight"""]
              UpperCAmelCase					=	unet_state_dict["""label_emb.0.2.bias"""]
       else:
              raise NotImplementedError(F'''Not implemented `class_embed_type`: {config['class_embed_type']}'''						)

       UpperCAmelCase					=	unet_state_dict["""input_blocks.0.0.weight"""]
       UpperCAmelCase					=	unet_state_dict["""input_blocks.0.0.bias"""]

       UpperCAmelCase					=	[v for v in unet_state_dict if v.startswith("""input_blocks.0.1"""						)]
       UpperCAmelCase					=	renew_attention_paths(lowerCAmelCase						)
       UpperCAmelCase					=	{"""old""": """input_blocks.0.1""", """new""": """transformer_in"""}
       assign_to_checkpoint(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

       UpperCAmelCase					=	unet_state_dict["""out.0.weight"""]
       UpperCAmelCase					=	unet_state_dict["""out.0.bias"""]
       UpperCAmelCase					=	unet_state_dict["""out.2.weight"""]
       UpperCAmelCase					=	unet_state_dict["""out.2.bias"""]

       # Retrieves the keys for the input blocks only
       UpperCAmelCase					=	len({""".""".join(layer.split("""."""						)[:2]						) for layer in unet_state_dict if """input_blocks""" in layer}						)
       UpperCAmelCase					=	{
           layer_id: [key for key in unet_state_dict if F'''input_blocks.{layer_id}''' in key]
           for layer_id in range(lowerCAmelCase						)
       }

       # Retrieves the keys for the middle blocks only
       UpperCAmelCase					=	len({""".""".join(layer.split("""."""						)[:2]						) for layer in unet_state_dict if """middle_block""" in layer}						)
       UpperCAmelCase					=	{
           layer_id: [key for key in unet_state_dict if F'''middle_block.{layer_id}''' in key]
           for layer_id in range(lowerCAmelCase						)
       }

       # Retrieves the keys for the output blocks only
       UpperCAmelCase					=	len({""".""".join(layer.split("""."""						)[:2]						) for layer in unet_state_dict if """output_blocks""" in layer}						)
       UpperCAmelCase					=	{
           layer_id: [key for key in unet_state_dict if F'''output_blocks.{layer_id}''' in key]
           for layer_id in range(lowerCAmelCase						)
       }

       for i in range(1		,					lowerCAmelCase						):
              UpperCAmelCase					=	(i - 1) // (config["""layers_per_block"""] + 1)
              UpperCAmelCase					=	(i - 1) % (config["""layers_per_block"""] + 1)

              UpperCAmelCase					=	[
                  key for key in input_blocks[i] if F'''input_blocks.{i}.0''' in key and F'''input_blocks.{i}.0.op''' not in key
              ]
              UpperCAmelCase					=	[key for key in input_blocks[i] if F'''input_blocks.{i}.1''' in key]
              UpperCAmelCase					=	[key for key in input_blocks[i] if F'''input_blocks.{i}.2''' in key]

              if F'''input_blocks.{i}.op.weight''' in unet_state_dict:
                     UpperCAmelCase					=	unet_state_dict.pop(
                         F'''input_blocks.{i}.op.weight'''						)
                     UpperCAmelCase					=	unet_state_dict.pop(
                         F'''input_blocks.{i}.op.bias'''						)

              UpperCAmelCase					=	renew_resnet_paths(lowerCAmelCase						)
              UpperCAmelCase					=	{"""old""": F'''input_blocks.{i}.0''', """new""": F'''down_blocks.{block_id}.resnets.{layer_in_block_id}'''}
              assign_to_checkpoint(
                  lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

              UpperCAmelCase					=	[key for key in resnets if """temopral_conv""" in key]
              UpperCAmelCase					=	renew_temp_conv_paths(lowerCAmelCase						)
              UpperCAmelCase					=	{
                  """old""": F'''input_blocks.{i}.0.temopral_conv''',
                  """new""": F'''down_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
              }
              assign_to_checkpoint(
                  lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

              if len(lowerCAmelCase						):
                     UpperCAmelCase					=	renew_attention_paths(lowerCAmelCase						)
                     UpperCAmelCase					=	{"""old""": F'''input_blocks.{i}.1''', """new""": F'''down_blocks.{block_id}.attentions.{layer_in_block_id}'''}
                     assign_to_checkpoint(
                         lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

              if len(lowerCAmelCase						):
                     UpperCAmelCase					=	renew_attention_paths(lowerCAmelCase						)
                     UpperCAmelCase					=	{
                         """old""": F'''input_blocks.{i}.2''',
                         """new""": F'''down_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
                     }
                     assign_to_checkpoint(
                         lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

       UpperCAmelCase					=	middle_blocks[0]
       UpperCAmelCase					=	[key for key in resnet_a if """temopral_conv""" in key]
       UpperCAmelCase					=	middle_blocks[1]
       UpperCAmelCase					=	middle_blocks[2]
       UpperCAmelCase					=	middle_blocks[3]
       UpperCAmelCase					=	[key for key in resnet_a if """temopral_conv""" in key]

       UpperCAmelCase					=	renew_resnet_paths(lowerCAmelCase						)
       UpperCAmelCase					=	{"""old""": """middle_block.0""", """new""": """mid_block.resnets.0"""}
       assign_to_checkpoint(
           lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					config=lowerCAmelCase		,					additional_replacements=[meta_path]						)

       UpperCAmelCase					=	renew_temp_conv_paths(lowerCAmelCase						)
       UpperCAmelCase					=	{"""old""": """middle_block.0.temopral_conv""", """new""": """mid_block.temp_convs.0"""}
       assign_to_checkpoint(
           lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					config=lowerCAmelCase		,					additional_replacements=[meta_path]						)

       UpperCAmelCase					=	renew_resnet_paths(lowerCAmelCase						)
       UpperCAmelCase					=	{"""old""": """middle_block.3""", """new""": """mid_block.resnets.1"""}
       assign_to_checkpoint(
           lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					config=lowerCAmelCase		,					additional_replacements=[meta_path]						)

       UpperCAmelCase					=	renew_temp_conv_paths(lowerCAmelCase						)
       UpperCAmelCase					=	{"""old""": """middle_block.3.temopral_conv""", """new""": """mid_block.temp_convs.1"""}
       assign_to_checkpoint(
           lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					config=lowerCAmelCase		,					additional_replacements=[meta_path]						)

       UpperCAmelCase					=	renew_attention_paths(lowerCAmelCase						)
       UpperCAmelCase					=	{"""old""": """middle_block.1""", """new""": """mid_block.attentions.0"""}
       assign_to_checkpoint(
           lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

       UpperCAmelCase					=	renew_attention_paths(lowerCAmelCase						)
       UpperCAmelCase					=	{"""old""": """middle_block.2""", """new""": """mid_block.temp_attentions.0"""}
       assign_to_checkpoint(
           lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

       for i in range(lowerCAmelCase						):
              UpperCAmelCase					=	i // (config["""layers_per_block"""] + 1)
              UpperCAmelCase					=	i % (config["""layers_per_block"""] + 1)
              UpperCAmelCase					=	[shave_segments(lowerCAmelCase		,					2						) for name in output_blocks[i]]
              UpperCAmelCase					=	{}

              for layer in output_block_layers:
                     UpperCAmelCase			,				UpperCAmelCase					=	layer.split("""."""						)[0], shave_segments(lowerCAmelCase		,					1						)
                     if layer_id in output_block_list:
                            output_block_list[layer_id].append(lowerCAmelCase						)
                     else:
                            UpperCAmelCase					=	[layer_name]

              if len(lowerCAmelCase						) > 1:
                     UpperCAmelCase					=	[key for key in output_blocks[i] if F'''output_blocks.{i}.0''' in key]
                     UpperCAmelCase					=	[key for key in output_blocks[i] if F'''output_blocks.{i}.1''' in key]
                     UpperCAmelCase					=	[key for key in output_blocks[i] if F'''output_blocks.{i}.2''' in key]

                     UpperCAmelCase					=	renew_resnet_paths(lowerCAmelCase						)
                     UpperCAmelCase					=	renew_resnet_paths(lowerCAmelCase						)

                     UpperCAmelCase					=	{"""old""": F'''output_blocks.{i}.0''', """new""": F'''up_blocks.{block_id}.resnets.{layer_in_block_id}'''}
                     assign_to_checkpoint(
                         lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

                     UpperCAmelCase					=	[key for key in resnets if """temopral_conv""" in key]
                     UpperCAmelCase					=	renew_temp_conv_paths(lowerCAmelCase						)
                     UpperCAmelCase					=	{
                         """old""": F'''output_blocks.{i}.0.temopral_conv''',
                         """new""": F'''up_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
                     }
                     assign_to_checkpoint(
                         lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

                     UpperCAmelCase					=	{k: sorted(lowerCAmelCase						) for k, v in output_block_list.items()}
                     if ["conv.bias", "conv.weight"] in output_block_list.values():
                            UpperCAmelCase					=	list(output_block_list.values()						).index(["""conv.bias""", """conv.weight"""]						)
                            UpperCAmelCase					=	unet_state_dict[
                                F'''output_blocks.{i}.{index}.conv.weight'''
                            ]
                            UpperCAmelCase					=	unet_state_dict[
                                F'''output_blocks.{i}.{index}.conv.bias'''
                            ]

                            # Clear attentions as they have been attributed above.
                            if len(lowerCAmelCase						) == 2:
                                   UpperCAmelCase					=	[]

                     if len(lowerCAmelCase						):
                            UpperCAmelCase					=	renew_attention_paths(lowerCAmelCase						)
                            UpperCAmelCase					=	{
                                """old""": F'''output_blocks.{i}.1''',
                                """new""": F'''up_blocks.{block_id}.attentions.{layer_in_block_id}''',
                            }
                            assign_to_checkpoint(
                                lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)

                     if len(lowerCAmelCase						):
                            UpperCAmelCase					=	renew_attention_paths(lowerCAmelCase						)
                            UpperCAmelCase					=	{
                                """old""": F'''output_blocks.{i}.2''',
                                """new""": F'''up_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
                            }
                            assign_to_checkpoint(
                                lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					additional_replacements=[meta_path]		,					config=lowerCAmelCase						)
              else:
                     UpperCAmelCase					=	renew_resnet_paths(lowerCAmelCase		,					n_shave_prefix_segments=1						)
                     for path in resnet_0_paths:
                            UpperCAmelCase					=	""".""".join(["""output_blocks""", str(lowerCAmelCase						), path["""old"""]]						)
                            UpperCAmelCase					=	""".""".join(["""up_blocks""", str(lowerCAmelCase						), """resnets""", str(lowerCAmelCase						), path["""new"""]]						)
                            UpperCAmelCase					=	unet_state_dict[old_path]

                     UpperCAmelCase					=	[l for l in output_block_layers if """temopral_conv""" in l]
                     for path in temopral_conv_paths:
                            UpperCAmelCase					=	path.split("""temopral_conv."""						)[-1]
                            UpperCAmelCase					=	""".""".join(["""output_blocks""", str(lowerCAmelCase						), str(lowerCAmelCase						), """temopral_conv""", pruned_path]						)
                            UpperCAmelCase					=	""".""".join(["""up_blocks""", str(lowerCAmelCase						), """temp_convs""", str(lowerCAmelCase						), pruned_path]						)
                            UpperCAmelCase					=	unet_state_dict[old_path]

       return new_checkpoint


if __name__ == "__main__":
 lowerCAmelCase_		:			List[str]				   =					argparse.ArgumentParser()

 parser.add_argument(
     '''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the checkpoint to convert.'''
 )
 parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')
 lowerCAmelCase_		:			List[Any]				   =					parser.parse_args()

 lowerCAmelCase_		:			Any				   =					torch.load(args.checkpoint_path, map_location='''cpu''')
 lowerCAmelCase_		:			Any				   =					UNetaDConditionModel()

 lowerCAmelCase_		:			Optional[int]				   =					convert_ldm_unet_checkpoint(unet_checkpoint, unet.config)

 lowerCAmelCase_		:			Any				   =					set(unet.state_dict().keys()) - set(converted_ckpt.keys())
 lowerCAmelCase_		:			Any				   =					set(converted_ckpt.keys()) - set(unet.state_dict().keys())

 assert len(diff_a) == len(diff_a) == 0, "Converted weights don't match"

 # load state_dict
 unet.load_state_dict(converted_ckpt)

 unet.save_pretrained(args.dump_path)

 # -- finish converting the unet --








"""simple docstring"""







import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class      UpperCamelCase_       (				a_      ,    a_      ,    unittest.TestCase     ):
  _A				:    str =		StableDiffusionParadigmsPipeline
  _A				:    Any =		TEXT_TO_IMAGE_PARAMS
  _A				:    List[Any] =		TEXT_TO_IMAGE_BATCH_PARAMS
  _A				:    str =		TEXT_TO_IMAGE_IMAGE_PARAMS
  _A				:    Tuple =		TEXT_TO_IMAGE_IMAGE_PARAMS



  def 						UpperCamelCase_ (					self )	->	Dict:



         """simple docstring"""





         torch.manual_seed(0 )
         UpperCAmelCase					=	UNetaDConditionModel(
             block_out_channels=(32, 64)       ,						layers_per_block=2       ,						sample_size=32       ,						in_channels=4       ,						out_channels=4       ,						down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")       ,						up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")       ,						cross_attention_dim=32       ,						attention_head_dim=(2, 4)       ,						use_linear_projection=snake_case__       ,						)
         UpperCAmelCase					=	DDIMParallelScheduler(
             beta_start=0.00_085       ,						beta_end=0.012       ,						beta_schedule="""scaled_linear"""       ,						clip_sample=snake_case__       ,						set_alpha_to_one=snake_case__       ,						)
         torch.manual_seed(0 )
         UpperCAmelCase					=	AutoencoderKL(
             block_out_channels=[32, 64]       ,						in_channels=3       ,						out_channels=3       ,						down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]       ,						up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]       ,						latent_channels=4       ,						sample_size=1_28       ,						)
         torch.manual_seed(0 )
         UpperCAmelCase					=	CLIPTextConfig(
             bos_token_id=0       ,						eos_token_id=2       ,						hidden_size=32       ,						intermediate_size=37       ,						layer_norm_eps=1e-05       ,						num_attention_heads=4       ,						num_hidden_layers=5       ,						pad_token_id=1       ,						vocab_size=10_00       ,						hidden_act="""gelu"""       ,						projection_dim=5_12       ,						)
         UpperCAmelCase					=	CLIPTextModel(snake_case__ )
         UpperCAmelCase					=	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" )

         UpperCAmelCase					=	{
             """unet""": unet,
             """scheduler""": scheduler,
             """vae""": vae,
             """text_encoder""": text_encoder,
             """tokenizer""": tokenizer,
             """safety_checker""": None,
             """feature_extractor""": None,
         }
         return components



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__=0 )	->	Any:



         """simple docstring"""





         if str(snake_case__ ).startswith("""mps""" ):
                UpperCAmelCase					=	torch.manual_seed(snake_case__ )
         else:
                UpperCAmelCase					=	torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )
         UpperCAmelCase					=	{
             """prompt""": """a photograph of an astronaut riding a horse""",
             """generator""": generator,
             """num_inference_steps""": 10,
             """guidance_scale""": 6.0,
             """output_type""": """numpy""",
             """parallel""": 3,
             """debug""": True,
         }
         return inputs



  def 						UpperCamelCase_ (					self )	->	Union[str, Any]:



         """simple docstring"""





         UpperCAmelCase					=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase					=	self.get_dummy_components()
         UpperCAmelCase					=	StableDiffusionParadigmsPipeline(**snake_case__ )
         UpperCAmelCase					=	sd_pipe.to(snake_case__ )
         sd_pipe.set_progress_bar_config(disable=snake_case__ )

         UpperCAmelCase					=	self.get_dummy_inputs(snake_case__ )
         UpperCAmelCase					=	sd_pipe(**snake_case__ ).images
         UpperCAmelCase					=	image[0, -3:, -3:, -1]
         assert image.shape == (1, 64, 64, 3)

         UpperCAmelCase					=	np.array([0.4_773, 0.5_417, 0.4_723, 0.4_925, 0.5_631, 0.4_752, 0.5_240, 0.4_935, 0.5_023] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2



  def 						UpperCamelCase_ (					self )	->	List[Any]:



         """simple docstring"""





         UpperCAmelCase					=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase					=	self.get_dummy_components()
         torch.manual_seed(0 )
         UpperCAmelCase					=	DDPMParallelScheduler()
         torch.manual_seed(0 )
         UpperCAmelCase					=	StableDiffusionParadigmsPipeline(**snake_case__ )
         UpperCAmelCase					=	sd_pipe.to(snake_case__ )
         sd_pipe.set_progress_bar_config(disable=snake_case__ )

         UpperCAmelCase					=	self.get_dummy_inputs(snake_case__ )
         UpperCAmelCase					=	sd_pipe(**snake_case__ ).images
         UpperCAmelCase					=	image[0, -3:, -3:, -1]
         assert image.shape == (1, 64, 64, 3)

         UpperCAmelCase					=	np.array([0.3_573, 0.4_420, 0.4_960, 0.4_799, 0.3_796, 0.3_879, 0.4_819, 0.4_365, 0.4_468] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2



  def 						UpperCamelCase_ (					self )	->	Tuple:



         """simple docstring"""





         super().test_inference_batch_consistent(batch_sizes=[1, 2] )



  def 						UpperCamelCase_ (					self )	->	Optional[int]:



         """simple docstring"""





         super().test_inference_batch_single_identical(batch_size=2       ,						expected_max_diff=3e-3 )



  def 						UpperCamelCase_ (					self )	->	Union[str, Any]:



         """simple docstring"""





         UpperCAmelCase					=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase					=	self.get_dummy_components()
         UpperCAmelCase					=	StableDiffusionParadigmsPipeline(**snake_case__ )
         UpperCAmelCase					=	sd_pipe.to(snake_case__ )
         sd_pipe.set_progress_bar_config(disable=snake_case__ )

         UpperCAmelCase					=	self.get_dummy_inputs(snake_case__ )
         UpperCAmelCase					=	"""french fries"""
         UpperCAmelCase					=	sd_pipe(**snake_case__       ,						negative_prompt=snake_case__ )
         UpperCAmelCase					=	output.images
         UpperCAmelCase					=	image[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)

         UpperCAmelCase					=	np.array([0.4_771, 0.5_420, 0.4_683, 0.4_918, 0.5_636, 0.4_725, 0.5_230, 0.4_923, 0.5_015] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2





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



  def 						UpperCamelCase_ (					self )	->	Union[str, Any]:



         """simple docstring"""





         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()



  def 						UpperCamelCase_ (					self       ,						snake_case__=0 )	->	int:



         """simple docstring"""





         UpperCAmelCase					=	torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )
         UpperCAmelCase					=	{
             """prompt""": """a photograph of an astronaut riding a horse""",
             """generator""": generator,
             """num_inference_steps""": 10,
             """guidance_scale""": 7.5,
             """output_type""": """numpy""",
             """parallel""": 3,
             """debug""": True,
         }
         return inputs



  def 						UpperCamelCase_ (					self )	->	Tuple:



         """simple docstring"""





         UpperCAmelCase					=	"""stabilityai/stable-diffusion-2-base"""
         UpperCAmelCase					=	DDIMParallelScheduler.from_pretrained(snake_case__       ,						subfolder="""scheduler""" )
         UpperCAmelCase					=	StableDiffusionParadigmsPipeline.from_pretrained(snake_case__       ,						scheduler=snake_case__       ,						safety_checker=snake_case__ )
         pipe.to(snake_case__ )
         pipe.set_progress_bar_config(disable=snake_case__ )
         pipe.enable_attention_slicing()

         UpperCAmelCase					=	self.get_inputs()
         UpperCAmelCase					=	pipe(**snake_case__ ).images
         UpperCAmelCase					=	image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 5_12, 5_12, 3)

         UpperCAmelCase					=	np.array([0.9_622, 0.9_602, 0.9_748, 0.9_591, 0.9_630, 0.9_691, 0.9_661, 0.9_631, 0.9_741] )

         assert np.abs(expected_slice - image_slice ).max() < 1e-2
