code: 
  
  
  '''simple docstring'''
  
  
  
  import argparse
  import pathlib
  from pathlib import Path
  from tempfile import TemporaryDirectory
  
  import esm as esm_module
  import torch
  from esm.esmfold.va.misc import batch_encode_sequences as esmfold_encode_sequences
  from esm.esmfold.va.pretrained import esmfold_va
  
  from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
  from transformers.models.esm.modeling_esm import (
      EsmForMaskedLM,
      EsmForSequenceClassification,
      EsmIntermediate,
      EsmLayer,
      EsmOutput,
      EsmSelfAttention,
      EsmSelfOutput,
  )
  from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
  from transformers.models.esm.tokenization_esm import EsmTokenizer
  from transformers.utils import logging
  
  
  logging.set_verbosity_info()
  __lowercase			:		Any =			logging.get_logger(__name__)
  
  __lowercase			:		Union[str, Any] =			[
      (
          'protein1',
          'MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA',
      ),
      ('protein2', 'MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA'),
      ('protein3', 'MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG'),
      ('protein4', 'MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA'),
  ]
  
  __lowercase			:		Tuple =			{
      'esm1b_t33_650M_UR50S': esm_module.pretrained.esm1b_t33_650M_UR50S,
      'esm1v_t33_650M_UR90S_1': esm_module.pretrained.esm1v_t33_650M_UR90S_1,
      'esm1v_t33_650M_UR90S_2': esm_module.pretrained.esm1v_t33_650M_UR90S_2,
      'esm1v_t33_650M_UR90S_3': esm_module.pretrained.esm1v_t33_650M_UR90S_3,
      'esm1v_t33_650M_UR90S_4': esm_module.pretrained.esm1v_t33_650M_UR90S_4,
      'esm1v_t33_650M_UR90S_5': esm_module.pretrained.esm1v_t33_650M_UR90S_5,
      'esm2_t48_15B_UR50D': esm_module.pretrained.esm2_t48_15B_UR50D,
      'esm2_t36_3B_UR50D': esm_module.pretrained.esm2_t36_3B_UR50D,
      'esm2_t33_650M_UR50D': esm_module.pretrained.esm2_t33_650M_UR50D,
      'esm2_t30_150M_UR50D': esm_module.pretrained.esm2_t30_150M_UR50D,
      'esm2_t12_35M_UR50D': esm_module.pretrained.esm2_t12_35M_UR50D,
      'esm2_t6_8M_UR50D': esm_module.pretrained.esm2_t6_8M_UR50D,
      'esmfold_v1': esmfold_va,
  }
  
  __lowercase			:		Dict =			list('ARNDCQEGHILKMFPSTWYV')
  
  __lowercase			:		int =			restypes + ['X']
  __lowercase			:		str =			restypes_with_x + ['<pad>', '<mask>', '<cls>', '<sep>', '<eos>']
  
  
  
  
  
  
  def 						lowerCamelCase							():
  					with TemporaryDirectory() as tempdir:
  										__a    : List[str]      					=						'\n'.join(_SCREAMING_SNAKE_CASE     )
  										__a    : Any      					=						Path(_SCREAMING_SNAKE_CASE     ) / 'vocab.txt'
  										vocab_file.write_text(_SCREAMING_SNAKE_CASE     )
  										__a    : int      					=						EsmTokenizer(vocab_file=str(_SCREAMING_SNAKE_CASE     )     )
  					__a    : Union[str, Any]      					=						0  # Overlaps with 'A' but that seems to be what they want
  					return hf_tokenizer
  
  
  
  
  
  
  def 						lowerCamelCase							(_SCREAMING_SNAKE_CASE   :     List[Any]		,   _SCREAMING_SNAKE_CASE   :     List[Any]     ):
  					__a    : Optional[Any]      					=						our_module.load_state_dict(original_module.state_dict()     )
  					if status.missing_keys:
  										raise ValueError(F"""Missing keys: {status.missing_keys}"""     )
  					if status.unexpected_keys:
  										raise ValueError(F"""Unexpected keys: {status.unexpected_keys}"""     )
  
  
  
  
  
  
  def 						lowerCamelCase							(_SCREAMING_SNAKE_CASE   :     str		,   _SCREAMING_SNAKE_CASE   :     str		,   _SCREAMING_SNAKE_CASE   :     bool		,   _SCREAMING_SNAKE_CASE   :     str		,   _SCREAMING_SNAKE_CASE   :     str     ):
  					if model.startswith('esmfold'     ):
  										__a    : Union[str, Any]      					=						MODEL_MAPPING[model]()
  					else:
  										__a    : Optional[Any]      					=						MODEL_MAPPING[model]()
  					esm.eval()  # disable dropout
  
  					if model.startswith('esmfold'     ):
  										__a    : int      					=						esm.esm.embed_dim
  										__a    : Any      					=						esm.esm.num_layers
  										__a    : List[Any]      					=						esm.esm.attention_heads
  										__a    : Optional[int]      					=						4 * embed_dim
  										__a    : List[str]      					=						esm.esm.token_dropout
  										__a    : str      					=						False  # This code path does not exist in ESM-2
  										__a    : Dict      					=						'rotary'
  										__a    : Any      					=						True
  										__a    : Optional[Any]      					=						EsmFoldConfig()
  										for key, val in esm.cfg.items():
  															if hasattr(_SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE     ) and key != "trunk":
  																				setattr(_SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE     )
  										for key, val in esm.cfg.trunk.items():
  															if hasattr(esmfold_config.trunk		,   _SCREAMING_SNAKE_CASE     ) and key != "structure_module":
  																				setattr(esmfold_config.trunk		,   _SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE     )
  										for key, val in esm.cfg.trunk.structure_module.items():
  															if hasattr(esmfold_config.trunk.structure_module		,   _SCREAMING_SNAKE_CASE     ):
  																				setattr(esmfold_config.trunk.structure_module		,   _SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE     )
  					elif hasattr(_SCREAMING_SNAKE_CASE		,   'args'     ):
  										# Indicates an ESM-1b or ESM-1v model
  										__a    : Optional[Any]      					=						esm.args.embed_dim
  										__a    : Dict      					=						esm.args.layers
  										__a    : List[str]      					=						esm.args.attention_heads
  										__a    : int      					=						esm.args.ffn_embed_dim
  										__a    : str      					=						esm.args.token_dropout
  										__a    : Union[str, Any]      					=						True if esm.emb_layer_norm_before else False
  										__a    : Dict      					=						'absolute'
  										__a    : Union[str, Any]      					=						False
  										__a    : List[Any]      					=						None
  					else:
  										# Indicates an ESM-2 model
  										__a    : str      					=						esm.embed_dim
  										__a    : Union[str, Any]      					=						esm.num_layers
  										__a    : Any      					=						esm.attention_heads
  										__a    : Union[str, Any]      					=						4 * embed_dim  # This is hardcoded in ESM-2
  										__a    : Union[str, Any]      					=						esm.token_dropout
  										__a    : List[str]      					=						False  # This code path does not exist in ESM-2
  										__a    : Any      					=						'rotary'
  										__a    : List[Any]      					=						False
  										__a    : Optional[int]      					=						None
  
  					if is_folding_model:
  										__a    : str      					=						esm.esm.alphabet
  					__a    : Any      					=						tuple(alphabet.all_toks     )
  					__a    : Optional[Any]      					=						alphabet.mask_idx
  					__a    : Dict      					=						alphabet.padding_idx
  
  					if is_folding_model:
  										__a    : Any      					=						esm.esm
  					else:
  										__a    : Optional[Any]      					=						esm
  
  					__a    : Any      					=						EsmConfig(
  					    vocab_size=original_esm_model.embed_tokens.num_embeddings		,   mask_token_id=_SCREAMING_SNAKE_CASE		,   hidden_size=_SCREAMING_SNAKE_CASE		,   num_hidden_layers=_SCREAMING_SNAKE_CASE		,   num_attention_heads=_SCREAMING_SNAKE_CASE		,   intermediate_size=_SCREAMING_SNAKE_CASE		,   max_position_embeddings=1_026		,   layer_norm_eps=1e-5		,   attention_probs_dropout_prob=0.0		,   hidden_dropout_prob=0.0		,   pad_token_id=_SCREAMING_SNAKE_CASE		,   emb_layer_norm_before=_SCREAMING_SNAKE_CASE		,   token_dropout=_SCREAMING_SNAKE_CASE		,   position_embedding_type=_SCREAMING_SNAKE_CASE		,   is_folding_model=_SCREAMING_SNAKE_CASE		,   esmfold_config=_SCREAMING_SNAKE_CASE		,   vocab_list=_SCREAMING_SNAKE_CASE		,   )
  					if classification_head:
  										__a    : Union[str, Any]      					=						esm.classification_heads['mnli'].out_proj.weight.shape[0]
  					print('Our ESM config:'		,   _SCREAMING_SNAKE_CASE     )
  
  					if model.startswith('esmfold'     ):
  										__a    : List[Any]      					=						EsmForProteinFolding
  					elif classification_head:
  										__a    : Tuple      					=						EsmForSequenceClassification
  					else:
  										__a    : Dict      					=						EsmForMaskedLM
  					__a    : int      					=						model_class(_SCREAMING_SNAKE_CASE     )
  					model.eval()
  
  					# Now let's copy all the weights.
  					# Embeddings
  					__a    : str      					=						original_esm_model.embed_tokens.weight
  					if position_embedding_type == "absolute":
  										__a    : Union[str, Any]      					=						original_esm_model.embed_positions.weight
  
  					if config.emb_layer_norm_before:
  										__a    : Any      					=						original_esm_model.emb_layer_norm_before.weight
  										__a    : Dict      					=						original_esm_model.emb_layer_norm_before.bias
  
  					__a    : List[Any]      					=						original_esm_model.emb_layer_norm_after.weight
  					__a    : List[str]      					=						original_esm_model.emb_layer_norm_after.bias
  
  					for i in range(config.num_hidden_layers     ):
  										# Encoder: start of layer
  										__a    : EsmLayer      					=						model.esm.encoder.layer[i]
  										# esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
  										__a    : str      					=						original_esm_model.layers[i]
  
  										# self attention
  										__a    : EsmSelfAttention      					=						layer.attention.self
  										assert (
  										    esm_layer.self_attn.k_proj.weight.data.shape
  										    == esm_layer.self_attn.q_proj.weight.data.shape
  										    == esm_layer.self_attn.v_proj.weight.data.shape
  										    == torch.Size((config.hidden_size, config.hidden_size)     )
  										)
  
  										__a    : Union[str, Any]      					=						esm_layer.self_attn.q_proj.weight
  										__a    : Any      					=						esm_layer.self_attn.q_proj.bias
  										__a    : List[Any]      					=						esm_layer.self_attn.k_proj.weight
  										__a    : Union[str, Any]      					=						esm_layer.self_attn.k_proj.bias
  										__a    : Optional[Any]      					=						esm_layer.self_attn.v_proj.weight
  										__a    : List[str]      					=						esm_layer.self_attn.v_proj.bias
  
  										if getattr(esm_layer.self_attn		,   'rot_emb'		,   _SCREAMING_SNAKE_CASE     ) is not None:
  															# Matt: Although inv_freq is not a trainable weight, it is computed at model init and cached.
  															# During the training of ESM-2 the model was converted to float16 precision, which also converts
  															# the inv_freq tensor, and the loss of precision remains even if the model is loaded later as float32.
  															# If we recompute inv_freq without this loss of precision then we will get subtly different rotary
  															# embeddings, which are enough to cause significant discrepancies in model outputs. To avoid this,
  															# we make sure the new model copies the data from the old inv_freq.
  															__a    : Optional[int]      					=						esm_layer.self_attn.rot_emb.inv_freq
  
  										# LayerNorm changes for pre-activation
  										__a    : str      					=						esm_layer.self_attn_layer_norm.weight
  										__a    : Dict      					=						esm_layer.self_attn_layer_norm.bias
  										__a    : Optional[Any]      					=						esm_layer.final_layer_norm.weight
  										__a    : List[str]      					=						esm_layer.final_layer_norm.bias
  
  										# self-attention output
  										__a    : EsmSelfOutput      					=						layer.attention.output
  										assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
  										__a    : int      					=						esm_layer.self_attn.out_proj.weight
  										__a    : int      					=						esm_layer.self_attn.out_proj.bias
  
  										# intermediate
  										__a    : EsmIntermediate      					=						layer.intermediate
  										assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
  										__a    : int      					=						esm_layer.fca.weight
  										__a    : Tuple      					=						esm_layer.fca.bias
  
  										# output
  										__a    : EsmOutput      					=						layer.output
  										assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
  										__a    : Tuple      					=						esm_layer.fca.weight
  										__a    : int      					=						esm_layer.fca.bias
  										# end of layer
  
  					if is_folding_model:
  										__a    : Optional[int]      					=						esm.esm_s_combine.data
  										__a    : Dict      					=						esm.afa_to_esm.data
  										transfer_and_check_weights(esm.embedding		,   model.embedding     )
  										transfer_and_check_weights(esm.esm_s_mlp		,   model.esm_s_mlp     )
  										transfer_and_check_weights(esm.trunk		,   model.trunk     )
  										transfer_and_check_weights(esm.distogram_head		,   model.distogram_head     )
  										transfer_and_check_weights(esm.ptm_head		,   model.ptm_head     )
  										transfer_and_check_weights(esm.lm_head		,   model.lm_head     )
  										transfer_and_check_weights(esm.lddt_head		,   model.lddt_head     )
  
  					elif classification_head:
  										__a    : int      					=						esm.esm.classification_heads['mnli'].dense.weight
  										__a    : Dict      					=						esm.classification_heads['mnli'].dense.bias
  										__a    : Union[str, Any]      					=						esm.classification_heads['mnli'].out_proj.weight
  										__a    : Any      					=						esm.classification_heads['mnli'].out_proj.bias
  					else:
  										# LM Head
  										__a    : Dict      					=						esm.lm_head.dense.weight
  										__a    : Optional[Any]      					=						esm.lm_head.dense.bias
  										__a    : List[Any]      					=						esm.lm_head.layer_norm.weight
  										__a    : Tuple      					=						esm.lm_head.layer_norm.bias
  										__a    : Union[str, Any]      					=						esm.lm_head.weight
  										__a    : Any      					=						esm.lm_head.bias
  
  					# Contact prediction head
  					transfer_and_check_weights(esm.contact_head		,   model.esm.contact_head     )
  
  					# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
  					if is_folding_model:
  										# Folding models aren't trained on masked inputs and don't like mask tokens.
  										__a    : List[Any]      					=						SAMPLE_DATA[:2]
  					else:
  										__a    : List[str]      					=						SAMPLE_DATA
  
  					if is_folding_model:
  										__a    : Dict      					=						get_esmfold_tokenizer()
  										__a    : Optional[Any]      					=						hf_tokenizer(
  										    [row[1] for row in sample_data]		,   return_tensors='pt'		,   padding=_SCREAMING_SNAKE_CASE		,   add_special_tokens=_SCREAMING_SNAKE_CASE     )
  										__a    : List[Any]      					=						esmfold_encode_sequences([row[1] for row in sample_data]     )
  										__a    : Tuple      					=						torch.all(hf_tokens['input_ids'] == esmfold_aas     ) and torch.all(
  										    hf_tokens['attention_mask'] == esmfold_mask     )
  					else:
  										# Let's check that we get the same results.
  										__a    : Tuple      					=						alphabet.get_batch_converter()
  										__a    : Any      					=						batch_converter(_SCREAMING_SNAKE_CASE     )
  										# Prepare tokenizer and make sure it matches
  										with TemporaryDirectory() as tempdir:
  															__a    : Optional[int]      					=						'\n'.join(alphabet.all_toks     )
  															__a    : List[str]      					=						Path(_SCREAMING_SNAKE_CASE     ) / 'vocab.txt'
  															vocab_file.write_text(_SCREAMING_SNAKE_CASE     )
  															__a    : Dict      					=						EsmTokenizer(vocab_file=str(_SCREAMING_SNAKE_CASE     )     )
  
  										__a    : Optional[Any]      					=						hf_tokenizer([row[1] for row in sample_data]		,   return_tensors='pt'		,   padding=_SCREAMING_SNAKE_CASE     )
  										__a    : Union[str, Any]      					=						torch.all(hf_tokens['input_ids'] == batch_tokens     )
  
  					print('Do both models tokenizers output the same tokens?'		,   '🔥' if success else '💩'     )
  					if not success:
  										raise Exception('Tokenization does not match!'     )
  
  					with torch.no_grad():
  										if is_folding_model:
  															# Let's test the model in parts
  															# ESMFold always converts the ESM stem to float16, which requires float16 ops
  															# that don't exist on CPU. Therefore, to test it we need to run it on GPU. However,
  															# ESMFold is what we in the community call a "big boy" and so we desperately avoid putting both the
  															# original and the converted model on the GPU at the same time.
  															__a    : Union[str, Any]      					=						esm.cuda().infer([row[1] for row in sample_data]     )
  															__a    : List[Any]      					=						model.cuda()(
  															    input_ids=hf_tokens['input_ids'].cuda()		,   attention_mask=hf_tokens['attention_mask'].cuda()     )
  										else:
  															__a    : Optional[int]      					=						model(**_SCREAMING_SNAKE_CASE		,   output_hidden_states=_SCREAMING_SNAKE_CASE     )
  															__a    : Optional[int]      					=						our_output['logits']
  															if classification_head:
  																				__a    : str      					=						esm.model.classification_heads['mnli'](esm.extract_features(_SCREAMING_SNAKE_CASE     )     )
  															else:
  																				__a    : str      					=						esm(hf_tokens['input_ids']		,   repr_layers=list(range(999     )     )     )
  																				__a    : Any      					=						their_output['logits']
  
  										if is_folding_model:
  															__a    : Optional[Any]      					=						torch.max(torch.abs(our_output['positions'] - their_output['positions']     )     ).item()
  															__a    : Any      					=						torch.allclose(our_output['positions']		,   their_output['positions']		,   atol=1e-5     )
  										else:
  															__a    : Optional[int]      					=						torch.max(torch.abs(our_output - their_output     )     ).item()
  															__a    : Optional[int]      					=						torch.allclose(_SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE		,   atol=1e-5     )
  
  										print(F"""max_absolute_diff = {max_absolute_diff}"""     )  # ~ 1e-5
  										print('Do both models output the same tensors?'		,   '🔥' if success else '💩'     )
  
  										if not success:
  															raise Exception('Something went wRoNg'     )
  
  										if not is_folding_model:
  															# Let's check contact prediction too
  															__a    : List[Any]      					=						model.predict_contacts(hf_tokens['input_ids']		,   hf_tokens['attention_mask']     )
  															__a    : Tuple      					=						esm.predict_contacts(hf_tokens['input_ids']     )
  															__a    : Any      					=						torch.max(torch.abs(our_output - their_output     )     ).item()
  															__a    : Union[str, Any]      					=						torch.allclose(_SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE		,   atol=1e-5     )
  
  															print('Contact prediction testing:'     )
  															print(F"""max_absolute_diff = {max_absolute_diff}"""     )  # ~ 1e-5
  															print('Do both models output the same tensors?'		,   '🔥' if success else '💩'     )
  
  															if not success:
  																				raise Exception('Something went wRoNg'     )
  
  										pathlib.Path(_SCREAMING_SNAKE_CASE     ).mkdir(parents=_SCREAMING_SNAKE_CASE		,   exist_ok=_SCREAMING_SNAKE_CASE     )
  										print(F"""Saving model to {pytorch_dump_folder_path}"""     )
  										model.save_pretrained(_SCREAMING_SNAKE_CASE     )
  
  										del esm  # Free up some memory before continuing
  
  					print(F"""Saving tokenizer to {pytorch_dump_folder_path}"""     )
  					hf_tokenizer.save_pretrained(_SCREAMING_SNAKE_CASE     )
  
  					if push_to_repo:
  										model.push_to_hub(repo_id=_SCREAMING_SNAKE_CASE		,   use_auth_token=_SCREAMING_SNAKE_CASE     )
  										hf_tokenizer.push_to_hub(repo_id=_SCREAMING_SNAKE_CASE		,   use_auth_token=_SCREAMING_SNAKE_CASE     )
  
  
  if __name__ == "__main__":
  	__lowercase			:		Union[str, Any] =			argparse.ArgumentParser()
  	# Required parameters
  	parser.add_argument(
  	    '--pytorch_dump_folder_path', type=str, required=True, help='Path to the output PyTorch model.'
  	)
  	parser.add_argument(
  	    '--classification_head', action='store_true', help='Whether to convert a final classification head.'
  	)
  	parser.add_argument('--model', default=None, type=str, required=True, help='Name of model to convert.')
  	parser.add_argument('--push_to_repo', type=str, help='Repo to upload to (including username!).')
  	parser.add_argument('--auth_token', type=str, help='HuggingFace auth token.')
  	__lowercase			:		int =			parser.parse_args()
  	convert_esm_checkpoint_to_pytorch(
  	    args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
  	)
  
