prettified all the code using "blue" at the urging of @tildebyte

scripts/dream.py

@@ -8,62 +8,66 @@ import sys
 import copy
 import warnings
 import ldm.dream.readline
-from   ldm.dream.pngwriter import PngWriter,PromptFormatter
+from ldm.dream.pngwriter import PngWriter, PromptFormatter
 
 debugging = False
 
+
 def main():
-    ''' Initialize command-line parsers and the diffusion model '''
+    """Initialize command-line parsers and the diffusion model"""
     arg_parser = create_argv_parser()
-    opt        = arg_parser.parse_args()
+    opt = arg_parser.parse_args()
     if opt.laion400m:
         # defaults suitable to the older latent diffusion weights
-        width   = 256
-        height  = 256
-        config  = "configs/latent-diffusion/txt2img-1p4B-eval.yaml"
-        weights = "models/ldm/text2img-large/model.ckpt"
+        width = 256
+        height = 256
+        config = 'configs/latent-diffusion/txt2img-1p4B-eval.yaml'
+        weights = 'models/ldm/text2img-large/model.ckpt'
     else:
         # some defaults suitable for stable diffusion weights
-        width   = 512
-        height  = 512
-        config  = "configs/stable-diffusion/v1-inference.yaml"
-        weights = "models/ldm/stable-diffusion-v1/model.ckpt"
+        width = 512
+        height = 512
+        config = 'configs/stable-diffusion/v1-inference.yaml'
+        weights = 'models/ldm/stable-diffusion-v1/model.ckpt'
 
-    print("* Initializing, be patient...\n")
+    print('* Initializing, be patient...\n')
     sys.path.append('.')
     from pytorch_lightning import logging
     from ldm.simplet2i import T2I
+
     # these two lines prevent a horrible warning message from appearing
     # when the frozen CLIP tokenizer is imported
     import transformers
+
     transformers.logging.set_verbosity_error()
-    
+
     # creating a simple text2image object with a handful of
     # defaults passed on the command line.
     # additional parameters will be added (or overriden) during
     # the user input loop
