





"""simple docstring"""





import argparse
import dataclasses
import json
import logging
import os
import shutil
from typing import List, Optional

import datasets
from accelerate import Accelerator
from datasets import load_dataset
from finetuning import finetune
from tqdm.auto import tqdm

import transformers
from transformers import AutoConfig, set_seed
from transformers.trainer_utils import IntervalStrategy


_UpperCamelCase							=      logging.getLogger(__name__)

_UpperCamelCase							=      'pytorch_model.bin'
@dataclasses.dataclass
class 		SCREAMING_SNAKE_CASE_							:
				"""simple docstring"""
				__snake_case   :       str	 =		dataclasses.field(
				    metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models."""}   )
				__snake_case   :       Optional[str]	 =		dataclasses.field(
				    default=snake_case__       ,  metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co."""}       ,  )

@dataclasses.dataclass
class 		SCREAMING_SNAKE_CASE_							:
				"""simple docstring"""
				__snake_case   :       str	 =		dataclasses.field(metadata={"""help""": """A csv or a json file containing the training data."""}   )
				__snake_case   :       str	 =		dataclasses.field(metadata={"""help""": """A csv or a json file containing the data to predict on."""}   )
				__snake_case   :       Optional[str]	 =		dataclasses.field(
				    default=snake_case__       ,  metadata={"""help""": """A csv or a json file containing the validation data."""}   )
				__snake_case   :       Optional[str]	 =		dataclasses.field(
				    default=snake_case__       ,  metadata={"""help""": """The name of the task to train on."""}       ,  )
				__snake_case   :       Optional[List[str]]	 =		dataclasses.field(
				    default=snake_case__       ,  metadata={"""help""": """The list of labels for the task."""}   )

