import argparse
import random

import joblib
import numpy as np
import torch
from igf.igf import (
    SecondaryLearner,
    collect_objective_set,
    compute_perplexity,
    generate_datasets,
    load_gpta,
    recopy_gpta,
    set_seed,
    train_secondary_learner,
)
from torch.utils.data import DataLoader, RandomSampler

from transformers import GPTaLMHeadModel
def       __magic_name__(  __UpperCAmelCase=32 ,						__UpperCAmelCase=10 ,						__UpperCAmelCase=100 ,						__UpperCAmelCase=1026 ,						__UpperCAmelCase=True ,						__UpperCAmelCase="data/tokenized_stories_train_wikitext103.jbl" ,						__UpperCAmelCase="igf_context_pairs.jbl" ,						)						->    Dict:



      '''simple docstring'''

      set_seed(3		)
      # generate train_data and objective_set
      _lowerCamelCase    ,					_lowerCamelCase			    =				generate_datasets(
          __UpperCAmelCase ,						__UpperCAmelCase ,						number=__UpperCAmelCase ,						min_len=1026 ,						trim=__UpperCAmelCase		)
      # keeps model same across runs
      set_seed(4		)
      # model, lm_optimizer, lm_scheduler = recopy_gpt2(model, device, max_steps) # store original model weights
      # can we train on GPU?
      _lowerCamelCase			    =				torch.device('''cuda:0''' if torch.cuda.is_available() else '''cpu'''		)

      # load pretrained model
      _lowerCamelCase			    =				load_gpta('''gpt2'''		).to(__UpperCAmelCase		)
      print('''computing perplexity on objective set'''		)
      _lowerCamelCase			    =				compute_perplexity(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		).item()
      print('''perplexity on objective set:''' ,						__UpperCAmelCase		)

      # collect igf pairs and save to file demo.jbl
      collect_objective_set(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)

      # clean up, delete model and data we don't need anymore
      del model, train_data, objective_set
      torch.cuda.empty_cache()
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase=15 ,						__UpperCAmelCase=128 ,						__UpperCAmelCase=100 ,						__UpperCAmelCase="igf_model.pt" ,						)						->    Optional[Any]:



      '''simple docstring'''

      set_seed(42		)

      # Load pre-trained model
      _lowerCamelCase			    =				GPTaLMHeadModel.from_pretrained('''gpt2'''		)

      # Initialize secondary learner to use embedding weights of model
      _lowerCamelCase			    =				SecondaryLearner(__UpperCAmelCase		)

      # Train secondary learner
      _lowerCamelCase			    =				train_secondary_learner(
          __UpperCAmelCase ,						__UpperCAmelCase ,						max_epochs=__UpperCAmelCase ,						batch_size=__UpperCAmelCase ,						eval_freq=100 ,						igf_model_path=__UpperCAmelCase ,						)

      del model, secondary_learner_train_data
      torch.cuda.empty_cache()

      return secondary_learner
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase=32 ,						__UpperCAmelCase=1000 ,						__UpperCAmelCase=16 ,						__UpperCAmelCase=1.0 ,						__UpperCAmelCase=recopy_gpta ,						__UpperCAmelCase=None ,						__UpperCAmelCase=10 ,						__UpperCAmelCase="gpt2_finetuned.pt" ,						)						->    str:



      '''simple docstring'''

      _lowerCamelCase			    =				torch.device('''cuda:0''' if torch.cuda.is_available() else '''cpu'''		)
      _lowerCamelCase			    =				RandomSampler(__UpperCAmelCase		)
      _lowerCamelCase			    =				DataLoader(__UpperCAmelCase ,						sampler=__UpperCAmelCase		)

      _lowerCamelCase			    =				max_steps // (len(__UpperCAmelCase		)) + 1
      _lowerCamelCase			    =				0
      _lowerCamelCase			    =				torch.zeros((1, context_len) ,						dtype=torch.long ,						device=__UpperCAmelCase		)
      _lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				recopy_model(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)

      model.train()
      if secondary_learner is not None:
            secondary_learner.to(__UpperCAmelCase		)
            secondary_learner.eval()
      _lowerCamelCase			    =				[]
      _lowerCamelCase			    =				0

      _lowerCamelCase			    =				[]
      _lowerCamelCase			    =				[]

      # Compute the performance of the transformer model at the beginning
      _lowerCamelCase			    =				compute_perplexity(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)
      test_perps.append(__UpperCAmelCase		)
      print('''Test perplexity, step''' ,						__UpperCAmelCase ,						''':''' ,						__UpperCAmelCase		)
      for epoch in range(int(__UpperCAmelCase		)		):
            for step, example in enumerate(__UpperCAmelCase		):
                  torch.cuda.empty_cache()
                  _lowerCamelCase			    =				random.randint(0 ,						example.size(2		) - context_len - 1		)
                  _lowerCamelCase			    =				example[0, 0, start : start + context_len]
                  lm_optimizer.zero_grad()
                  _lowerCamelCase			    =				model(__UpperCAmelCase ,						labels=__UpperCAmelCase		)
                  _lowerCamelCase			    =				True

                  if secondary_learner is not None:
                        _lowerCamelCase			    =				secondary_learner.forward(
                            torch.tensor(__UpperCAmelCase ,						dtype=torch.long ,						device=__UpperCAmelCase		).unsqueeze(0		)		)[0].item()
                        observed_qs.append(float(__UpperCAmelCase		)		)

                        # Here we implement the simple non-constant threshold for the predicted IG(X) value
                        # We will decay the selectivity of our secondary learner filter from
                        # 1 standard deviation above average to 1 below average after 10 batches.

                        if global_step == 10:
                              _lowerCamelCase			    =				-1
                        if predicted_q < threshold:
                              _lowerCamelCase			    =				False

            # If we passed the filter, add the context to the batch!
                  if do_backprop:
                        contexts.append(np.array(context.cpu()		)		)
                        _lowerCamelCase			    =				outputs[0]
                        lm_loss.backward()
                        examples += 1

                  del outputs

                  # Once the batch is filled with enough contexts, backprop on the batch.
                  if examples == batch_size:
                        torch.cuda.empty_cache()
                        _lowerCamelCase			    =				0
                        # Do LM backprop
                        torch.nn.utils.clip_grad_norm_(model.parameters() ,						3.0		)
                        lm_optimizer.step()
                        lm_scheduler.step()  # Update learning rate schedule
                        global_step += 1
                        # Compute the performance of the transformer model at this batch
                        if global_step % eval_interval == 0:
                              _lowerCamelCase			    =				compute_perplexity(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)
                              test_perps.append(__UpperCAmelCase		)

                              print('''Test perplexity, step''' ,						__UpperCAmelCase ,						''':''' ,						__UpperCAmelCase		)
            # Break out of the loop after 60 batches
                  if max_steps > 0 and global_step > 60:
                        break
            if max_steps > 0 and global_step > 60:
                  break

    # save finetuned transformer model
      torch.save(model.state_dict() ,						__UpperCAmelCase		)
      torch.cuda.empty_cache()
      # Do some cleaning up so we can reinitialize for the next run of this function
      del lm_optimizer
      del lm_scheduler
      return model