-    t2i = T2I(width=width,
-              height=height,
-              sampler_name=opt.sampler_name,
-              weights=weights,
-              full_precision=opt.full_precision,
-              config=config,
-              latent_diffusion_weights=opt.laion400m, # this is solely for recreating the prompt
-              embedding_path=opt.embedding_path,
-              device=opt.device
+    t2i = T2I(
+        width=width,
+        height=height,
+        sampler_name=opt.sampler_name,
+        weights=weights,
+        full_precision=opt.full_precision,
+        config=config,
+        latent_diffusion_weights=opt.laion400m,  # this is solely for recreating the prompt
+        embedding_path=opt.embedding_path,
+        device=opt.device,
     )
 
     # make sure the output directory exists
     if not os.path.exists(opt.outdir):
         os.makedirs(opt.outdir)
-        
+
     # gets rid of annoying messages about random seed
-    logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
+    logging.getLogger('pytorch_lightning').setLevel(logging.ERROR)
 
     infile = None
     try:
         if opt.infile is not None:
-            infile = open(opt.infile,'r')
+            infile = open(opt.infile, 'r')
     except FileNotFoundError as e:
         print(e)
         exit(-1)
@@ -73,135 +77,156 @@ def main():
 
     # load GFPGAN if requested
     if opt.use_gfpgan:
-        print("\n* --gfpgan was specified, loading gfpgan...")
+        print('\n* --gfpgan was specified, loading gfpgan...')
         with warnings.catch_warnings():
-            warnings.filterwarnings("ignore", category=DeprecationWarning) 
+            warnings.filterwarnings('ignore', category=DeprecationWarning)
 
             try:
-                model_path = os.path.join(opt.gfpgan_dir, opt.gfpgan_model_path)
+                model_path = os.path.join(
+                    opt.gfpgan_dir, opt.gfpgan_model_path
+                )
                 if not os.path.isfile(model_path):
-                    raise Exception("GFPGAN model not found at path "+model_path)
+                    raise Exception(
+                        'GFPGAN model not found at path ' + model_path
+                    )
 
                 sys.path.append(os.path.abspath(opt.gfpgan_dir))
                 from gfpgan import GFPGANer
 
-                bg_upsampler = load_gfpgan_bg_upsampler(opt.gfpgan_bg_upsampler, opt.gfpgan_bg_tile)
+                bg_upsampler = load_gfpgan_bg_upsampler(
+                    opt.gfpgan_bg_upsampler, opt.gfpgan_bg_tile
+                )
 
-                t2i.gfpgan = GFPGANer(model_path=model_path, upscale=opt.gfpgan_upscale, arch='clean', channel_multiplier=2, bg_upsampler=bg_upsampler)
+                t2i.gfpgan = GFPGANer(
+                    model_path=model_path,
+                    upscale=opt.gfpgan_upscale,
+                    arch='clean',
+                    channel_multiplier=2,
+                    bg_upsampler=bg_upsampler,
+                )
             except Exception:
                 import traceback
-                print("Error loading GFPGAN:", file=sys.stderr)
+
+                print('Error loading GFPGAN:', file=sys.stderr)
                 print(traceback.format_exc(), file=sys.stderr)
 
-    print("\n* Initialization done! Awaiting your command (-h for help, 'q' to quit, 'cd' to change output dir, 'pwd' to print output dir)...")
+    print(
+        "\n* Initialization done! Awaiting your command (-h for help, 'q' to quit, 'cd' to change output dir, 'pwd' to print output dir)..."
+    )
 
-    log_path   = os.path.join(opt.outdir,'dream_log.txt')
-    with open(log_path,'a') as log:
+    log_path = os.path.join(opt.outdir, 'dream_log.txt')
+    with open(log_path, 'a') as log:
         cmd_parser = create_cmd_parser()
-        main_loop(t2i,opt.outdir,cmd_parser,log,infile)
+        main_loop(t2i, opt.outdir, cmd_parser, log, infile)
         log.close()
     if infile:
         infile.close()
 
 
-def main_loop(t2i,outdir,parser,log,infile):
-    ''' prompt/read/execute loop '''
-    done       = False
+def main_loop(t2i, outdir, parser, log, infile):
+    """prompt/read/execute loop"""
+    done = False
     last_seeds = []
-    
+
     while not done:
         try:
-            command = infile.readline() if infile else input("dream> ") 
+            command = infile.readline() if infile else input('dream> ')
         except EOFError:
             done = True
             break
 
-        if infile and len(command)==0:
+        if infile and len(command) == 0:
             done = True
             break
 
-        if command.startswith(('#','//')):
+        if command.startswith(('#', '//')):
             continue
 
         # before splitting, escape single quotes so as not to mess
         # up the parser
-        command = command.replace("'","\\'")
+        command = command.replace("'", "\\'")
 
         try:
             elements = shlex.split(command)
         except ValueError as e:
             print(str(e))
             continue
-        
-        if len(elements)==0:
+
+        if len(elements) == 0:
             continue
 
-        if elements[0]=='q':
+        if elements[0] == 'q':
             done = True
             break
 
-        if elements[0]=='cd' and len(elements)>1:
+        if elements[0] == 'cd' and len(elements) > 1:
             if os.path.exists(elements[1]):
-                print(f"setting image output directory to {elements[1]}")
-                outdir=elements[1]
+                print(f'setting image output directory to {elements[1]}')
+                outdir = elements[1]
             else:
-                print(f"directory {elements[1]} does not exist")
+                print(f'directory {elements[1]} does not exist')
             continue
 
-        if elements[0]=='pwd':
-            print(f"current output directory is {outdir}")
+        if elements[0] == 'pwd':
+            print(f'current output directory is {outdir}')
             continue
-        
-        if elements[0].startswith('!dream'): # in case a stored prompt still contains the !dream command
+
+        if elements[0].startswith(
+            '!dream'
+        ):   # in case a stored prompt still contains the !dream command
             elements.pop(0)
-            
+
         # rearrange the arguments to mimic how it works in the Dream bot.
         switches = ['']
         switches_started = False
 
         for el in elements:
-            if el[0]=='-' and not switches_started:
+            if el[0] == '-' and not switches_started:
                 switches_started = True
             if switches_started:
                 switches.append(el)
             else:
                 switches[0] += el
                 switches[0] += ' '
-        switches[0] = switches[0][:len(switches[0])-1]
+        switches[0] = switches[0][: len(switches[0]) - 1]
 
         try:
-            opt      = parser.parse_args(switches)
+            opt = parser.parse_args(switches)
         except SystemExit:
             parser.print_help()
             continue
-        if len(opt.prompt)==0:
-            print("Try again with a prompt!")
+        if len(opt.prompt) == 0:
+            print('Try again with a prompt!')
             continue
-        if opt.seed is not None and opt.seed<0: # retrieve previous value!
+        if opt.seed is not None and opt.seed < 0:   # retrieve previous value!
             try:
                 opt.seed = last_seeds[opt.seed]
-                print(f"reusing previous seed {opt.seed}")
+                print(f'reusing previous seed {opt.seed}')
             except IndexError:
-                print(f"No previous seed at position {opt.seed} found")
+                print(f'No previous seed at position {opt.seed} found')
                 opt.seed = None
-                
-        normalized_prompt      = PromptFormatter(t2i,opt).normalize_prompt()
-        individual_images      = not opt.grid
+
+        normalized_prompt = PromptFormatter(t2i, opt).normalize_prompt()
+        individual_images = not opt.grid
 
         try:
-            file_writer        = PngWriter(outdir,normalized_prompt,opt.batch_size)
-            callback           = file_writer.write_image if individual_images else None
+            file_writer = PngWriter(outdir, normalized_prompt, opt.batch_size)
+            callback = file_writer.write_image if individual_images else None
 
-            image_list   = t2i.prompt2image(image_callback=callback,**vars(opt))
-            results      = file_writer.files_written     if individual_images else image_list
+            image_list = t2i.prompt2image(image_callback=callback, **vars(opt))
+            results = (
+                file_writer.files_written if individual_images else image_list
+            )
 
             if opt.grid and len(results) > 0:
                 grid_img = file_writer.make_grid([r[0] for r in results])
                 filename = file_writer.unique_filename(results[0][1])
-                seeds    = [a[1] for a in results]
-                results  = [[filename,seeds]]
-                metadata_prompt   = f'{normalized_prompt} -S{results[0][1]}'
-                file_writer.save_image_and_prompt_to_png(grid_img,metadata_prompt,filename)
+                seeds = [a[1] for a in results]
+                results = [[filename, seeds]]
+                metadata_prompt = f'{normalized_prompt} -S{results[0][1]}'
+                file_writer.save_image_and_prompt_to_png(
+                    grid_img, metadata_prompt, filename
+                )
 
             last_seeds = [r[1] for r in results]
 
@@ -213,10 +238,11 @@ def main_loop(t2i,outdir,parser,log,infile):
             print(e)
             continue
 
-        print("Outputs:")
-        write_log_message(t2i,normalized_prompt,results,log)
+        print('Outputs:')
+        write_log_message(t2i, normalized_prompt, results, log)
+
+    print('goodbye!')
 
-    print("goodbye!")
 
 def load_gfpgan_bg_upsampler(bg_upsampler, bg_tile=400):
     import torch
@@ -224,13 +250,24 @@ def load_gfpgan_bg_upsampler(bg_upsampler, bg_tile=400):
     if bg_upsampler == 'realesrgan':
         if not torch.cuda.is_available():  # CPU
             import warnings
-            warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. '
-                          'If you really want to use it, please modify the corresponding codes.')
+
+            warnings.warn(
+                'The unoptimized RealESRGAN is slow on CPU. We do not use it. '
+                'If you really want to use it, please modify the corresponding codes.'
+            )
             bg_upsampler = None
         else:
             from basicsr.archs.rrdbnet_arch import RRDBNet
             from realesrgan import RealESRGANer