style_context: 
  
  
  '''simple docstring'''
  
  
  
  from typing import Callable, List, Optional, Union
  
  import torch
  
  from ...models import UNetaDModel
  from ...schedulers import CMStochasticIterativeScheduler
  from ...utils import (
      is_accelerate_available,
      is_accelerate_version,
      logging,
      randn_tensor,
      replace_example_docstring,
  )
  from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
  
  
  __lowercase			:		List[str] =			logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  __lowercase			:		Optional[int] =			'\n    Examples:\n        ```py\n        >>> import torch\n\n        >>> from diffusers import ConsistencyModelPipeline\n\n        >>> device = "cuda"\n        >>> # Load the cd_imagenet64_l2 checkpoint.\n        >>> model_id_or_path = "openai/diffusers-cd_imagenet64_l2"\n        >>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)\n        >>> pipe.to(device)\n\n        >>> # Onestep Sampling\n        >>> image = pipe(num_inference_steps=1).images[0]\n        >>> image.save("cd_imagenet64_l2_onestep_sample.png")\n\n        >>> # Onestep sampling, class-conditional image generation\n        >>> # ImageNet-64 class label 145 corresponds to king penguins\n        >>> image = pipe(num_inference_steps=1, class_labels=145).images[0]\n        >>> image.save("cd_imagenet64_l2_onestep_sample_penguin.png")\n\n        >>> # Multistep sampling, class-conditional image generation\n        >>> # Timesteps can be explicitly specified; the particular timesteps below are from the original Github repo:\n        >>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77\n        >>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]\n        >>> image.save("cd_imagenet64_l2_multistep_sample_penguin.png")\n        ```\n'
  
  
  
  
  
  class 							__UpperCamelCase       (	lowerCAmelCase_      ):
  
  
  
  
  	def __init__(     self		,       __a		,       __a							):
  
  						'''simple docstring'''
  
  
  
  
  
  						super().__init__()
  
  						self.register_modules(
  						    unet=__a		,       scheduler=__a		,       )
  
  						__a    : List[str]      					=						None
  
  
  
  
  	def    __UpperCAmelCase							(     self		,       __a=0							):
  
  						'''simple docstring'''
  
  
  
  
  
  						if is_accelerate_available() and is_accelerate_version('>='		,       '0.14.0'							):
  											from accelerate import cpu_offload
  						else:
  											raise ImportError('`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher'							)
  
  						__a    : str      					=						torch.device(f"""cuda:{gpu_id}"""							)
  
  						if self.device.type != "cpu":
  											self.to('cpu'		,       silence_dtype_warnings=__a							)
  											torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
  
  						for cpu_offloaded_model in [self.unet]:
  											cpu_offload(__a		,       __a							)
  
  						if self.safety_checker is not None:
  											cpu_offload(self.safety_checker		,       execution_device=__a		,       offload_buffers=__a							)
  
  
  
  
  	def    __UpperCAmelCase							(     self		,       __a=0							):
  
  						'''simple docstring'''
  
  
  
  
  
  						if is_accelerate_available() and is_accelerate_version('>='		,       '0.17.0.dev0'							):
  											from accelerate import cpu_offload_with_hook
  						else:
  											raise ImportError('`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'							)
  
  						__a    : Tuple      					=						torch.device(f"""cuda:{gpu_id}"""							)
  
  						if self.device.type != "cpu":
  											self.to('cpu'		,       silence_dtype_warnings=__a							)
  											torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
  
  						__a    : List[Any]      					=						None
  						for cpu_offloaded_model in [self.unet]:
  											__a					,       __a    : List[str]      					=						cpu_offload_with_hook(__a		,       __a		,       prev_module_hook=__a							)
  
  						if self.safety_checker is not None:
  											__a					,       __a    : str      					=						cpu_offload_with_hook(self.safety_checker		,       __a		,       prev_module_hook=__a							)
  
  						# We'll offload the last model manually.
  						__a    : List[Any]      					=						hook
  
  
  
  
  	@property
  	# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  	def    __UpperCAmelCase							(     self							):
  
  						'''simple docstring'''
  
  
  
  
  
  						if not hasattr(self.unet		,       '_hf_hook'							):
  											return self.device
  						for module in self.unet.modules():
  											if (
  											    hasattr(__a		,       '_hf_hook'							)
  											    and hasattr(module._hf_hook		,       'execution_device'							)
  											    and module._hf_hook.execution_device is not None
  											):
  																return torch.device(module._hf_hook.execution_device							)
  						return self.device
  
  
  
  
  	def    __UpperCAmelCase							(     self		,       __a		,       __a		,       __a		,       __a		,       __a		,       __a		,       __a		,       __a=None							):
  
  						'''simple docstring'''
  
  
  
  
  
  						__a    : str      					=						(batch_size, num_channels, height, width)
  						if isinstance(__a		,       __a							) and len(__a							) != batch_size:
  											raise ValueError(
  											    f"""You have passed a list of generators of length {len(__a							)}, but requested an effective batch"""
  											    f""" size of {batch_size}. Make sure the batch size matches the length of the generators."""							)
  
  						if latents is None:
  											__a    : Optional[Any]      					=						randn_tensor(__a		,       generator=__a		,       device=__a		,       dtype=__a							)
  						else:
  											__a    : Dict      					=						latents.to(device=__a		,       dtype=__a							)
  
  						# scale the initial noise by the standard deviation required by the scheduler
  						__a    : Union[str, Any]      					=						latents * self.scheduler.init_noise_sigma
  						return latents
  
  
  
  
  	def    __UpperCAmelCase							(     self		,       __a		,       __a = "pil"							):
  
  						'''simple docstring'''
  
  
  
  
  
  						if output_type not in ["pt", "np", "pil"]:
  											raise ValueError(
  											    f"""output_type={output_type} is not supported. Make sure to choose one of ['pt', 'np', or 'pil']"""							)
  
  						# Equivalent to diffusers.VaeImageProcessor.denormalize
  						__a    : List[Any]      					=						(sample / 2 + 0.5).clamp(0		,       1							)
  						if output_type == "pt":
  											return sample
  
  						# Equivalent to diffusers.VaeImageProcessor.pt_to_numpy
  						__a    : str      					=						sample.cpu().permute(0		,       2		,       3		,       1							).numpy()
  						if output_type == "np":
  											return sample
  
  						# Output_type must be 'pil'
  						__a    : Union[str, Any]      					=						self.numpy_to_pil(__a							)
  						return sample
  
  
  
  
  	def    __UpperCAmelCase							(     self		,       __a		,       __a		,       __a=None							):
  
  						'''simple docstring'''
  
  
  
  
  
  						if self.unet.config.num_class_embeds is not None:
  											if isinstance(__a		,       __a							):
  																__a    : Tuple      					=						torch.tensor(__a		,       dtype=torch.int							)
  											elif isinstance(__a		,       __a							):
  																assert batch_size == 1, "Batch size must be 1 if classes is an int"
  																__a    : int      					=						torch.tensor([class_labels]		,       dtype=torch.int							)
  											elif class_labels is None:
  																# Randomly generate batch_size class labels
  																# TODO: should use generator here? int analogue of randn_tensor is not exposed in ...utils
  																__a    : str      					=						torch.randint(0		,       self.unet.config.num_class_embeds		,       size=(batch_size,)							)
  											__a    : str      					=						class_labels.to(__a							)
  						else:
  											__a    : Optional[int]      					=						None
  						return class_labels
  
  
  
  
  	def    __UpperCAmelCase							(     self		,       __a		,       __a		,       __a		,       __a		,       __a		,       __a							):
  
  						'''simple docstring'''
  
  
  
  
  
  						if num_inference_steps is None and timesteps is None:
  											raise ValueError('Exactly one of `num_inference_steps` or `timesteps` must be supplied.'							)
  
  						if num_inference_steps is not None and timesteps is not None:
  											logger.warning(
  											    f"""Both `num_inference_steps`: {num_inference_steps} and `timesteps`: {timesteps} are supplied;"""
  											    ' `timesteps` will be used over `num_inference_steps`.'							)
  
  						if latents is not None:
  											__a    : Optional[Any]      					=						(batch_size, 3, img_size, img_size)
  											if latents.shape != expected_shape:
  																raise ValueError(f"""The shape of latents is {latents.shape} but is expected to be {expected_shape}."""							)
  
  						if (callback_steps is None) or (
  						    callback_steps is not None and (not isinstance(__a		,       __a							) or callback_steps <= 0)
  						):
  											raise ValueError(
  											    f"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
  											    f""" {type(__a							)}."""							)
  
  
  
  
  
  	@torch.no_grad()
  	@replace_example_docstring(__a							)
  	def __call__(     self		,       __a = 1		,       __a = None		,       __a = 1		,       __a = None		,       __a = None		,       __a = None		,       __a = "pil"		,       __a = True		,       __a = None		,       __a = 1		,       ):
  
  						'''simple docstring'''
  
  
  
  
  
  						__a    : Dict      					=						self.unet.config.sample_size
  						__a    : Union[str, Any]      					=						self._execution_device
  
  						# 1. Check inputs
  						self.check_inputs(__a		,       __a		,       __a		,       __a		,       __a		,       __a							)
  
  						# 2. Prepare image latents
  						# Sample image latents x_0 ~ N(0, sigma_0^2 * I)
  						__a    : Union[str, Any]      					=						self.prepare_latents(
  						    batch_size=__a		,       num_channels=self.unet.config.in_channels		,       height=__a		,       width=__a		,       dtype=self.unet.dtype		,       device=__a		,       generator=__a		,       latents=__a		,       )
  
  						# 3. Handle class_labels for class-conditional models
  						__a    : Optional[Any]      					=						self.prepare_class_labels(__a		,       __a		,       class_labels=__a							)
  
  						# 4. Prepare timesteps
  						if timesteps is not None:
  											self.scheduler.set_timesteps(timesteps=__a		,       device=__a							)
  											__a    : Union[str, Any]      					=						self.scheduler.timesteps
  											__a    : Tuple      					=						len(__a							)
  						else:
  											self.scheduler.set_timesteps(__a							)
  											__a    : Tuple      					=						self.scheduler.timesteps
  
  						# 5. Denoising loop
  						# Multistep sampling: implements Algorithm 1 in the paper
  						with self.progress_bar(total=__a							) as progress_bar:
  											for i, t in enumerate(__a							):
  																__a    : Optional[int]      					=						self.scheduler.scale_model_input(__a		,       __a							)
  																__a    : List[str]      					=						self.unet(__a		,       __a		,       class_labels=__a		,       return_dict=__a							)[0]
  
  																__a    : int      					=						self.scheduler.step(__a		,       __a		,       __a		,       generator=__a							)[0]
  
  																# call the callback, if provided
  																progress_bar.update()
  																if callback is not None and i % callback_steps == 0:
  																					callback(__a		,       __a		,       __a							)
  
          # 6. Post-process image sample
  						__a    : Tuple      					=						self.postprocess_image(__a		,       output_type=__a							)
  
  						# Offload last model to CPU
  						if hasattr(self		,       'final_offload_hook'							) and self.final_offload_hook is not None:
  											self.final_offload_hook.offload()
  
  						if not return_dict:
  											return (image,)
  
  						return ImagePipelineOutput(images=__a							)
  
  
  
  