





'''simple docstring'''



import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


_SCREAMING_SNAKE_CASE   =      logging.get_logger(__name__)  # pylint: disable=invalid-name






class 							__lowercase    ( lowerCAmelCase__  ):







	'''simple docstring'''

	def __init__(self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,)							->							List[str]:







				'''simple docstring'''





				super().__init__()
				self.register_modules(
				    vae=_lowerCamelCase       ,text_encoder=_lowerCamelCase       ,tokenizer=_lowerCamelCase       ,unet=_lowerCamelCase       ,scheduler=_lowerCamelCase       ,safety_checker=_lowerCamelCase       ,feature_extractor=_lowerCamelCase       ,)


	def   _UpperCAmelCase       (self       ,_lowerCamelCase = "auto"	)							->							Optional[Any]:







				'''simple docstring'''





				if slice_size == "auto":
							# half the attention head size is usually a good trade-off between
							# speed and memory
							__lowercase			  =							self.unet.config.attention_head_dim // 2
				self.unet.set_attention_slice(_lowerCamelCase	)


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







				'''simple docstring'''





				self.enable_attention_slicing(_lowerCamelCase	)


	@torch.no_grad()
	def __call__(self       ,_lowerCamelCase       ,_lowerCamelCase = 512       ,_lowerCamelCase = 512       ,_lowerCamelCase = 50       ,_lowerCamelCase = 7.5       ,_lowerCamelCase = None       ,_lowerCamelCase = 1       ,_lowerCamelCase = 0.0       ,_lowerCamelCase = None       ,_lowerCamelCase = None       ,_lowerCamelCase = "pil"       ,_lowerCamelCase = True       ,_lowerCamelCase = None       ,_lowerCamelCase = 1       ,_lowerCamelCase = None       ,**_lowerCamelCase       ,)							->							Union[str, Any]:







				'''simple docstring'''





				if isinstance(_lowerCamelCase       ,_lowerCamelCase	):
							__lowercase			  =							1
				elif isinstance(_lowerCamelCase       ,_lowerCamelCase	):
							__lowercase			  =							len(_lowerCamelCase	)
				else:
							raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(_lowerCamelCase	)}"	)

				if height % 8 != 0 or width % 8 != 0:
							raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}."	)

				if (callback_steps is None) or (
				    callback_steps is not None and (not isinstance(_lowerCamelCase       ,_lowerCamelCase	) or callback_steps <= 0)
				):
							raise ValueError(
							    f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
							    f" {type(_lowerCamelCase	)}."	)

				# get prompt text embeddings
				__lowercase			  =							self.tokenizer(
				    _lowerCamelCase       ,padding='''max_length'''       ,max_length=self.tokenizer.model_max_length       ,return_tensors='''pt'''       ,)
				__lowercase			  =							text_inputs.input_ids

				if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
							__lowercase			  =							self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]	)
							logger.warning(
							    '''The following part of your input was truncated because CLIP can only handle sequences up to'''
							    f" {self.tokenizer.model_max_length} tokens: {removed_text}"	)
							__lowercase			  =							text_input_ids[:, : self.tokenizer.model_max_length]

				if text_embeddings is None:
							__lowercase			  =							self.text_encoder(text_input_ids.to(self.device	)	)[0]

				# duplicate text embeddings for each generation per prompt, using mps friendly method
				__lowercase  ,						__lowercase  ,						__lowercase			  =							text_embeddings.shape
				__lowercase			  =							text_embeddings.repeat(1       ,_lowerCamelCase       ,1	)
				__lowercase			  =							text_embeddings.view(bs_embed * num_images_per_prompt       ,_lowerCamelCase       ,-1	)

				# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
				# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
				# corresponds to doing no classifier free guidance.
				__lowercase			  =							guidance_scale > 1.0
				# get unconditional embeddings for classifier free guidance
				if do_classifier_free_guidance:
							__lowercase			  =							42
							if negative_prompt is None:
										__lowercase			  =							['''''']
							elif type(_lowerCamelCase	) is not type(_lowerCamelCase	):
										raise TypeError(
										    f"`negative_prompt` should be the same type to `prompt`, but got {type(_lowerCamelCase	)} !="
										    f" {type(_lowerCamelCase	)}."	)
							elif isinstance(_lowerCamelCase       ,_lowerCamelCase	):
										__lowercase			  =							[negative_prompt]
							elif batch_size != len(_lowerCamelCase	):
										raise ValueError(
										    f"`negative_prompt`: {negative_prompt} has batch size {len(_lowerCamelCase	)}, but `prompt`:"
										    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
										    ''' the batch size of `prompt`.'''	)
							else:
										__lowercase			  =							negative_prompt

							__lowercase			  =							text_input_ids.shape[-1]
							__lowercase			  =							self.tokenizer(
							    _lowerCamelCase       ,padding='''max_length'''       ,max_length=_lowerCamelCase       ,truncation=_lowerCamelCase       ,return_tensors='''pt'''       ,)
							__lowercase			  =							self.text_encoder(uncond_input.input_ids.to(self.device	)	)[0]

							# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
							__lowercase			  =							uncond_embeddings.shape[1]
							__lowercase			  =							uncond_embeddings.repeat(_lowerCamelCase       ,_lowerCamelCase       ,1	)
							__lowercase			  =							uncond_embeddings.view(batch_size * num_images_per_prompt       ,_lowerCamelCase       ,-1	)

							# For classifier free guidance, we need to do two forward passes.
							# Here we concatenate the unconditional and text embeddings into a single batch
							# to avoid doing two forward passes
							__lowercase			  =							torch.cat([uncond_embeddings, text_embeddings]	)

				# get the initial random noise unless the user supplied it

				# Unlike in other pipelines, latents need to be generated in the target device
				# for 1-to-1 results reproducibility with the CompVis implementation.
				# However this currently doesn't work in `mps`.
				__lowercase			  =							(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
				__lowercase			  =							(batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
				__lowercase			  =							text_embeddings.dtype
				if latents is None:
							if self.device.type == "mps":
										# randn does not exist on mps
										__lowercase			  =							torch.randn(
										    _lowerCamelCase       ,generator=_lowerCamelCase       ,device='''cpu'''       ,dtype=_lowerCamelCase	).to(self.device	)
										__lowercase			  =							torch.randn(_lowerCamelCase       ,generator=_lowerCamelCase       ,device='''cpu'''       ,dtype=_lowerCamelCase	).to(
										    self.device	)
							else:
										__lowercase			  =							torch.randn(
										    _lowerCamelCase       ,generator=_lowerCamelCase       ,device=self.device       ,dtype=_lowerCamelCase	)
										__lowercase			  =							torch.randn(_lowerCamelCase       ,generator=_lowerCamelCase       ,device=self.device       ,dtype=_lowerCamelCase	)
				else:
							if latents_reference.shape != latents_shape:
										raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"	)
							__lowercase			  =							latents_reference.to(self.device	)
							__lowercase			  =							latents.to(self.device	)

				# This is the key part of the pipeline where we
				# try to ensure that the generated images w/ the same seed
				# but different sizes actually result in similar images
				__lowercase			  =							(latents_shape[3] - latents_shape_reference[3]) // 2
				__lowercase			  =							(latents_shape[2] - latents_shape_reference[2]) // 2
				__lowercase			  =							latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
				__lowercase			  =							latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
				__lowercase			  =							0 if dx < 0 else dx
				__lowercase			  =							0 if dy < 0 else dy
				__lowercase			  =							max(-dx       ,0	)
				__lowercase			  =							max(-dy       ,0	)
				# import pdb
				# pdb.set_trace()
				__lowercase			  =							latents_reference[:, :, dy : dy + h, dx : dx + w]

				# set timesteps
				self.scheduler.set_timesteps(_lowerCamelCase	)

				# Some schedulers like PNDM have timesteps as arrays
				# It's more optimized to move all timesteps to correct device beforehand
				__lowercase			  =							self.scheduler.timesteps.to(self.device	)

				# scale the initial noise by the standard deviation required by the scheduler
				__lowercase			  =							latents * self.scheduler.init_noise_sigma

				# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
				# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
				# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
				# and should be between [0, 1]
				__lowercase			  =							'''eta''' in set(inspect.signature(self.scheduler.step	).parameters.keys()	)
				__lowercase			  =							{}
				if accepts_eta:
							__lowercase			  =							eta

				for i, t in enumerate(self.progress_bar(_lowerCamelCase	)	):
							# expand the latents if we are doing classifier free guidance
							__lowercase			  =							torch.cat([latents] * 2	) if do_classifier_free_guidance else latents
							__lowercase			  =							self.scheduler.scale_model_input(_lowerCamelCase       ,_lowerCamelCase	)

							# predict the noise residual
							__lowercase			  =							self.unet(_lowerCamelCase       ,_lowerCamelCase       ,encoder_hidden_states=_lowerCamelCase	).sample

							# perform guidance
							if do_classifier_free_guidance:
										__lowercase  ,						__lowercase			  =							noise_pred.chunk(2	)
										__lowercase			  =							noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

							# compute the previous noisy sample x_t -> x_t-1
							__lowercase			  =							self.scheduler.step(_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,**_lowerCamelCase	).prev_sample

							# call the callback, if provided
							if callback is not None and i % callback_steps == 0:
										callback(_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)

				__lowercase			  =							1 / 0.1_8_2_1_5 * latents
				__lowercase			  =							self.vae.decode(_lowerCamelCase	).sample

				__lowercase			  =							(image / 2 + 0.5).clamp(0       ,1	)

				# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
				__lowercase			  =							image.cpu().permute(0       ,2       ,3       ,1	).float().numpy()

				if self.safety_checker is not None:
							__lowercase			  =							self.feature_extractor(self.numpy_to_pil(_lowerCamelCase	)       ,return_tensors='''pt'''	).to(
							    self.device	)
							__lowercase  ,						__lowercase			  =							self.safety_checker(
							    images=_lowerCamelCase       ,clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype	)	)
				else:
							__lowercase			  =							None

				if output_type == "pil":
							__lowercase			  =							self.numpy_to_pil(_lowerCamelCase	)

				if not return_dict:
							return (image, has_nsfw_concept)

				return StableDiffusionPipelineOutput(images=_lowerCamelCase       ,nsfw_content_detected=_lowerCamelCase	)












'''simple docstring'''



import argparse
import gc
import json
import os
import re

import torch
from huggingface_hub import hf_hub_download

from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig
from transformers.modeling_utils import WEIGHTS_INDEX_NAME, shard_checkpoint


_SCREAMING_SNAKE_CASE   =      {
    '''169M''': 1_2,
    '''430M''': 2_4,
    '''1B5''': 2_4,
    '''3B''': 3_2,
    '''7B''': 3_2,
    '''14B''': 4_0,
}

_SCREAMING_SNAKE_CASE   =      {
    '''169M''': 7_6_8,
    '''430M''': 1_0_2_4,
    '''1B5''': 2_0_4_8,
    '''3B''': 2_5_6_0,
    '''7B''': 4_0_9_6,
    '''14B''': 5_1_2_0,
}




def 				_lowerCAmelCase  (			lowerCamelCase_     :					Dict			):
			__lowercase			  =							list(state_dict.keys()			)
			for name in state_dict_keys:
						__lowercase			  =							state_dict.pop(lowerCamelCase_			)
						# emb -> embedding
						if name.startswith('''emb.'''			):
									__lowercase			  =							name.replace('''emb.'''						,    '''embeddings.'''			)
						# ln_0 -> pre_ln (only present at block 0)
						if name.startswith('''blocks.0.ln0'''			):
									__lowercase			  =							name.replace('''blocks.0.ln0'''						,    '''blocks.0.pre_ln'''			)
						# att -> attention
						__lowercase			  =							re.sub(r'''blocks\.(\d+)\.att'''						,    r'''blocks.\1.attention'''						,    lowerCamelCase_			)
						# ffn -> feed_forward
						__lowercase			  =							re.sub(r'''blocks\.(\d+)\.ffn'''						,    r'''blocks.\1.feed_forward'''						,    lowerCamelCase_			)
						# time_mix_k -> time_mix_key and reshape
						if name.endswith('''.time_mix_k'''			):
									__lowercase			  =							name.replace('''.time_mix_k'''						,    '''.time_mix_key'''			)
						# time_mix_v -> time_mix_value and reshape
						if name.endswith('''.time_mix_v'''			):
									__lowercase			  =							name.replace('''.time_mix_v'''						,    '''.time_mix_value'''			)
						# time_mix_r -> time_mix_key and reshape
						if name.endswith('''.time_mix_r'''			):
									__lowercase			  =							name.replace('''.time_mix_r'''						,    '''.time_mix_receptance'''			)

						if name != "head.weight":
									__lowercase			  =							'''rwkv.''' + name

						__lowercase			  =							weight
			return state_dict




def 				_lowerCAmelCase  (			lowerCamelCase_     :					Optional[int]						,    lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					List[str]						,    lowerCamelCase_     :					List[Any]=None						,    lowerCamelCase_     :					List[Any]=None						,    lowerCamelCase_     :					Any=False						,    lowerCamelCase_     :					int=None			):
			# 1. If possible, build the tokenizer.
			if tokenizer_file is None:
						print('''No `--tokenizer_file` provided, we will use the default tokenizer.'''			)
						__lowercase			  =							5_0_2_7_7
						__lowercase			  =							AutoTokenizer.from_pretrained('''EleutherAI/gpt-neox-20b'''			)
			else:
						__lowercase			  =							PreTrainedTokenizerFast(tokenizer_file=lowerCamelCase_			)
						__lowercase			  =							len(lowerCamelCase_			)
			tokenizer.save_pretrained(lowerCamelCase_			)

			# 2. Build the config
			__lowercase			  =							list(NUM_HIDDEN_LAYERS_MAPPING.keys()			)
			if size is None:
						# Try to infer size from the checkpoint name
						for candidate in possible_sizes:
									if candidate in checkpoint_file:
												__lowercase			  =							candidate
												break
						if size is None:
									raise ValueError('''Could not infer the size, please provide it with the `--size` argument.'''			)
			if size not in possible_sizes:
						raise ValueError(f"`size` should be one of {possible_sizes}, got {size}."			)

			__lowercase			  =							RwkvConfig(
			    vocab_size=lowerCamelCase_						,    num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]						,    hidden_size=HIDEN_SIZE_MAPPING[size]						,    )
			config.save_pretrained(lowerCamelCase_			)

			# 3. Download model file then convert state_dict
			__lowercase			  =							hf_hub_download(lowerCamelCase_						,    lowerCamelCase_			)
			__lowercase			  =							torch.load(lowerCamelCase_						,    map_location='''cpu'''			)
			__lowercase			  =							convert_state_dict(lowerCamelCase_			)

			# 4. Split in shards and save
			__lowercase  ,						__lowercase			  =							shard_checkpoint(lowerCamelCase_			)
			for shard_file, shard in shards.items():
						torch.save(lowerCamelCase_						,    os.path.join(lowerCamelCase_						,    lowerCamelCase_			)			)

			if index is not None:
						__lowercase			  =							os.path.join(lowerCamelCase_						,    lowerCamelCase_			)
						# Save the index as well
						with open(lowerCamelCase_						,    '''w'''						,    encoding='''utf-8'''			) as f:
									__lowercase			  =							json.dumps(lowerCamelCase_						,    indent=2						,    sort_keys=lowerCamelCase_			) + '''\n'''
									f.write(lowerCamelCase_			)

						# 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict
						print(
						    '''Cleaning up shards. This may error with an OOM error, it this is the case don\'t worry you still have converted the model.'''			)
						__lowercase			  =							list(shards.keys()			)

						del state_dict
						del shards
						gc.collect()

						for shard_file in shard_files:
									__lowercase			  =							torch.load(os.path.join(lowerCamelCase_						,    lowerCamelCase_			)			)
									torch.save({k: v.cpu().clone() for k, v in state_dict.items()}						,    os.path.join(lowerCamelCase_						,    lowerCamelCase_			)			)

			del state_dict
			gc.collect()

			if push_to_hub:
						if model_name is None:
									raise ValueError('''Please provide a `model_name` to push the model to the Hub.'''			)
						__lowercase			  =							AutoModelForCausalLM.from_pretrained(lowerCamelCase_			)
						model.push_to_hub(lowerCamelCase_						,    max_shard_size='''2GB'''			)
						tokenizer.push_to_hub(lowerCamelCase_			)