-            model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+
+            model = RRDBNet(
+                num_in_ch=3,
+                num_out_ch=3,
+                num_feat=64,
+                num_block=23,
+                num_grow_ch=32,
+                scale=2,
+            )
             bg_upsampler = RealESRGANer(
                 scale=2,
                 model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
@@ -238,12 +275,14 @@ def load_gfpgan_bg_upsampler(bg_upsampler, bg_tile=400):
                 tile=bg_tile,
                 tile_pad=10,
                 pre_pad=0,
-                half=True)  # need to set False in CPU mode
+                half=True,
+            )  # need to set False in CPU mode
     else:
         bg_upsampler = None
 
     return bg_upsampler
 
+
 # variant generation is going to be superseded by a generalized
 # "prompt-morph" functionality
 # def generate_variants(t2i,outdir,opt,previous_gens):
@@ -268,110 +307,209 @@ def load_gfpgan_bg_upsampler(bg_upsampler, bg_tile=400):
 #                 continue
 #     print(f'{opt.variants} variants generated')
 #     return variants
-                
 
-def write_log_message(t2i,prompt,results,logfile):
-    ''' logs the name of the output image, its prompt and seed to the terminal, log file, and a Dream text chunk in the PNG metadata'''
-    last_seed  = None
-    img_num    = 1
-    seenit     = {}
+
+def write_log_message(t2i, prompt, results, logfile):
+    """logs the name of the output image, its prompt and seed to the terminal, log file, and a Dream text chunk in the PNG metadata"""
+    last_seed = None
+    img_num = 1
+    seenit = {}
 
     for r in results:
         seed = r[1]
