

import logging
import os
import time
from argparse import Namespace
from pathlib import Path

import datasets
import torch
from accelerate import Accelerator, DistributedType
from accelerate.utils import ProjectConfiguration
from arguments import TrainingArguments
from datasets import load_dataset
from huggingface_hub import Repository
from torch.optim import AdamW
from torch.utils.data import IterableDataset
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, get_scheduler, set_seed



class 						lowercase_					(						lowerCAmelCase_  ):



						def __init__(    self   :       List[str]    ,      A__   :       Union[str, Any]    ,      A__   :       Optional[int]    ,      A__   :       Any=False    ,      A__   :       List[Any]=1024    ,      A__   :       Any=1024    ,      A__   :       Optional[Any]=3.6    ,      A__   :       str=False    ,      )					->  List[str]:
								_snake_case		=							tokenizer
								_snake_case		=							tokenizer.bos_token_id
								_snake_case		=							dataset
								_snake_case		=							seq_length
								_snake_case		=							0
								_snake_case		=							infinite
								_snake_case		=							0
								_snake_case		=							tokenized

								if self.tokenized:
										_snake_case		=							seq_length * num_of_sequences
										_snake_case		=							'''input_ids'''
								else:
										_snake_case		=							seq_length * chars_per_token * num_of_sequences
										_snake_case		=							'''content'''



						def __iter__(    self   :       Dict	)					->  Optional[Any]:
								_snake_case		=							iter(self.dataset	)
								_snake_case		=							True
								while more_examples:
										_snake_case,  _snake_case		=							[], 0
										while True:
												if buffer_len >= self.max_buffer_size:
														break
												try:
														buffer.append(next(lowerCamelCase_	)[self.content_field]	)
														buffer_len += len(buffer[-1]	)
												except StopIteration:
														if self.infinite:
																_snake_case		=							iter(self.dataset	)
																self.epoch += 1
																logger.info(f"""Dataset epoch: {self.epoch}"""	)
														else:
																_snake_case		=							False
																break
										if self.tokenized:
												_snake_case		=							buffer
										else:
												_snake_case		=							self.tokenizer(lowerCamelCase_    ,      truncation=lowerCamelCase_	)['''input_ids''']
										_snake_case		=							[]
										for tokenized_input in tokenized_inputs:
												all_token_ids.extend(tokenized_input + [self.concat_token_id]	)
										for i in range(0    ,      len(lowerCamelCase_	)    ,      self.seq_length	):
												_snake_case		=							all_token_ids[i : i + self.seq_length]
												if len(lowerCamelCase_	) == self.seq_length:
														self.current_size += 1
														yield torch.tensor(lowerCamelCase_	)







						def 		UpperCamelCase_	(    self   :       List[Any]    ,      A__   :       List[Any]=1000	)					->  List[Any]:
								return ShufflerIterDataPipe(self    ,      buffer_size=lowerCamelCase_	)



def    snake_case_(_UpperCamelCase  )							->			Dict:



		"""simple docstring"""






		_snake_case		=							args.model_ckpt.split('''/'''  )[-1]
		_snake_case		=							logging.getLogger(__name__  )
		_snake_case		=							Path(args.save_dir  ) / '''log/'''
		log_dir.mkdir(exist_ok=a_  )
		_snake_case		=							F"""debug_{accelerator.process_index}.log"""
		logging.basicConfig(
		    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''				, datefmt='''%m/%d/%Y %H:%M:%S'''				, level=logging.INFO				, handlers=[logging.FileHandler(log_dir / filename  ), logging.StreamHandler()]				, )
		if accelerator.is_main_process:  # we only want to setup logging once
				accelerator.init_trackers(a_				, vars(a_  )  )
				_snake_case		=							accelerator.trackers[0].run.name
				logger.setLevel(logging.INFO  )
				datasets.utils.logging.set_verbosity_info()
				transformers.utils.logging.set_verbosity_info()
		else:
				_snake_case		=							''''''
				logger.setLevel(logging.ERROR  )
				datasets.utils.logging.set_verbosity_error()
				transformers.utils.logging.set_verbosity_error()
		return logger, run_name



def    snake_case_(_UpperCamelCase  )							->			Dict:



		"""simple docstring"""






		_snake_case		=							{'''streaming''': True}
		_snake_case		=							load_dataset(args.dataset_name_train				, split='''train'''				, **a_  )
		_snake_case		=							train_data.shuffle(buffer_size=args.shuffle_buffer				, seed=args.seed  )
		_snake_case		=							load_dataset(args.dataset_name_valid				, split='''train'''				, **a_  )
		_snake_case		=							ConstantLengthDataset(
		    a_				, a_				, infinite=a_				, seq_length=args.seq_length				, tokenized=args.tokenized  )
		_snake_case		=							ConstantLengthDataset(
		    a_				, a_				, infinite=a_				, seq_length=args.seq_length				, tokenized=args.tokenized  )
		_snake_case		=							train_dataset.shuffle(buffer_size=args.shuffle_buffer  )
		_snake_case		=							DataLoader(a_				, batch_size=args.train_batch_size				, shuffle=a_  )
		_snake_case		=							DataLoader(a_				, batch_size=args.valid_batch_size  )
		return train_dataloader, eval_dataloader