def       __magic_name__(  )						->    Tuple:



      '''simple docstring'''

      _lowerCamelCase			    =				argparse.ArgumentParser(description='''Fine-tune a transformer model with IGF on a language modeling task'''		)

      # Required parameters
      parser.add_argument(
          '''--data_dir''' ,						default=__UpperCAmelCase ,						type=__UpperCAmelCase ,						required=__UpperCAmelCase ,						help='''The input data dir. Should contain data files for WikiText.''' ,						)
      parser.add_argument(
          '''--model_name_or_path''' ,						default=__UpperCAmelCase ,						type=__UpperCAmelCase ,						required=__UpperCAmelCase ,						help='''Path to pretrained model or model identifier from huggingface.co/models''' ,						)
      parser.add_argument(
          '''--data_file''' ,						type=__UpperCAmelCase ,						default=__UpperCAmelCase ,						help=(
              '''A jbl file containing tokenized data which can be split as objective dataset, '''
              '''train_dataset and test_dataset.'''
          ) ,						)

      parser.add_argument(
          '''--igf_data_file''' ,						type=__UpperCAmelCase ,						default=__UpperCAmelCase ,						help='''A jbl file containing the context and information gain pairs to train secondary learner.''' ,						)

      parser.add_argument(
          '''--output_dir''' ,						default=__UpperCAmelCase ,						type=__UpperCAmelCase ,						required=__UpperCAmelCase ,						help='''The output directory where the final fine-tuned model is stored.''' ,						)

      parser.add_argument(
          '''--tokenizer_name''' ,						default=__UpperCAmelCase ,						type=__UpperCAmelCase ,						help='''Pretrained tokenizer name or path if not the same as model_name''' ,						)
      parser.add_argument('''--seed''' ,						type=__UpperCAmelCase ,						default=__UpperCAmelCase ,						help='''A seed for reproducible training.'''		)

      parser.add_argument(
          '''--context_len''' ,						default=32 ,						type=__UpperCAmelCase ,						help=(
              '''The maximum total input sequence length after tokenization. Sequences longer '''
              '''than this will be truncated, sequences shorter will be padded.'''
          ) ,						)

      parser.add_argument(
          '''--size_objective_set''' ,						default=100 ,						type=__UpperCAmelCase ,						help='''number of articles that are long enough to be used as our objective set''' ,						)
      parser.add_argument(
          '''--eval_freq''' ,						default=100 ,						type=__UpperCAmelCase ,						help='''secondary model evaluation is triggered at eval_freq'''		)

      parser.add_argument('''--max_steps''' ,						default=1000 ,						type=__UpperCAmelCase ,						help='''To calculate training epochs'''		)

      parser.add_argument(
          '''--secondary_learner_batch_size''' ,						default=128 ,						type=__UpperCAmelCase ,						help='''batch size of training data for secondary learner''' ,						)

      parser.add_argument(
          '''--batch_size''' ,						default=16 ,						type=__UpperCAmelCase ,						help='''batch size of training data of language model(gpt2) '''		)

      parser.add_argument(
          '''--eval_interval''' ,						default=10 ,						type=__UpperCAmelCase ,						help=(
              '''decay the selectivity of our secondary learner filter from'''
              '''1 standard deviation above average to 1 below average after 10 batches'''
          ) ,						)

      parser.add_argument(
          '''--number''' ,						default=100 ,						type=__UpperCAmelCase ,						help='''The number of examples split to be used as objective_set/test_data'''		)

      parser.add_argument(
          '''--min_len''' ,						default=1026 ,						type=__UpperCAmelCase ,						help='''The minimum length of the article to be used as objective set'''		)

      parser.add_argument(
          '''--secondary_learner_max_epochs''' ,						default=15 ,						type=__UpperCAmelCase ,						help='''number of epochs to train secondary learner'''		)

      parser.add_argument('''--trim''' ,						default=__UpperCAmelCase ,						type=__UpperCAmelCase ,						help='''truncate the example if it exceeds context length'''		)

      parser.add_argument(
          '''--threshold''' ,						default=1.0 ,						type=__UpperCAmelCase ,						help=(
              '''The threshold value used by secondary learner to filter the train_data and allow only'''
              ''' informative data as input to the model'''
          ) ,						)

      parser.add_argument('''--finetuned_model_name''' ,						default='''gpt2_finetuned.pt''' ,						type=__UpperCAmelCase ,						help='''finetuned_model_name'''		)

      parser.add_argument(
          '''--recopy_model''' ,						default=__UpperCAmelCase ,						type=__UpperCAmelCase ,						help='''Reset the model to the original pretrained GPT-2 weights after each iteration''' ,						)

      # function calls
      # Collecting *n* pairs of context and information gain(X, IG(X)) for training the secondary learner
      generate_n_pairs(
          context_len=32 ,						max_steps=10 ,						size_objective_set=100 ,						min_len=1026 ,						trim=__UpperCAmelCase ,						data_file='''data/tokenized_stories_train_wikitext103.jbl''' ,						igf_data_file='''igf_context_pairs.jbl''' ,						)

      # Load train data for secondary learner
      _lowerCamelCase			    =				joblib.load('''data/IGF_values.jbl'''		)

      # Train secondary learner
      _lowerCamelCase			    =				training_secondary_learner(
          __UpperCAmelCase ,						secondary_learner_max_epochs=15 ,						secondary_learner_batch_size=128 ,						eval_freq=100 ,						igf_model_path='''igf_model.pt''' ,						)

      # load pretrained gpt2 model
      _lowerCamelCase			    =				GPTaLMHeadModel.from_pretrained('''gpt2'''		)
      set_seed(42		)

      # Generate train and test data to train and evaluate gpt2 model
      _lowerCamelCase    ,					_lowerCamelCase			    =				generate_datasets(
          context_len=32 ,						file='''data/tokenized_stories_train_wikitext103.jbl''' ,						number=100 ,						min_len=1026 ,						trim=__UpperCAmelCase		)

      # fine-tuning of the gpt2 model using igf (Information Gain Filtration)
      finetune(
          __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						context_len=32 ,						max_steps=1000 ,						batch_size=16 ,						threshold=1.0 ,						recopy_model=__UpperCAmelCase ,						secondary_learner=__UpperCAmelCase ,						eval_interval=10 ,						finetuned_model_name='''gpt2_finetuned.pt''' ,						)