-        log_message = (f'{r[0]}: {prompt} -S{seed}')
+        log_message = f'{r[0]}: {prompt} -S{seed}'
 
         print(log_message)
-        logfile.write(log_message+"\n")
+        logfile.write(log_message + '\n')
         logfile.flush()
 
+
 def create_argv_parser():
-    parser = argparse.ArgumentParser(description="Parse script's command line args")
-    parser.add_argument("--laion400m",
-                        "--latent_diffusion",
-                        "-l",
-                        dest='laion400m',
-                        action='store_true',
-                        help="fallback to the latent diffusion (laion400m) weights and config")
-    parser.add_argument("--from_file",
-                        dest='infile',
-                        type=str,
-                        help="if specified, load prompts from this file")
-    parser.add_argument('-n','--iterations',
-                        type=int,
-                        default=1,
-                        help="number of images to generate")
-    parser.add_argument('-F','--full_precision',
-                        dest='full_precision',
-                        action='store_true',
-                        help="use slower full precision math for calculations")
-    parser.add_argument('--sampler','-m',
-                        dest="sampler_name",
-                        choices=['ddim', 'k_dpm_2_a', 'k_dpm_2', 'k_euler_a', 'k_euler', 'k_heun', 'k_lms', 'plms'],
-                        default='k_lms',
-                        help="which sampler to use (k_lms) - can only be set on command line")
-    parser.add_argument('--outdir',
-                        '-o',
-                        type=str,
-                        default="outputs/img-samples",
-                        help="directory in which to place generated images and a log of prompts and seeds (outputs/img-samples")
-    parser.add_argument('--embedding_path',
-                        type=str,
-                        help="Path to a pre-trained embedding manager checkpoint - can only be set on command line")
-    parser.add_argument('--device',
-                        '-d',
-                        type=str,
-                        default="cuda",
-                        help="device to run stable diffusion on. defaults to cuda `torch.cuda.current_device()` if avalible")
+    parser = argparse.ArgumentParser(
+        description="Parse script's command line args"
+    )
+    parser.add_argument(
+        '--laion400m',
+        '--latent_diffusion',
+        '-l',
+        dest='laion400m',
+        action='store_true',
+        help='fallback to the latent diffusion (laion400m) weights and config',
+    )
+    parser.add_argument(
+        '--from_file',
+        dest='infile',
+        type=str,
+        help='if specified, load prompts from this file',
+    )
+    parser.add_argument(
+        '-n',
+        '--iterations',
+        type=int,
+        default=1,
+        help='number of images to generate',
+    )
+    parser.add_argument(
+        '-F',
+        '--full_precision',
+        dest='full_precision',
+        action='store_true',
+        help='use slower full precision math for calculations',
+    )
+    parser.add_argument(
+        '--sampler',
+        '-m',
+        dest='sampler_name',
+        choices=[
+            'ddim',
+            'k_dpm_2_a',
+            'k_dpm_2',
+            'k_euler_a',
+            'k_euler',
+            'k_heun',
+            'k_lms',
+            'plms',
+        ],
+        default='k_lms',
+        help='which sampler to use (k_lms) - can only be set on command line',
+    )
+    parser.add_argument(
+        '--outdir',
+        '-o',
+        type=str,
+        default='outputs/img-samples',
+        help='directory in which to place generated images and a log of prompts and seeds (outputs/img-samples',
+    )
+    parser.add_argument(
+        '--embedding_path',
+        type=str,
+        help='Path to a pre-trained embedding manager checkpoint - can only be set on command line',
+    )
+    parser.add_argument(
+        '--device',
+        '-d',
+        type=str,
+        default='cuda',
+        help='device to run stable diffusion on. defaults to cuda `torch.cuda.current_device()` if avalible',
+    )
     # GFPGAN related args