if __name__ == "__main__":
	_SCREAMING_SNAKE_CASE   =      argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '''--repo_id''', default=None, type=str, required=True, help='''Repo ID from which to pull the checkpoint.'''
	)
	parser.add_argument(
	    '''--checkpoint_file''', default=None, type=str, required=True, help='''Name of the checkpoint file in the repo.'''
	)
	parser.add_argument(
	    '''--output_dir''', default=None, type=str, required=True, help='''Where to save the converted model.'''
	)
	parser.add_argument(
	    '''--tokenizer_file''',
	    default=None,
	    type=str,
	    help='''Path to the tokenizer file to use (if not provided, only the model is converted).''',
	)
	parser.add_argument(
	    '''--size''',
	    default=None,
	    type=str,
	    help='''Size of the model. Will be inferred from the `checkpoint_file` if not passed.''',
	)
	parser.add_argument(
	    '''--push_to_hub''',
	    action='''store_true''',
	    help='''Push to the Hub the converted model.''',
	)
	parser.add_argument(
	    '''--model_name''',
	    default=None,
	    type=str,
	    help='''Name of the pushed model on the Hub, including the username / organization.''',
	)

	_SCREAMING_SNAKE_CASE   =      parser.parse_args()
	convert_rmkv_checkpoint_to_hf_format(
	    args.repo_id,
	    args.checkpoint_file,
	    args.output_dir,
	    size=args.size,
	    tokenizer_file=args.tokenizer_file,
	    push_to_hub=args.push_to_hub,
	    model_name=args.model_name,
	)