if __name__ == "__main__":
     main()
import os
from glob import glob

import imageio
import torch
import torchvision
import wandb
from img_processing import custom_to_pil, loop_post_process, preprocess, preprocess_vqgan
from loaders import load_vqgan
from PIL import Image
from torch import nn

from transformers import CLIPModel, CLIPTokenizerFast
from utils import get_device, get_timestamp, show_pil



class  UpperCamelCase				:

     '''simple docstring'''

     def __init__(					self				,				A_ = "cpu"				,				A_ = "openai/clip-vit-large-patch14"						) ->				None:






           """simple docstring"""
           _lowerCamelCase			    =				device
           _lowerCamelCase			    =				CLIPTokenizerFast.from_pretrained(A_						)
           _lowerCamelCase			    =				[0.48145466, 0.4578275, 0.40821073]
           _lowerCamelCase			    =				[0.26862954, 0.26130258, 0.27577711]
           _lowerCamelCase			    =				torchvision.transforms.Normalize(self.image_mean				,				self.image_std						)
           _lowerCamelCase			    =				torchvision.transforms.Resize(2_24						)
           _lowerCamelCase			    =				torchvision.transforms.CenterCrop(2_24						)



     def 		UpperCamelCase_						(					self				,				A_						) ->				int:






           """simple docstring"""
           _lowerCamelCase			    =				self.resize(A_						)
           _lowerCamelCase			    =				self.center_crop(A_						)
           _lowerCamelCase			    =				self.normalize(A_						)
           return images





     def __call__(					self				,				A_=None				,				A_=None				,				**A_						) ->				Optional[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				self.tokenizer(text=A_				,				**A_						)
           _lowerCamelCase			    =				self.preprocess_img(A_						)
           _lowerCamelCase			    =				{key: value.to(self.device						) for (key, value) in encoding.items()}
           return encoding





class  UpperCamelCase				( nn.Module ):

     '''simple docstring'''

     def __init__(					self				,				A_=10				,				A_=0.01				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=False				,				A_=True				,				A_="image"				,				A_=True				,				A_=False				,				A_=False				,				A_=False				,				) ->				None:






           """simple docstring"""
           super().__init__()
           _lowerCamelCase			    =				None
           _lowerCamelCase			    =				device if device else get_device()
           if vqgan:
                 _lowerCamelCase			    =				vqgan
           else:
                 _lowerCamelCase			    =				load_vqgan(self.device				,				conf_path=A_				,				ckpt_path=A_						)
           self.vqgan.eval()
           if clip:
                 _lowerCamelCase			    =				clip
           else:
                 _lowerCamelCase			    =				CLIPModel.from_pretrained('''openai/clip-vit-base-patch32'''						)
           self.clip.to(self.device						)
           _lowerCamelCase			    =				ProcessorGradientFlow(device=self.device						)

           _lowerCamelCase			    =				iterations
           _lowerCamelCase			    =				lr
           _lowerCamelCase			    =				log
           _lowerCamelCase			    =				make_grid
           _lowerCamelCase			    =				return_val
           _lowerCamelCase			    =				quantize
           _lowerCamelCase			    =				self.vqgan.decoder.z_shape



     def 		UpperCamelCase_						(					self				,				A_=None				,				A_=None				,				A_=5				,				A_=True						) ->				Any:






           """simple docstring"""
           _lowerCamelCase			    =				[]
           if output_path is None:
                 _lowerCamelCase			    =				'''./animation.gif'''
           if input_path is None:
                 _lowerCamelCase			    =				self.save_path
           _lowerCamelCase			    =				sorted(glob(input_path + '''/*'''						)						)
           if not len(A_						):
                 raise ValueError(
                     '''No images found in save path, aborting (did you pass save_intermediate=True to the generate'''
                     ''' function?)'''						)
           if len(A_						) == 1:
                 print('''Only one image found in save path, (did you pass save_intermediate=True to the generate function?)'''						)
           _lowerCamelCase			    =				total_duration / len(A_						)
           _lowerCamelCase			    =				[frame_duration] * len(A_						)
           if extend_frames:
                 _lowerCamelCase			    =				1.5
                 _lowerCamelCase			    =				3
           for file_name in paths:
                 if file_name.endswith('''.png'''						):
                       images.append(imageio.imread(A_						)						)
           imageio.mimsave(A_				,				A_				,				duration=A_						)
           print(F'gif saved to {output_path}'						)



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






           """simple docstring"""
           if not (path or img):
                 raise ValueError('''Input either path or tensor'''						)
           if img is not None:
                 raise NotImplementedError
           _lowerCamelCase			    =				preprocess(Image.open(A_						)				,				target_image_size=2_56						).to(self.device						)
           _lowerCamelCase			    =				preprocess_vqgan(A_						)
           _lowerCamelCase    ,					*_lowerCamelCase			    =				self.vqgan.encode(A_						)
           return z



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






           """simple docstring"""
           _lowerCamelCase			    =				self.latent.detach().requires_grad_()
           _lowerCamelCase			    =				base_latent + transform_vector
           if self.quantize:
                 _lowerCamelCase    ,					*_lowerCamelCase			    =				self.vqgan.quantize(A_						)
           else:
                 _lowerCamelCase			    =				trans_latent
           return self.vqgan.decode(A_						)



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_=None						) ->				Any:






           """simple docstring"""
           _lowerCamelCase			    =				self.clip_preprocessor(text=A_				,				images=A_				,				return_tensors='''pt'''				,				padding=A_						)
           _lowerCamelCase			    =				self.clip(**A_						)
           _lowerCamelCase			    =				clip_outputs.logits_per_image
           if weights is not None:
                 _lowerCamelCase			    =				similarity_logits * weights
           return similarity_logits.sum()



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				Dict:






           """simple docstring"""
           _lowerCamelCase			    =				self._get_clip_similarity(pos_prompts['''prompts''']				,				A_				,				weights=(1 / pos_prompts['''weights'''])						)
           if neg_prompts:
                 _lowerCamelCase			    =				self._get_clip_similarity(neg_prompts['''prompts''']				,				A_				,				weights=neg_prompts['''weights''']						)
           else:
                 _lowerCamelCase			    =				torch.tensor([1]				,				device=self.device						)
           _lowerCamelCase			    =				-torch.log(A_						) + torch.log(A_						)
           return loss



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				str:






           """simple docstring"""
           _lowerCamelCase			    =				torch.randn_like(self.latent				,				requires_grad=A_				,				device=self.device						)
           _lowerCamelCase			    =				torch.optim.Adam([vector]				,				lr=self.lr						)

           for i in range(self.iterations						):
                 optim.zero_grad()
                 _lowerCamelCase			    =				self._add_vector(A_						)
                 _lowerCamelCase			    =				loop_post_process(A_						)
                 _lowerCamelCase			    =				self._get_CLIP_loss(A_				,				A_				,				A_						)
                 print('''CLIP loss'''				,				A_						)
                 if self.log:
                       wandb.log({'''CLIP Loss''': clip_loss}						)
                 clip_loss.backward(retain_graph=A_						)
                 optim.step()
                 if self.return_val == "image":
                       yield custom_to_pil(transformed_img[0]						)
                 else:
                       yield vector



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				Any:






           """simple docstring"""
           wandb.init(reinit=A_				,				project='''face-editor'''						)
           wandb.config.update({'''Positive Prompts''': positive_prompts}						)
           wandb.config.update({'''Negative Prompts''': negative_prompts}						)
           wandb.config.update({'''lr''': self.lr, '''iterations''': self.iterations}						)
           if image_path:
                 _lowerCamelCase			    =				Image.open(A_						)
                 _lowerCamelCase			    =				image.resize((2_56, 2_56)						)
                 wandb.log('''Original Image'''				,				wandb.Image(A_						)						)



     def 		UpperCamelCase_						(					self				,				A_						) ->				int:






           """simple docstring"""
           if not prompts:
                 return []
           _lowerCamelCase			    =				[]
           _lowerCamelCase			    =				[]
           if isinstance(A_				,				A_						):
                 _lowerCamelCase			    =				[prompt.strip() for prompt in prompts.split('''|'''						)]
           for prompt in prompts:
                 if isinstance(A_				,				(tuple, list)						):
                       _lowerCamelCase			    =				prompt[0]
                       _lowerCamelCase			    =				float(prompt[1]						)
                 elif ":" in prompt:
                       _lowerCamelCase    ,					_lowerCamelCase			    =				prompt.split(''':'''						)
                       _lowerCamelCase			    =				float(A_						)
                 else:
                       _lowerCamelCase			    =				prompt
                       _lowerCamelCase			    =				1.0
                 processed_prompts.append(A_						)
                 weights.append(A_						)
           return {
               "prompts": processed_prompts,
               "weights": torch.tensor(A_				,				device=self.device						),
           }





     def 		UpperCamelCase_						(					self				,				A_				,				A_=None				,				A_=None				,				A_=True				,				A_=False				,				A_=True				,				A_=True				,				A_=None				,				) ->				str:






           """simple docstring"""
           if image_path:
                 _lowerCamelCase			    =				self._get_latent(A_						)
           else:
                 _lowerCamelCase			    =				torch.randn(self.latent_dim				,				device=self.device						)
           if self.log:
                 self._init_logging(A_				,				A_				,				A_						)

           assert pos_prompts, "You must provide at least one positive prompt."
           _lowerCamelCase			    =				self.process_prompts(A_						)
           _lowerCamelCase			    =				self.process_prompts(A_						)
           if save_final and save_path is None:
                 _lowerCamelCase			    =				os.path.join('''./outputs/'''				,				'''_'''.join(pos_prompts['''prompts''']						)						)
           if not os.path.exists(A_						):
                 os.makedirs(A_						)
           else:
                 _lowerCamelCase			    =				save_path + '''_''' + get_timestamp()
                 os.makedirs(A_						)
           _lowerCamelCase			    =				save_path

           _lowerCamelCase			    =				self.vqgan.decode(self.latent						)[0]
           if show_intermediate:
                 print('''Original Image'''						)
                 show_pil(custom_to_pil(A_						)						)

           _lowerCamelCase			    =				loop_post_process(A_						)
           for iter, transformed_img in enumerate(self._optimize_CLIP(A_				,				A_				,				A_						)						):
                 if show_intermediate:
                       show_pil(A_						)
                 if save_intermediate:
                       transformed_img.save(os.path.join(self.save_path				,				F'iter_{iter:03d}.png'						)						)
                 if self.log:
                       wandb.log({'''Image''': wandb.Image(A_						)}						)
           if show_final:
                 show_pil(A_						)
           if save_final:
                 transformed_img.save(os.path.join(self.save_path				,				F'iter_{iter:03d}_final.png'						)						)