-    parser.add_argument('--gfpgan',
-                        dest='use_gfpgan',
-                        action='store_true',
-                        help="load gfpgan for use in the dreambot. Note: Enabling GFPGAN will require more GPU memory")
-    parser.add_argument("--gfpgan_upscale",
-                        type=int,
-                        default=2,
-                        help="The final upsampling scale of the image. Default: 2. Only used if --gfpgan is specified")
-    parser.add_argument("--gfpgan_bg_upsampler",
-                        type=str,
-                        default='realesrgan',
-                        help="Background upsampler. Default: None. Options: realesrgan, none. Only used if --gfpgan is specified")
-    parser.add_argument("--gfpgan_bg_tile",
-                        type=int,
-                        default=400,
-                        help="Tile size for background sampler, 0 for no tile during testing. Default: 400. Only used if --gfpgan is specified")
-    parser.add_argument("--gfpgan_model_path",
-                        type=str,
-                        default='experiments/pretrained_models/GFPGANv1.3.pth',
-                        help="indicates the path to the GFPGAN model, relative to --gfpgan_dir. Only used if --gfpgan is specified")
-    parser.add_argument("--gfpgan_dir",
-                        type=str,
-                        default='../GFPGAN',
-                        help="indicates the directory containing the GFPGAN code. Only used if --gfpgan is specified")
+    parser.add_argument(
+        '--gfpgan',
+        dest='use_gfpgan',
+        action='store_true',
+        help='load gfpgan for use in the dreambot. Note: Enabling GFPGAN will require more GPU memory',
+    )
+    parser.add_argument(
+        '--gfpgan_upscale',
+        type=int,
+        default=2,
+        help='The final upsampling scale of the image. Default: 2. Only used if --gfpgan is specified',
+    )
+    parser.add_argument(
+        '--gfpgan_bg_upsampler',
+        type=str,
+        default='realesrgan',
+        help='Background upsampler. Default: None. Options: realesrgan, none. Only used if --gfpgan is specified',
+    )
+    parser.add_argument(
+        '--gfpgan_bg_tile',
+        type=int,
+        default=400,
+        help='Tile size for background sampler, 0 for no tile during testing. Default: 400. Only used if --gfpgan is specified',
+    )
+    parser.add_argument(
+        '--gfpgan_model_path',
+        type=str,
+        default='experiments/pretrained_models/GFPGANv1.3.pth',
+        help='indicates the path to the GFPGAN model, relative to --gfpgan_dir. Only used if --gfpgan is specified',
+    )
+    parser.add_argument(
+        '--gfpgan_dir',
+        type=str,
+        default='../GFPGAN',
+        help='indicates the directory containing the GFPGAN code. Only used if --gfpgan is specified',
+    )
     return parser
-                        
-    
+
+
 def create_cmd_parser():
-    parser = argparse.ArgumentParser(description='Example: dream> a fantastic alien landscape -W1024 -H960 -s100 -n12')
+    parser = argparse.ArgumentParser(
+        description='Example: dream> a fantastic alien landscape -W1024 -H960 -s100 -n12'
+    )
     parser.add_argument('prompt')