def    snake_case_(_UpperCamelCase				, _UpperCamelCase				, _UpperCamelCase=["bias", "ln_1.weight", "ln_2.weight", "ln_f.weight"]  )							->			Dict:



		"""simple docstring"""






		_snake_case,  _snake_case		=							[], []
		for n, p in model.named_parameters():
				if any(nd in n for nd in no_decay  ):
						params_without_wd.append(a_  )
				else:
						params_with_wd.append(a_  )
		return [
		    {"params": params_with_wd, "weight_decay": args.weight_decay},
		    {"params": params_without_wd, "weight_decay": 0.0},
		]



def    snake_case_(_UpperCamelCase				, _UpperCamelCase  )							->			Optional[Any]:



		"""simple docstring"""






		logger.info(F"""Step {step}: {metrics}"""  )
		if accelerator.is_main_process:
				accelerator.log(a_				, a_  )



def    snake_case_(_UpperCamelCase				, _UpperCamelCase				, _UpperCamelCase  )							->			int:



		"""simple docstring"""






		_snake_case		=							accelerator.unwrap_model(a_  ).config
		_snake_case		=							4 if args.gradient_checkpointing else 3
		_snake_case		=							args.train_batch_size * accelerator.state.num_processes * args.gradient_accumulation_steps
		_snake_case		=							24 * checkpoint_factor * batch_size * args.seq_length * config_model.n_layer * (config_model.n_embd**2)
		_snake_case		=							factor * (
		    1.0
		    + (args.seq_length / (6.0 * config_model.n_embd))
		    + (tokenizer.vocab_size / (16.0 * config_model.n_layer * config_model.n_embd))
		)
		_snake_case		=							flops_per_iteration / (elapsed_time * accelerator.state.num_processes * (10**12))
		return tflops



def    snake_case_(_UpperCamelCase  )							->			Dict:



		"""simple docstring"""






		model.eval()
		_snake_case		=							[]
		for step, batch in enumerate(a_  ):
				with torch.no_grad():
						_snake_case		=							model(a_				, labels=a_  )
				_snake_case		=							outputs.loss.repeat(args.valid_batch_size  )
				losses.append(accelerator.gather(a_  )  )
				if args.max_eval_steps > 0 and step >= args.max_eval_steps:
						break
		_snake_case		=							torch.cat(a_  )
		_snake_case		=							losses[: eval_dataloader.dataset.current_size].mean()
		try:
				_snake_case		=							torch.exp(a_  )
		except OverflowError:
				_snake_case		=							float('''inf'''  )
		return loss.item(), perplexity.item()


# Settings
__A    	=			HfArgumentParser(TrainingArguments)
__A    	=			parser.parse_args()

# Accelerator
__A    	=			ProjectConfiguration(project_dir=args.save_dir, logging_dir='''log''')
__A    	=			Accelerator(log_with=['''wandb''', '''tensorboard'''], project_config=config)
__A    	=			{str(k): str(v) for k, v in accelerator.state.__dict__.items()}

__A    	=			Namespace(**vars(args), **acc_state)
__A    	=			accelerator.state.num_processes * args.train_batch_size
set_seed(args.seed)

# Clone model repository
if accelerator.is_main_process:
	__A    	=			Repository(args.save_dir, clone_from=args.model_ckpt)

# Logging
__A	,  __A    	=			setup_logging(args)
logger.info(accelerator.state)

# Checkout new branch on repo
if accelerator.is_main_process:
	hf_repo.git_checkout(run_name, create_branch_ok=True)

# Load model and tokenizer
__A    	=			AutoModelForCausalLM.from_pretrained(args.save_dir)
if args.gradient_checkpointing:
	model.gradient_checkpointing_enable()
__A    	=			AutoTokenizer.from_pretrained(args.save_dir)

# Load dataset and dataloader
__A	,  __A    	=			create_dataloaders(args)

# Prepare the optimizer and learning rate scheduler
__A    	=			AdamW(get_grouped_params(model, args), lr=args.learning_rate)
__A    	=			get_scheduler(
    name=args.lr_scheduler_type,
    optimizer=optimizer,
    num_warmup_steps=args.num_warmup_steps,
    num_training_steps=args.max_train_steps,
)
accelerator.register_for_checkpointing(lr_scheduler)