@dataclasses.dataclass
class 		SCREAMING_SNAKE_CASE_							:
				"""simple docstring"""
				__snake_case   :       str	 =		dataclasses.field(
				    metadata={"""help""": """The output directory where the model predictions and checkpoints will be written."""}   )
				__snake_case   :       Optional[str]	 =		dataclasses.field(
				    default="""accuracy"""       ,  metadata={"""help""": """The evaluation metric used for the task."""}   )
				__snake_case   :       Optional[str]	 =		dataclasses.field(
				    default="""no"""       ,  metadata={
				        """help""": """The evaluation strategy to adopt during training. Possible values are: [\"no\", \"step\", \"epoch]"""
				    }       ,  )
				__snake_case   :       Optional[int]	 =		dataclasses.field(
				    default=10       ,  metadata={"""help""": """Number of evaluation calls with no improvement after which training will be stopped."""}       ,  )
				__snake_case   :       Optional[float]	 =		dataclasses.field(
				    default=0.0       ,  metadata={
				        """help""": """How much the specified evaluation metric must improve to satisfy early stopping conditions."""
				    }       ,  )
				__snake_case   :       Optional[bool]	 =		dataclasses.field(
				    default=snake_case__       ,  metadata={"""help""": """Whether to filter the pseudo-labeled data based on the confidence score."""}       ,  )
				__snake_case   :       Optional[bool]	 =		dataclasses.field(
				    default=snake_case__       ,  metadata={"""help""": """Whether to filter the pseudo-labeled data based on the validation performance."""}       ,  )
				__snake_case   :       Optional[bool]	 =		dataclasses.field(
				    default=snake_case__       ,  metadata={"""help""": """Whether to fine-tune on labeled data after pseudo training."""}       ,  )
				__snake_case   :       Optional[float]	 =		dataclasses.field(
				    default=0.0       ,  metadata={"""help""": """Confidence threshold for pseudo-labeled data filtering."""}       ,  )
				__snake_case   :       Optional[int]	 =		dataclasses.field(
				    default=1_00       ,  metadata={"""help""": """Number of evaluation calls with no improvement after which training will be stopped."""}       ,  )
				__snake_case   :       Optional[int]	 =		dataclasses.field(
				    default=snake_case__       ,  metadata={"""help""": """Random seed for initialization."""}       ,  )



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Dict    ,						SCREAMING_SNAKE_CASE       :							Optional[int]    ,						SCREAMING_SNAKE_CASE       :							Any    ,						SCREAMING_SNAKE_CASE       :							Tuple    ,						SCREAMING_SNAKE_CASE       :							Optional[int]    ,						SCREAMING_SNAKE_CASE       :							Optional[int]   ):


							'''simple docstring'''







							__lowerCamelCase      :      List[str]         =datasets.concatenate_datasets([infer_input, infer_output]    ,						axis=1   )

							if args.do_filter_by_confidence:
														__lowerCamelCase      :      Optional[Any]         =dataset.filter(lambda SCREAMING_SNAKE_CASE   : example["probability"] > args.confidence_threshold   )

							if args.do_filter_by_val_performance:
														assert eval_result >= 0.0 and eval_result <= 1.0
														__lowerCamelCase      :      Tuple         =int(eval_result * len(SCREAMING_SNAKE_CASE   )   )
														print(SCREAMING_SNAKE_CASE   )
														__lowerCamelCase      :      int         =dataset.sort('''probability'''    ,						reverse=SCREAMING_SNAKE_CASE   )
														__lowerCamelCase      :      Any         =dataset.select(range(SCREAMING_SNAKE_CASE   )   )

							__lowerCamelCase      :      int         =dataset.remove_columns(['''label''', '''probability''']   )
							__lowerCamelCase      :      int         =dataset.rename_column('''prediction'''    ,						'''label'''   )
							__lowerCamelCase      :      int         =dataset.map(lambda SCREAMING_SNAKE_CASE   : {"label": idalabel[example["label"]]}   )
							__lowerCamelCase      :      Dict         =dataset.shuffle(seed=args.seed   )

							__lowerCamelCase      :      Tuple         =os.path.join(SCREAMING_SNAKE_CASE    ,						F'train_pseudo.{args.data_file_extension}'   )
							if args.data_file_extension == "csv":
														dataset.to_csv(SCREAMING_SNAKE_CASE    ,						index=SCREAMING_SNAKE_CASE   )
							else:
														dataset.to_json(SCREAMING_SNAKE_CASE   )



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Any    ,						SCREAMING_SNAKE_CASE       :							Optional[Any]    ,						SCREAMING_SNAKE_CASE       :							Optional[int]    ,						SCREAMING_SNAKE_CASE       :							Any    ,						**SCREAMING_SNAKE_CASE       :							Optional[int]   ):


							'''simple docstring'''







							__lowerCamelCase      :      Optional[Any]         =Accelerator()
							# Make one log on every process with the configuration for debugging.
							logging.basicConfig(
							    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''    ,						datefmt='''%m/%d/%Y %H:%M:%S'''    ,						level=logging.INFO    ,						)
							logger.info(accelerator.state   )

							# Setup logging, we only want one process per machine to log things on the
							# screen. accelerator.is_local_main_process is only True for one process per
							# machine.
							logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR   )

							if accelerator.is_local_main_process:
														datasets.utils.logging.set_verbosity_warning()
														transformers.utils.logging.set_verbosity_info()
							else:
														datasets.utils.logging.set_verbosity_error()
														transformers.utils.logging.set_verbosity_error()

							__lowerCamelCase      :      Union[str, Any]         =STModelArguments(model_name_or_path=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      int         =STDataArguments(train_file=SCREAMING_SNAKE_CASE    ,						infer_file=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      int         =STTrainingArguments(output_dir=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      int         =argparse.Namespace()

							for arg_class in (model_args, data_args, training_args):
														for key, value in vars(SCREAMING_SNAKE_CASE   ).items():
																					setattr(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )

							for key, value in kwargs.items():
														if hasattr(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   ):
																					setattr(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )

    # Sanity checks
							__lowerCamelCase      :      Optional[int]         ={}
							__lowerCamelCase      :      Any         =None

							# You need to provide the training data and the data to predict on
							assert args.train_file is not None
							assert args.infer_file is not None
							__lowerCamelCase      :      List[str]         =args.train_file
							__lowerCamelCase      :      int         =args.infer_file

							if args.evaluation_strategy != IntervalStrategy.NO.value:
														assert args.eval_file is not None
														__lowerCamelCase      :      Tuple         =args.eval_file

							for key in data_files:
														__lowerCamelCase      :      str         =data_files[key].split('''.'''   )[-1]
														assert extension in ["csv", "json"], F'`{key}_file` should be a csv or a json file.'
														if args.data_file_extension is None:
																					__lowerCamelCase      :      str         =extension
														else:
																					assert extension == args.data_file_extension, F'`{key}_file` should be a {args.data_file_extension} file`.'

							assert (
							    args.eval_metric in datasets.list_metrics()
							), F'{args.eval_metric} not in the list of supported metrics {datasets.list_metrics()}.'

							# If passed along, set the training seed now.
							if args.seed is not None:
														set_seed(args.seed   )

							logger.info('''Creating the initial data directory for self-training...'''   )
							__lowerCamelCase      :      Dict         =F'{args.output_dir}/self-train_iter-{{}}'.format
							__lowerCamelCase      :      str         =data_dir_format(0   )

							if accelerator.is_main_process:
														if args.output_dir is not None:
																					os.makedirs(args.output_dir    ,						exist_ok=SCREAMING_SNAKE_CASE   )
																					os.makedirs(SCREAMING_SNAKE_CASE    ,						exist_ok=SCREAMING_SNAKE_CASE   )
							accelerator.wait_for_everyone()

							__lowerCamelCase      :      Tuple         =None
							__lowerCamelCase      :      Union[str, Any]         =None
							__lowerCamelCase      :      Any         =0
							__lowerCamelCase      :      str         =False
							# Show the progress bar
							__lowerCamelCase      :      Union[str, Any]         =tqdm(range(args.max_selftrain_iterations   )    ,						disable=not accelerator.is_local_main_process   )

							# Self-train
							for iteration in range(0    ,						int(args.max_selftrain_iterations   )   ):
														__lowerCamelCase      :      Any         =data_dir_format(SCREAMING_SNAKE_CASE   )
														assert os.path.exists(SCREAMING_SNAKE_CASE   )

														# Stage 1: initial fine-tuning for iteration = 0 or pseudo-training for
														# iteration > 0
														__lowerCamelCase      :      Optional[int]         =os.path.join(SCREAMING_SNAKE_CASE    ,						'''stage-1'''   )
														__lowerCamelCase      :      Optional[Any]         ={
														    '''accelerator''': accelerator,
														    '''model_name_or_path''': args.model_name_or_path,
														    '''cache_dir''': args.cache_dir,
														    '''do_train''': True,
														    '''train_file''': data_files['''train'''] if iteration == 0 else data_files['''train_pseudo'''],
														    '''do_eval''': True if args.eval_file is not None else False,
														    '''eval_file''': data_files['''eval'''],
														    '''do_predict''': True,
														    '''infer_file''': data_files['''infer'''],
														    '''task_name''': args.task_name,
														    '''label_list''': args.label_list,
														    '''output_dir''': current_output_dir,
														    '''eval_metric''': args.eval_metric,
														    '''evaluation_strategy''': args.evaluation_strategy,
														    '''early_stopping_patience''': args.early_stopping_patience,
														    '''early_stopping_threshold''': args.early_stopping_threshold,
														    '''seed''': args.seed,
														}
														# Add additional training arguments
														for key, value in kwargs.items():
																					if key not in arguments_dict and not hasattr(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   ):
																												arguments_dict.update({key: value}   )

														__lowerCamelCase      :      List[Any]         =os.path.join(SCREAMING_SNAKE_CASE    ,						'''best-checkpoint'''    ,						SCREAMING_SNAKE_CASE   )
														if os.path.exists(SCREAMING_SNAKE_CASE   ):
																					logger.info(
																					    '''Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 1.'''    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						)
														else:
																					logger.info('''***** Running self-training: iteration: %d, stage: 1 *****'''    ,						SCREAMING_SNAKE_CASE   )
																					finetune(**SCREAMING_SNAKE_CASE   )
																					accelerator.wait_for_everyone()
																					assert os.path.exists(SCREAMING_SNAKE_CASE   )
																					logger.info('''Self-training job completed: iteration: %d, stage: 1.'''    ,						SCREAMING_SNAKE_CASE   )

														if iteration > 0 and args.finetune_on_labeled_data:
																					# Stage 2 (optional): fine-tuning on the original labeled data
																					__lowerCamelCase      :      Optional[Any]         =os.path.join(SCREAMING_SNAKE_CASE    ,						'''best-checkpoint'''   )
																					__lowerCamelCase      :      str         =os.path.join(SCREAMING_SNAKE_CASE    ,						'''stage-2'''   )
																					# Update arguments_dict
																					__lowerCamelCase      :      Union[str, Any]         =model_path
																					__lowerCamelCase      :      List[str]         =data_files['''train''']
																					__lowerCamelCase      :      Dict         =current_output_dir

																					__lowerCamelCase      :      Dict         =os.path.join(SCREAMING_SNAKE_CASE    ,						'''best-checkpoint'''    ,						SCREAMING_SNAKE_CASE   )
																					if os.path.exists(SCREAMING_SNAKE_CASE   ):
																												logger.info(
																												    '''Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 2.'''    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						)
																					else:
																												logger.info('''***** Running self-training: iteration: %d, stage: 2 *****'''    ,						SCREAMING_SNAKE_CASE   )
																												finetune(**SCREAMING_SNAKE_CASE   )
																												accelerator.wait_for_everyone()
																												assert os.path.exists(SCREAMING_SNAKE_CASE   )
																												logger.info('''Self-training job completed: iteration: %d, stage: 2.'''    ,						SCREAMING_SNAKE_CASE   )

														__lowerCamelCase      :      str         =iteration
														__lowerCamelCase      :      Union[str, Any]         =data_dir_format(iteration + 1   )

														__lowerCamelCase      :      Dict         =AutoConfig.from_pretrained(os.path.join(SCREAMING_SNAKE_CASE    ,						'''best-checkpoint'''   )   )
														__lowerCamelCase      :      Optional[Any]         =config.idalabel
														__lowerCamelCase      :      Any         =os.path.join(SCREAMING_SNAKE_CASE    ,						'''eval_results_best-checkpoint.json'''   )
														__lowerCamelCase      :      int         =os.path.join(SCREAMING_SNAKE_CASE    ,						'''test_results_best-checkpoint.json'''   )
														assert os.path.exists(SCREAMING_SNAKE_CASE   )

														with open(SCREAMING_SNAKE_CASE    ,						'''r'''   ) as f:
																					__lowerCamelCase      :      Tuple         =float(json.load(SCREAMING_SNAKE_CASE   )[args.eval_metric]   )
														__lowerCamelCase      :      str         =os.path.join(SCREAMING_SNAKE_CASE    ,						'''infer_output_best-checkpoint.csv'''   )
														assert os.path.exists(SCREAMING_SNAKE_CASE   )
														# Loading the dataset from local csv or json files.
														__lowerCamelCase      :      List[Any]         =load_dataset(args.data_file_extension    ,						data_files={'''data''': data_files['''infer''']}   )['''data''']
														__lowerCamelCase      :      Tuple         =load_dataset('''csv'''    ,						data_files={'''data''': infer_output_file}   )['''data''']

														if accelerator.is_main_process:
																					os.makedirs(SCREAMING_SNAKE_CASE    ,						exist_ok=SCREAMING_SNAKE_CASE   )
																					shutil.copy(SCREAMING_SNAKE_CASE    ,						os.path.join(SCREAMING_SNAKE_CASE    ,						F'eval_results_iter-{iteration}.json'   )   )
																					if os.path.exists(SCREAMING_SNAKE_CASE   ):
																												shutil.copy(SCREAMING_SNAKE_CASE    ,						os.path.join(SCREAMING_SNAKE_CASE    ,						F'test_results_iter-{iteration}.json'   )   )
																					create_pseudo_labeled_data(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )
														accelerator.wait_for_everyone()

														__lowerCamelCase      :      Optional[Any]         =os.path.join(SCREAMING_SNAKE_CASE    ,						F'train_pseudo.{args.data_file_extension}'   )

														if args.evaluation_strategy != IntervalStrategy.NO.value:
																					__lowerCamelCase      :      List[Any]         =eval_result

																					if best_iteration is None:
																												__lowerCamelCase      :      List[str]         =new_iteration
																												__lowerCamelCase      :      Tuple         =new_eval_result
																					else:
																												if new_eval_result - best_eval_result > args.early_stopping_threshold:
																																			__lowerCamelCase      :      Any         =new_iteration
																																			__lowerCamelCase      :      int         =new_eval_result
																																			__lowerCamelCase      :      Tuple         =0
																												else:
																																			if new_eval_result == best_eval_result:
																																										__lowerCamelCase      :      Tuple         =new_iteration
																																										__lowerCamelCase      :      str         =new_eval_result
																																			early_stopping_patience_counter += 1

																												if early_stopping_patience_counter >= args.early_stopping_patience:
																																			__lowerCamelCase      :      str         =True

														progress_bar.update(1   )

														if should_training_stop:
																					break

							if best_iteration is not None:
														# Save the best iteration
														logger.info('''Best iteration: %d'''    ,						SCREAMING_SNAKE_CASE   )
														logger.info('''Best evaluation result: %s = %f'''    ,						args.eval_metric    ,						SCREAMING_SNAKE_CASE   )
														accelerator.wait_for_everyone()
														if accelerator.is_main_process:
																					shutil.copy(
																					    os.path.join(SCREAMING_SNAKE_CASE    ,						F'eval_results_iter-{iteration}.json'   )    ,						os.path.join(SCREAMING_SNAKE_CASE    ,						'''eval_results_best-iteration.json'''   )    ,						)
							else:
														# Assume that the last iteration is the best
														logger.info('''Best iteration: %d'''    ,						args.max_selftrain_iterations - 1   )
														logger.info('''Best evaluation result: %s = %f'''    ,						args.eval_metric    ,						SCREAMING_SNAKE_CASE   )
														accelerator.wait_for_everyone()
														if accelerator.is_main_process:
																					shutil.copy(
																					    os.path.join(SCREAMING_SNAKE_CASE    ,						F'eval_results_iter-{args.max_selftrain_iterations - 1}.json'   )    ,						os.path.join(SCREAMING_SNAKE_CASE    ,						'''eval_results_best-iteration.json'''   )    ,						)











"""simple docstring"""





from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


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

_UpperCamelCase							=      '\n    Examples:\n        ```py\n        >>> import torch\n        >>> import numpy as np\n\n        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline\n        >>> from transformers import pipeline\n        >>> from diffusers.utils import load_image\n\n\n        >>> def make_hint(image, depth_estimator):\n        ...     image = depth_estimator(image)["depth"]\n        ...     image = np.array(image)\n        ...     image = image[:, :, None]\n        ...     image = np.concatenate([image, image, image], axis=2)\n        ...     detected_map = torch.from_numpy(image).float() / 255.0\n        ...     hint = detected_map.permute(2, 0, 1)\n        ...     return hint\n\n\n        >>> depth_estimator = pipeline("depth-estimation")\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior = pipe_prior.to("cuda")\n\n        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16\n        ... )\n        >>> pipe = pipe.to("cuda")\n\n\n        >>> img = load_image(\n        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"\n        ...     "/kandinsky/cat.png"\n        ... ).resize((768, 768))\n\n        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")\n\n        >>> prompt = "A robot, 4k photo"\n        >>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"\n\n        >>> generator = torch.Generator(device="cuda").manual_seed(43)\n\n        >>> image_emb, zero_image_emb = pipe_prior(\n        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator\n        ... ).to_tuple()\n\n        >>> images = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     hint=hint,\n        ...     num_inference_steps=50,\n        ...     generator=generator,\n        ...     height=768,\n        ...     width=768,\n        ... ).images\n\n        >>> images[0].save("robot_cat.png")\n        ```\n'



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Any    ,						SCREAMING_SNAKE_CASE       :							Tuple    ,						SCREAMING_SNAKE_CASE       :							Tuple=8   ):


							'''simple docstring'''







							__lowerCamelCase      :      Tuple         =height // scale_factor**2
							if height % scale_factor**2 != 0:
														new_height += 1
							__lowerCamelCase      :      Tuple         =width // scale_factor**2
							if width % scale_factor**2 != 0:
														new_width += 1
							return new_height * scale_factor, new_width * scale_factor
class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :Union[str, Any]						,    __lowercase  :UNetaDConditionModel						,    __lowercase  :DDPMScheduler						,    __lowercase  :VQModel						,    ):
											super().__init__()

											self.register_modules(
											    unet=__lowercase						,    scheduler=__lowercase						,    movq=__lowercase						,    )
											__lowerCamelCase      :      Any         =2 ** (len(self.movq.config.block_out_channels			) - 1)




				def 					__lowercase  (     self  :int						,    __lowercase  :Tuple						,    __lowercase  :Dict						,    __lowercase  :str						,    __lowercase  :Dict						,    __lowercase  :Union[str, Any]						,    __lowercase  :Dict			):
											if latents is None:
																		__lowerCamelCase      :      str         =randn_tensor(__lowercase						,    generator=__lowercase						,    device=__lowercase						,    dtype=__lowercase			)
											else:
																		if latents.shape != shape:
																									raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {shape}'			)
																		__lowerCamelCase      :      str         =latents.to(__lowercase			)

											__lowerCamelCase      :      Any         =latents * scheduler.init_noise_sigma
											return latents




				def 					__lowercase  (     self  :Dict						,    __lowercase  :Any=0			):
											if is_accelerate_available():
																		from accelerate import cpu_offload
											else:
																		raise ImportError('''Please install accelerate via `pip install accelerate`'''			)

											__lowerCamelCase      :      Any         =torch.device(f'cuda:{gpu_id}'			)

											__lowerCamelCase      :      Dict         =[
											    self.unet,
											    self.movq,
											]
											for cpu_offloaded_model in models:
																		if cpu_offloaded_model is not None:
																									cpu_offload(__lowercase						,    __lowercase			)




				def 					__lowercase  (     self  :Optional[int]						,    __lowercase  :List[Any]=0			):
											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.'''			)

											__lowerCamelCase      :      str         =torch.device(f'cuda:{gpu_id}'			)

											if self.device.type != "cpu":
																		self.to('''cpu'''						,    silence_dtype_warnings=__lowercase			)
																		torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

											__lowerCamelCase      :      Optional[Any]         =None
											for cpu_offloaded_model in [self.unet, self.movq]:
																		__lowerCamelCase    ,						__lowerCamelCase      :      Union[str, Any]         =cpu_offload_with_hook(__lowercase						,    __lowercase						,    prev_module_hook=__lowercase			)

											# We'll offload the last model manually.
											__lowerCamelCase      :      Dict         =hook




				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def 					__lowercase  (     self  :Tuple			):
											if not hasattr(self.unet						,    '''_hf_hook'''			):
																		return self.device
											for module in self.unet.modules():
																		if (
																		    hasattr(__lowercase						,    '''_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




				@torch.no_grad()
				@replace_example_docstring(__lowercase			)
				def __call__(     self  :Union[str, Any]						,    __lowercase  :Union[torch.FloatTensor, List[torch.FloatTensor]]						,    __lowercase  :Union[torch.FloatTensor, List[torch.FloatTensor]]						,    __lowercase  :torch.FloatTensor						,    __lowercase  :int = 512						,    __lowercase  :int = 512						,    __lowercase  :int = 100						,    __lowercase  :float = 4.0						,    __lowercase  :int = 1						,    __lowercase  :Optional[Union[torch.Generator, List[torch.Generator]]] = None						,    __lowercase  :Optional[torch.FloatTensor] = None						,    __lowercase  :Optional[str] = "pil"						,    __lowercase  :bool = True						,    ):
											__lowerCamelCase      :      Any         =self._execution_device

											__lowerCamelCase      :      Dict         =guidance_scale > 1.0

											if isinstance(__lowercase						,    __lowercase			):
																		__lowerCamelCase      :      Any         =torch.cat(__lowercase						,    dim=0			)
											if isinstance(__lowercase						,    __lowercase			):
																		__lowerCamelCase      :      Dict         =torch.cat(__lowercase						,    dim=0			)
											if isinstance(__lowercase						,    __lowercase			):
																		__lowerCamelCase      :      Dict         =torch.cat(__lowercase						,    dim=0			)

											__lowerCamelCase      :      List[Any]         =image_embeds.shape[0] * num_images_per_prompt

											if do_classifier_free_guidance:
																		__lowerCamelCase      :      Optional[int]         =image_embeds.repeat_interleave(__lowercase						,    dim=0			)
																		__lowerCamelCase      :      Any         =negative_image_embeds.repeat_interleave(__lowercase						,    dim=0			)
																		__lowerCamelCase      :      List[Any]         =hint.repeat_interleave(__lowercase						,    dim=0			)

											__lowerCamelCase      :      List[str]         =torch.cat([negative_image_embeds, image_embeds]						,    dim=0			).to(dtype=self.unet.dtype						,    device=__lowercase			)
											__lowerCamelCase      :      List[str]         =torch.cat([hint, hint]						,    dim=0			).to(dtype=self.unet.dtype						,    device=__lowercase			)
											self.scheduler.set_timesteps(__lowercase						,    device=__lowercase			)
											__lowerCamelCase      :      str         =self.scheduler.timesteps

											__lowerCamelCase      :      Tuple         =self.movq.config.latent_channels

											__lowerCamelCase    ,						__lowerCamelCase      :      List[str]         =downscale_height_and_width(__lowercase						,    __lowercase						,    self.movq_scale_factor			)

											# create initial latent
											__lowerCamelCase      :      Optional[int]         =self.prepare_latents(
											    (batch_size, num_channels_latents, height, width)						,    image_embeds.dtype						,    __lowercase						,    __lowercase						,    __lowercase						,    self.scheduler						,    )

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

																		__lowerCamelCase      :      List[Any]         ={'''image_embeds''': image_embeds, '''hint''': hint}
																		__lowerCamelCase      :      Union[str, Any]         =self.unet(
																		    sample=__lowercase						,    timestep=__lowercase						,    encoder_hidden_states=__lowercase						,    added_cond_kwargs=__lowercase						,    return_dict=__lowercase						,    )[0]

																		if do_classifier_free_guidance:
																									__lowerCamelCase    ,						__lowerCamelCase      :      Dict         =noise_pred.split(latents.shape[1]						,    dim=1			)
																									__lowerCamelCase    ,						__lowerCamelCase      :      int         =noise_pred.chunk(2			)
																									__lowerCamelCase    ,						__lowerCamelCase      :      Tuple         =variance_pred.chunk(2			)
																									__lowerCamelCase      :      Tuple         =noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																									__lowerCamelCase      :      Any         =torch.cat([noise_pred, variance_pred_text]						,    dim=1			)

																		if not (
																		    hasattr(self.scheduler.config						,    '''variance_type'''			)
																		    and self.scheduler.config.variance_type in ["learned", "learned_range"]
																		):
																									__lowerCamelCase    ,						__lowerCamelCase      :      Any         =noise_pred.split(latents.shape[1]						,    dim=1			)

																		# compute the previous noisy sample x_t -> x_t-1
																		__lowerCamelCase      :      Tuple         =self.scheduler.step(
																		    __lowercase						,    __lowercase						,    __lowercase						,    generator=__lowercase						,    )[0]
											# post-processing
											__lowerCamelCase      :      List[str]         =self.movq.decode(__lowercase						,    force_not_quantize=__lowercase			)['''sample''']

											if output_type not in ["pt", "np", "pil"]:
																		raise ValueError(f'Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}'			)

											if output_type in ["np", "pil"]:
																		__lowerCamelCase      :      List[str]         =image * 0.5 + 0.5
																		__lowerCamelCase      :      Optional[int]         =image.clamp(0						,    1			)
																		__lowerCamelCase      :      Union[str, Any]         =image.cpu().permute(0						,    2						,    3						,    1			).float().numpy()

											if output_type == "pil":
																		__lowerCamelCase      :      Optional[Any]         =self.numpy_to_pil(__lowercase			)

											if not return_dict:
																		return (image,)

											return ImagePipelineOutput(images=__lowercase			)