-    parser.add_argument('-s','--steps',type=int,help="number of steps")
-    parser.add_argument('-S','--seed',type=int,help="image seed; a +ve integer, or use -1 for the previous seed, -2 for the one before that, etc")
-    parser.add_argument('-n','--iterations',type=int,default=1,help="number of samplings to perform (slower, but will provide seeds for individual images)")
-    parser.add_argument('-b','--batch_size',type=int,default=1,help="number of images to produce per sampling (will not provide seeds for individual images!)")
-    parser.add_argument('-W','--width',type=int,help="image width, multiple of 64")
-    parser.add_argument('-H','--height',type=int,help="image height, multiple of 64")
-    parser.add_argument('-C','--cfg_scale',default=7.5,type=float,help="prompt configuration scale")
-    parser.add_argument('-g','--grid',action='store_true',help="generate a grid")
-    parser.add_argument('-i','--individual',action='store_true',help="generate individual files (default)")
-    parser.add_argument('-I','--init_img',type=str,help="path to input image for img2img mode (supersedes width and height)")
-    parser.add_argument('-f','--strength',default=0.75,type=float,help="strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely")
-    parser.add_argument('-G','--gfpgan_strength', default=0.5, type=float, help="The strength at which to apply the GFPGAN model to the result, in order to improve faces.")
-# variants is going to be superseded by a generalized "prompt-morph" function
-#    parser.add_argument('-v','--variants',type=int,help="in img2img mode, the first generated image will get passed back to img2img to generate the requested number of variants")
-    parser.add_argument('-x','--skip_normalize',action='store_true',help="skip subprompt weight normalization")
+    parser.add_argument('-s', '--steps', type=int, help='number of steps')
+    parser.add_argument(
+        '-S',
+        '--seed',
+        type=int,
+        help='image seed; a +ve integer, or use -1 for the previous seed, -2 for the one before that, etc',
+    )
+    parser.add_argument(
+        '-n',
+        '--iterations',
+        type=int,
+        default=1,
+        help='number of samplings to perform (slower, but will provide seeds for individual images)',
+    )
+    parser.add_argument(
+        '-b',
+        '--batch_size',
+        type=int,
+        default=1,
+        help='number of images to produce per sampling (will not provide seeds for individual images!)',
+    )
+    parser.add_argument(
+        '-W', '--width', type=int, help='image width, multiple of 64'
+    )
+    parser.add_argument(
+        '-H', '--height', type=int, help='image height, multiple of 64'
+    )
+    parser.add_argument(
+        '-C',
+        '--cfg_scale',
+        default=7.5,
+        type=float,
+        help='prompt configuration scale',
+    )
+    parser.add_argument(
+        '-g', '--grid', action='store_true', help='generate a grid'
+    )
+    parser.add_argument(
+        '-i',
+        '--individual',
+        action='store_true',
+        help='generate individual files (default)',
+    )
+    parser.add_argument(
+        '-I',
+        '--init_img',
+        type=str,
+        help='path to input image for img2img mode (supersedes width and height)',
+    )
+    parser.add_argument(
+        '-f',
+        '--strength',
+        default=0.75,
+        type=float,
+        help='strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely',
+    )
+    parser.add_argument(
+        '-G',
+        '--gfpgan_strength',
+        default=0.5,
+        type=float,
+        help='The strength at which to apply the GFPGAN model to the result, in order to improve faces.',
+    )
+    # variants is going to be superseded by a generalized "prompt-morph" function
+    #    parser.add_argument('-v','--variants',type=int,help="in img2img mode, the first generated image will get passed back to img2img to generate the requested number of variants")
+    parser.add_argument(
+        '-x',
+        '--skip_normalize',
+        action='store_true',
+        help='skip subprompt weight normalization',
+    )
     return parser
 
 
-
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()
-

scripts/preload_models.py

@@ -11,26 +11,28 @@ import warnings
 transformers.logging.set_verbosity_error()
 
 # this will preload the Bert tokenizer fles
-print("preloading bert tokenizer...")
+print('preloading bert tokenizer...')
 from transformers import BertTokenizerFast
-tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
-print("...success")
+
+tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
+print('...success')
 
 # this will download requirements for Kornia
-print("preloading Kornia requirements (ignore the deprecation warnings)...")
+print('preloading Kornia requirements (ignore the deprecation warnings)...')
 with warnings.catch_warnings():
-    warnings.filterwarnings("ignore", category=DeprecationWarning) 
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
     import kornia
-print("...success")
+print('...success')
 
-version='openai/clip-vit-large-patch14'
+version = 'openai/clip-vit-large-patch14'
 
 print('preloading CLIP model (Ignore the deprecation warnings)...')
 sys.stdout.flush()
 import clip
 from transformers import CLIPTokenizer, CLIPTextModel
-tokenizer  =CLIPTokenizer.from_pretrained(version)
-transformer=CLIPTextModel.from_pretrained(version)
+
+tokenizer = CLIPTokenizer.from_pretrained(version)
+transformer = CLIPTextModel.from_pretrained(version)
 print('\n\n...success')
 
 # In the event that the user has installed GFPGAN and also elected to use
@@ -38,23 +40,33 @@ print('\n\n...success')
 gfpgan = False
 try:
     from realesrgan import RealESRGANer
+
     gfpgan = True
 except ModuleNotFoundError:
     pass
 
 if gfpgan:
-    print("Loading models from RealESRGAN and facexlib")
+    print('Loading models from RealESRGAN and facexlib')
     try:
         from basicsr.archs.rrdbnet_arch import RRDBNet
         from facexlib.utils.face_restoration_helper import FaceRestoreHelper
-        RealESRGANer(scale=2,
-                     model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
-                     model=RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2))
-        FaceRestoreHelper(1,det_model='retinaface_resnet50')
-        print("...success")
+
+        RealESRGANer(
+            scale=2,
+            model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
+            model=RRDBNet(
+                num_in_ch=3,
+                num_out_ch=3,
+                num_feat=64,
+                num_block=23,
+                num_grow_ch=32,
+                scale=2,
+            ),
+        )
+        FaceRestoreHelper(1, det_model='retinaface_resnet50')
+        print('...success')
     except Exception:
         import traceback
-        print("Error loading GFPGAN:")
+
+        print('Error loading GFPGAN:')
         print(traceback.format_exc())
-
-