def    snake_case_()							->			int:



		"""simple docstring"""






		return optimizer.param_groups[0]["lr"]


# Prepare everything with our `accelerator`.
__A	,  __A	,  __A	,  __A    	=			accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)

# load in the weights and states from a previous save
if args.resume_from_checkpoint:
	if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
		accelerator.print(f'''Resumed from checkpoint: {args.resume_from_checkpoint}''')
		accelerator.load_state(args.resume_from_checkpoint)
		__A    	=			os.path.basename(args.resume_from_checkpoint)
	else:
		# Get the most recent checkpoint
		__A    	=			[f.name for f in os.scandir(args.save_dir) if f.is_dir() and '''step''' in str(f)]
		dirs.sort(key=os.path.getctime)
		__A    	=			dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
	# Extract the step of the checkpoint to continue from there
	__A    	=			os.path.splitext(path)[0]
	__A    	=			int(training_difference.replace('''step_''', ''''''))

# Train model
model.train()
__A    	=			0
__A    	=			time.time()
__A    	=			0
for step, batch in enumerate(train_dataloader, start=1):
	if args.resume_from_checkpoint and step < resume_step:
		continue  # we need to skip steps until we reach the resumed step
	__A    	=			model(batch, labels=batch, use_cache=False).loss
	__A    	=			accelerator.gather(loss.repeat(args.train_batch_size)).mean()
	loss_tracking += avg_loss.item() / args.gradient_accumulation_steps
	log_metrics(step, {'''samples''': step * samples_per_step, '''loss_per_step/train''': loss.item()})
	__A    	=			loss / args.gradient_accumulation_steps
	if step % args.gradient_accumulation_steps != 0:
		# Prevent backward from doing gradient all_reduce in every step
		if accelerator.distributed_type == DistributedType.MULTI_GPU:
			with model.no_sync():
				accelerator.backward(loss)
		else:
			accelerator.backward(loss)
	else:
		__A    	=			get_lr()
		accelerator.backward(loss)
		accelerator.clip_grad_norm_(model.parameters(), 1.0)
		optimizer.step()
		lr_scheduler.step()
		optimizer.zero_grad()
		__A    	=			time.time() - t_start
		__A    	=			compute_tflops(elapsed_time, accelerator, args)
		log_metrics(
		    step,
		    {
		        '''steps''': completed_steps,
		        '''loss/train''': loss_tracking,
		        '''lr''': lr,
		        '''tflops''': tflops,
		        '''time_per_iteration''': elapsed_time,
		    },
		)
		__A    	=			time.time()
		__A    	=			0
		completed_steps += 1
	if step % args.save_checkpoint_steps == 0:
		logger.info('''Evaluating and saving model checkpoint''')
		__A	,  __A    	=			evaluate(args)
		log_metrics(step, {'''loss/eval''': eval_loss, '''perplexity''': perplexity})
		accelerator.wait_for_everyone()
		__A    	=			os.path.join(args.save_dir, f'''step_{step}''')
		accelerator.save_state(save_dir)
		if accelerator.is_main_process:
			hf_repo.push_to_hub(commit_message=f'''step {step}''')
		model.train()
	if completed_steps >= args.max_train_steps:
		break

# Evaluate and save the last checkpoint
logger.info('''Evaluating and saving model after training''')
__A	,  __A    	=			evaluate(args)
log_metrics(step, {'''loss/eval''': eval_loss, '''perplexity''': perplexity})
accelerator.wait_for_everyone()
__A    	=			accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(args.save_dir, save_function=accelerator.save)
__A    	=			os.path.join(args.save_dir, f'''step_{step}''')
accelerator.save_state(save_dir)
if accelerator.is_main_process:
	hf_repo.push_to_hub(commit_message='''final model''')




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


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


__A    	=			'''
    Examples:
        ```py
        >>> import torch

        >>> from diffusers import ConsistencyModelPipeline

        >>> device = "cuda"
        >>> # Load the cd_imagenet64_l2 checkpoint.
        >>> model_id_or_path = "openai/diffusers-cd_imagenet64_l2"
        >>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
        >>> pipe.to(device)

        >>> # Onestep Sampling
        >>> image = pipe(num_inference_steps=1).images[0]
        >>> image.save("cd_imagenet64_l2_onestep_sample.png")

        >>> # Onestep sampling, class-conditional image generation
        >>> # ImageNet-64 class label 145 corresponds to king penguins
        >>> image = pipe(num_inference_steps=1, class_labels=145).images[0]
        >>> image.save("cd_imagenet64_l2_onestep_sample_penguin.png")

        >>> # Multistep sampling, class-conditional image generation
        >>> # Timesteps can be explicitly specified; the particular timesteps below are from the original Github repo:
        >>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77
        >>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]
        >>> image.save("cd_imagenet64_l2_multistep_sample_penguin.png")
        ```
'''



class 						lowercase_					(						__lowercase  ):



      def __init__(    self   :       str    ,      A__   :       UNetaDModel    ,      A__   :       CMStochasticIterativeScheduler	)					->  None:
        super().__init__()

        self.register_modules(
            unet=A__    ,      scheduler=A__    ,      )

        _snake_case		=							None



      def 		UpperCamelCase_	(    self   :       Union[str, Any]    ,      A__   :       Optional[int]=0	)					->  List[str]:
        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'''	)

        _snake_case		=							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   :       Optional[int]    ,      A__   :       Optional[Any]=0	)					->  Union[str, Any]:
        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.'''	)

        _snake_case		=							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)

        _snake_case		=							None
        for cpu_offloaded_model in [self.unet]:
          _snake_case,  _snake_case		=							cpu_offload_with_hook(A__    ,      A__    ,      prev_module_hook=A__	)

        if self.safety_checker is not None:
          _snake_case,  _snake_case		=							cpu_offload_with_hook(self.safety_checker    ,      A__    ,      prev_module_hook=A__	)

        # We'll offload the last model manually.
        _snake_case		=							hook



      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def 		UpperCamelCase_	(    self   :       int	)					->  List[str]:
        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   :       List[str]    ,      A__   :       Tuple    ,      A__   :       str    ,      A__   :       Tuple    ,      A__   :       Union[str, Any]    ,      A__   :       Union[str, Any]    ,      A__   :       Union[str, Any]    ,      A__   :       Tuple    ,      A__   :       Union[str, Any]=None	)					->  Optional[int]:
        _snake_case		=							(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:
          _snake_case		=							randn_tensor(A__    ,      generator=A__    ,      device=A__    ,      dtype=A__	)
        else:
          _snake_case		=							latents.to(device=A__    ,      dtype=A__	)

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



      def 		UpperCamelCase_	(    self   :       Optional[Any]    ,      A__   :       torch.FloatTensor    ,      A__   :       str = "pil"	)					->  Optional[int]:
        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
        _snake_case		=							(sample / 2 + 0.5).clamp(0    ,      1	)
        if output_type == "pt":
          return sample

        # Equivalent to diffusers.VaeImageProcessor.pt_to_numpy
        _snake_case		=							sample.cpu().permute(0    ,      2    ,      3    ,      1	).numpy()
        if output_type == "np":
          return sample

        # Output_type must be 'pil'
        _snake_case		=							self.numpy_to_pil(A__	)
        return sample



      def 		UpperCamelCase_	(    self   :       Any    ,      A__   :       Optional[Any]    ,      A__   :       Dict    ,      A__   :       int=None	)					->  Dict:
        if self.unet.config.num_class_embeds is not None:
          if isinstance(A__    ,      A__	):
            _snake_case		=							torch.tensor(A__    ,      dtype=torch.int	)
          elif isinstance(A__    ,      A__	):
            assert batch_size == 1, "Batch size must be 1 if classes is an int"
            _snake_case		=							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
            _snake_case		=							torch.randint(0    ,      self.unet.config.num_class_embeds    ,      size=(batch_size,)	)
          _snake_case		=							class_labels.to(A__	)
        else:
          _snake_case		=							None
        return class_labels



      def 		UpperCamelCase_	(    self   :       List[str]    ,      A__   :       List[str]    ,      A__   :       List[str]    ,      A__   :       str    ,      A__   :       List[Any]    ,      A__   :       Optional[Any]    ,      A__   :       Dict	)					->  Optional[Any]:
        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:
          _snake_case		=							(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   :       Any    ,      A__   :       int = 1    ,      A__   :       Optional[Union[torch.Tensor, List[int], int]] = None    ,      A__   :       int = 1    ,      A__   :       List[int] = None    ,      A__   :       Optional[Union[torch.Generator, List[torch.Generator]]] = None    ,      A__   :       Optional[torch.FloatTensor] = None    ,      A__   :       Optional[str] = "pil"    ,      A__   :       bool = True    ,      A__   :       Optional[Callable[[int, int, torch.FloatTensor], None]] = None    ,      A__   :       int = 1    ,      )					->  Dict:
        _snake_case		=							self.unet.config.sample_size
        _snake_case		=							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)
        _snake_case		=							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
        _snake_case		=							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__	)
          _snake_case		=							self.scheduler.timesteps
          _snake_case		=							len(A__	)
        else:
          self.scheduler.set_timesteps(A__	)
          _snake_case		=							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__	):
            _snake_case		=							self.scheduler.scale_model_input(A__    ,      A__	)
            _snake_case		=							self.unet(A__    ,      A__    ,      class_labels=A__    ,      return_dict=A__	)[0]

            _snake_case		=							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
        _snake_case		=							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__	)

