From 84459b72bd7e0881cd39ac2090369e0584dbd55d Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Sat, 13 Aug 2022 20:14:30 -0700 Subject: [PATCH 01/11] Updating the environment.yaml file so we can use a newer version of pytorch, torchvision and pytorch-lightning. --- environment.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/environment.yml b/environment.yml index 618063f..08cc1b5 100644 --- a/environment.yml +++ b/environment.yml @@ -4,11 +4,11 @@ channels: - conda-forge - defaults dependencies: - - pytorch::pytorch=1.10.0 - - pytorch::torchvision=0.11.1 - - cudatoolkit=10.2 + - pytorch::pytorch=1.12.0 + - pytorch::torchvision=0.13.1 + - cudatoolkit=10.2 # The cudatoolkit library could also be updated to 11.3 but might give some troubles with older GPUs, for an RTX 3050 or higher cudatoolkit=11.3 is recommended. - omegaconf - - pytorch-lightning + - pytorch-lightning=1.5.8 # For compatibility - tqdm - regex - kornia From e8067a9c58695189942302b70024011bef16a8c4 Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Sat, 13 Aug 2022 21:06:50 -0700 Subject: [PATCH 02/11] Add an option to customize the value of cutn through the the sidebar on the web interface. --- app.py | 32 ++++++++++++++++++++++++++++++-- defaults.yaml | 6 ++++-- logic.py | 6 +++++- 3 files changed, 39 insertions(+), 5 deletions(-) diff --git a/app.py b/app.py index 38558df..d2b7630 100644 --- a/app.py +++ b/app.py @@ -54,6 +54,7 @@ def generate_image( image_prompts: List[Image.Image] = [], continue_prev_run: bool = False, seed: Optional[int] = None, + cutn: int = 32, mse_weight: float = 0, mse_weight_decay: float = 0, mse_weight_decay_steps: int = 0, @@ -76,6 +77,7 @@ def generate_image( image_x=image_x, image_y=image_y, seed=seed, + cutn=cutn, init_image=init_image, image_prompts=image_prompts, continue_prev_run=continue_prev_run, @@ -136,7 +138,7 @@ def generate_image( frames = [] try: - # Try block catches st.script_runner.StopExecution, no need of a dedicated stop button + # Try block catches st.StopExecution, no need of a dedicated stop button # Reason is st.form is meant to be self-contained either within sidebar, or in main body # The way the form is implemented in this app splits the form across both regions # This is intended to prevent the model settings from crowding the main body @@ -208,6 +210,7 @@ def generate_image( "continue_prev_run": continue_prev_run, "prev_run_id": prev_run_id, "seed": run.seed, + "cutn": cutn, "Xdim": image_x, "ydim": image_y, "vqgan_ckpt": vqgan_ckpt, @@ -246,7 +249,7 @@ def generate_image( status_text.text("Done!") # End of run - except st.script_runner.StopException as e: + except st.StopException as e: # Dump output to dashboard print(f"Received Streamlit StopException") status_text.text("Execution interruped, dumping outputs ...") @@ -289,6 +292,7 @@ def generate_image( "continue_prev_run": continue_prev_run, "prev_run_id": prev_run_id, "seed": run.seed, + "cutn": cutn, "Xdim": image_x, "ydim": image_y, "vqgan_ckpt": vqgan_ckpt, @@ -430,6 +434,29 @@ def generate_image( else: seed = None + #cutn = st.sidebar.checkbox( + # "Set Cutn", + # value=defaults["cutn"], + # help="Check to set the number of cuts to pass to CLIP, lower values uses less VRAM, higher values increases the image quality. Will add option to specify the number of cuts", + #) + + use_cutn = st.sidebar.checkbox( + "Use Cutn", + value=defaults["use_cutn"], + help="Check to set the number of cuts to pass to CLIP, lower values uses less VRAM, higher values increase the image quality. Will add option to specify the number of cuts", + ) + cutn = st.sidebar.empty() + if use_cutn is True: + cutn = cutn.number_input( + "Number of Cuts sent to CLIP", + value=defaults["cutn"], + min_value=1, + step=1, + help="Specify the number of cuts to pass to CLIP, lower values uses less VRAM but higher values increases the image quality", + ) + else: + cutn = 32 + use_custom_starting_image = st.sidebar.checkbox( "Use starting image", value=defaults["use_starting_image"], @@ -645,6 +672,7 @@ def generate_image( image_x=int(image_x), image_y=int(image_y), seed=int(seed) if set_seed is True else None, + cutn=int(cutn), init_image=init_image, image_prompts=image_prompts, continue_prev_run=continue_prev_run, diff --git a/defaults.yaml b/defaults.yaml index 6633c53..9ccdb11 100644 --- a/defaults.yaml +++ b/defaults.yaml @@ -1,9 +1,11 @@ # Modify for different systems, e.g. larger default xdim/ydim for more powerful GPUs num_steps: 500 -Xdim: 640 -ydim: 480 +Xdim: 511 +ydim: 511 set_seed: false seed: 0 +use_cutn: true +cutn: 32 use_starting_image: false use_image_prompts: false continue_prev_run: false diff --git a/logic.py b/logic.py index 1f774b8..460d96c 100644 --- a/logic.py +++ b/logic.py @@ -70,6 +70,7 @@ def __init__( image_prompts: List[Image.Image] = [], continue_prev_run: bool = False, seed: Optional[int] = None, + cutn: int = 32, mse_weight=0.5, mse_weight_decay=0.1, mse_weight_decay_steps=50, @@ -98,6 +99,7 @@ def __init__( self.image_prompts = image_prompts self.continue_prev_run = continue_prev_run self.seed = seed + self.cutn = cutn # Setup ------------------------------------------------------------------------------ # Split text by "|" symbol @@ -121,7 +123,8 @@ def __init__( vqgan_config=f"assets/{vqgan_ckpt}.yaml", vqgan_checkpoint=f"assets/{vqgan_ckpt}.ckpt", step_size=0.05, - cutn=64, + #cutn=32, + cutn=cutn, cut_pow=1.0, display_freq=50, seed=seed, @@ -214,6 +217,7 @@ def model_init(self, init_image: Image.Image = None) -> None: None, :, None, None ] + if self.seed is not None: torch.manual_seed(self.seed) else: From b985a53048745a29a9857c6d41895b6d5671a22f Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Tue, 23 Aug 2022 10:51:06 -0700 Subject: [PATCH 03/11] - Added option to the sidebar to configure the clip model, cut power and step size/learning rate. - Added a small try block for the prev_image.png to reduce an issue that happens when the image is open on an external program and the model tries to load or save the file which resulted in a error, this should happen less frequently. - Added some extra information to the Streamlit progress bar, as I was not able to replace it by the stqdm library I tried to implement some of the features of said library like the number of iterations per seconds. - Added some lines to save the image of each step/frame that was created, these are the same images used to create the video when we hit the stop button, this should be useful to see step by step what the app has done as sometimes the video might be too fast or too low quality to see correctly what it has done. - Fixed an issue where the app could not run using the CPU only as it was by default trying to use at least one GPU. - Changed the format of the Zoom Factor field on the sidebar from "%.1e" to "%.4f", this should give better control of small increments and a better idea of what value it currently has. --- app.py | 211 +++++++++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 169 insertions(+), 42 deletions(-) diff --git a/app.py b/app.py index d2b7630..ede29e7 100644 --- a/app.py +++ b/app.py @@ -11,10 +11,12 @@ import datetime import shutil import torch -import json -import os +import json, time +import os, io, timeit import base64 import traceback +import clip +from stqdm import stqdm import argparse @@ -43,10 +45,12 @@ except ModuleNotFoundError: pass +#os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:30" def generate_image( text_input: str = "the first day of the waters", vqgan_ckpt: str = "vqgan_imagenet_f16_16384", + clip_model: str = "ViT-B/32", num_steps: int = 300, image_x: int = 300, image_y: int = 300, @@ -55,6 +59,8 @@ def generate_image( continue_prev_run: bool = False, seed: Optional[int] = None, cutn: int = 32, + cut_pow: float = 1.0, + step_size: float = 0.05, mse_weight: float = 0, mse_weight_decay: float = 0, mse_weight_decay_steps: int = 0, @@ -73,11 +79,14 @@ def generate_image( run = VQGANCLIPRun( text_input=text_input, vqgan_ckpt=vqgan_ckpt, + clip_model=clip_model, num_steps=num_steps, image_x=image_x, image_y=image_y, seed=seed, cutn=cutn, + cut_pow=cut_pow, + step_size=step_size, init_image=init_image, image_prompts=image_prompts, continue_prev_run=continue_prev_run, @@ -127,7 +136,10 @@ def generate_image( ### Model init ------------------------------------------------------------- if continue_prev_run is True: - run.model_init(init_image=st.session_state["prev_im"]) + try: + run.model_init(init_image=st.session_state["prev_im"]) + except KeyError: + run.model_init(init_image=Image.open("prev_image.png")) elif init_image is not None: run.model_init(init_image=init_image) else: @@ -145,25 +157,52 @@ def generate_image( # However, touching any button resets the app state, making it impossible to # implement a stop button that can still dump output # Thankfully there's a built-in stop button :) + while True: # While loop to accomodate running predetermined steps or running indefinitely - status_text.text(f"Running step {step_counter}") + + # trying to check how long it takes to execute each loop which should give us how long it took to process each image generation iteration. + start = timeit.default_timer() - _, im = run.iterate() + _, im = run.iterate() if num_steps > 0: # skip when num_steps = -1 step_progress_bar.progress((step_counter + 1) / num_steps) else: step_progress_bar.progress(100) + + duration = timeit.default_timer() - start + + if duration >= 1: + speed = "s/it" + else: + speed = "it/s" + duration = 1 / duration + + if num_steps > 0: + total_number_steps = f"/{num_steps}" + percent = f"%{100 * float(step_counter)/float(num_steps)}" + else: + total_number_steps = "" + percent = "" + + status_text.text(f"Running step: {step_counter}{total_number_steps} {percent} | {duration:.2f}{speed}") # At every step, display and save image im_display_slot.image(im, caption="Output image", output_format="PNG") st.session_state["prev_im"] = im + + try: + # Save prev_im.png + im.save(f"prev_image.png", format='PNG') + except (PermissionError,OSError): + # Save prev_im.png + im.save(f"prev_image.png", format='PNG') # ref: https://stackoverflow.com/a/33117447/13095028 - # im_byte_arr = io.BytesIO() - # im.save(im_byte_arr, format="JPEG") - # frames.append(im_byte_arr.getvalue()) # read() + #im_byte_arr = io.BytesIO() + #im.save(im_byte_arr, format="JPEG") + #frames.append(im_byte_arr.getvalue()) # read() frames.append(np.asarray(im)) step_counter += 1 @@ -184,16 +223,28 @@ def generate_image( runoutputdir.mkdir() # Save final image - im.save(runoutputdir / "output.PNG", format="PNG") + im.save(runoutputdir / "output.png", format="PNG") + + # Save all the frames into the step folder so we can see the process better if we want to. + runoutputdir_step_folder = Path(f"{runoutputdir}/steps") + print (f"Saving frames to folder: {runoutputdir_step_folder}") + runoutputdir_step_folder.mkdir() + + frame_number = 0 + for frame in frames: + #print (frame) + frame_data = Image.fromarray(frame) + frame_data.save(f"{runoutputdir_step_folder}/{frame_number}.png", format='PNG') + frame_number += 1 # Save init image if init_image is not None: - init_image.save(runoutputdir / "init-image.JPEG", format="JPEG") + init_image.save(runoutputdir / "init-image.jpeg", format="JPEG") # Save image prompts for count, image_prompt in enumerate(image_prompts): image_prompt.save( - runoutputdir / f"image-prompt-{count}.JPEG", format="JPEG" + runoutputdir / f"image-prompt-{count}.jpeg", format="JPEG" ) # Save animation @@ -211,9 +262,12 @@ def generate_image( "prev_run_id": prev_run_id, "seed": run.seed, "cutn": cutn, + "cut_pow":cut_pow, + "step_size": step_size, "Xdim": image_x, "ydim": image_y, "vqgan_ckpt": vqgan_ckpt, + "clip_model": clip_model, "start_time": run_start_dt.strftime("%Y%m%dT%H%M%S"), "end_time": datetime.datetime.now().strftime("%Y%m%dT%H%M%S"), "mse_weight": mse_weight, @@ -252,8 +306,10 @@ def generate_image( except st.StopException as e: # Dump output to dashboard print(f"Received Streamlit StopException") - status_text.text("Execution interruped, dumping outputs ...") + #status_text.text("Execution interruped, dumping outputs ...") + print("Execution interruped, dumping outputs ...") writer = imageio.get_writer("temp.mp4", fps=24) + for frame in frames: writer.append_data(frame) writer.close() @@ -263,19 +319,32 @@ def generate_image( runoutputdir = outputdir / ( run_start_dt.strftime("%Y%m%dT%H%M%S") + "-" + run_id ) + print ("Saving to folder: ", runoutputdir) runoutputdir.mkdir() # Save final image - im.save(runoutputdir / "output.PNG", format="PNG") + im.save(runoutputdir / "output.png", format="PNG") + + # Save all the frames into the step folder so we can see the process better if we want to. + runoutputdir_step_folder = Path(f"{runoutputdir}/steps") + print (f"Saving frames to folder: {runoutputdir_step_folder}") + runoutputdir_step_folder.mkdir() + + frame_number = 0 + for frame in frames: + #print (frame) + frame_data = Image.fromarray(frame) + frame_data.save(f"{runoutputdir_step_folder}/{frame_number}.png", format='PNG') + frame_number += 1 # Save init image if init_image is not None: - init_image.save(runoutputdir / "init-image.JPEG", format="JPEG") + init_image.save(runoutputdir / "init-image.jpeg", format="JPEG") # Save image prompts for count, image_prompt in enumerate(image_prompts): image_prompt.save( - runoutputdir / f"image-prompt-{count}.JPEG", format="JPEG" + runoutputdir / f"image-prompt-{count}.jpeg", format="JPEG" ) # Save animation @@ -293,9 +362,12 @@ def generate_image( "prev_run_id": prev_run_id, "seed": run.seed, "cutn": cutn, + "cut_pow": cut_pow, + "step_size": step_size, "Xdim": image_x, "ydim": image_y, "vqgan_ckpt": vqgan_ckpt, + "clip_model": clip_model, "start_time": run_start_dt.strftime("%Y%m%dT%H%M%S"), "end_time": datetime.datetime.now().strftime("%Y%m%dT%H%M%S"), "mse_weight": mse_weight, @@ -329,7 +401,9 @@ def generate_image( with open(runoutputdir / "details.json", "w") as f: json.dump(details, f, indent=4) - status_text.text("Done!") # End of run + status_text.text("Done!") # End of run + + if __name__ == "__main__": @@ -343,7 +417,7 @@ def generate_image( args = parser.parse_args() # Select specific GPU if chosen - if args.gpu is not None: + if args.gpu is not None and args.gpu != "cpu": for i in args.gpu.split(","): assert ( int(i) < torch.cuda.device_count() @@ -354,10 +428,13 @@ def generate_image( except RuntimeError: print(traceback.format_exc()) else: + os.environ["CUDA_VISIBLE_DEVICES"] = "" device = None defaults = OmegaConf.load("defaults.yaml") outputdir = Path("output") + #print (outputdir) + if not outputdir.exists(): outputdir.mkdir() @@ -398,6 +475,31 @@ def generate_image( index=default_weight_index, help="Choose which weights to load, trained on different datasets. Make sure the weights and configs are downloaded to `assets/` as per the README!", ) + + use_clip_model = st.sidebar.checkbox( + "Clip Model", + value=defaults["use_clip_model"], + help="Clip Model to use", + ) + clip_model = st.sidebar.empty() + if use_clip_model is True: + clip_model = clip_model.text_input( + "Clip Model:", + value=defaults["clip_model"], + help="""Model Versions:‏‏‎ + ‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎ + Initially, we’ve released one CLIP model based on the Vision Transformer architecture equivalent to ViT-B/32, + along with the RN50 model, using the architecture equivalent to ResNet-50. + As part of the staged release process, we have also released the RN101 model, as well as RN50x4, + a RN50 scaled up 4x according to the EfficientNet scaling rule. In July 2021, we additionally released the RN50x16 and ViT-B/16 models, + and in January 2022, the RN50x64 and ViT-L/14 models were released. Lastly, the ViT-L/14@336px model was released in April 2022. + ‎‎‎‎‎‎‎‎‎‎‎‎‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ + Defautl = ViT-B/32 . + Available Models: """ + str(clip.available_models()), + ) + else: + clip_model = "ViT-B/32" + num_steps = st.sidebar.number_input( "Num steps", value=defaults["num_steps"], @@ -434,11 +536,17 @@ def generate_image( else: seed = None - #cutn = st.sidebar.checkbox( - # "Set Cutn", - # value=defaults["cutn"], - # help="Check to set the number of cuts to pass to CLIP, lower values uses less VRAM, higher values increases the image quality. Will add option to specify the number of cuts", - #) + continue_prev_run = st.sidebar.checkbox( + "Continue previous run", + value=defaults["continue_prev_run"], + help="Use existing image and existing weights for the next run. If yes, ignores 'Use starting image'", + ) + + use_cutout_augmentations = st.sidebar.checkbox( + "Use cutout augmentations", + value=True, + help="Adds cutout augmentatinos in the image generation process. Uses up to additional 4 GiB of GPU memory. Greatly improves image quality. Toggled on by default.", + ) use_cutn = st.sidebar.checkbox( "Use Cutn", @@ -446,6 +554,7 @@ def generate_image( help="Check to set the number of cuts to pass to CLIP, lower values uses less VRAM, higher values increase the image quality. Will add option to specify the number of cuts", ) cutn = st.sidebar.empty() + cut_pow = st.sidebar.empty() if use_cutn is True: cutn = cutn.number_input( "Number of Cuts sent to CLIP", @@ -454,8 +563,35 @@ def generate_image( step=1, help="Specify the number of cuts to pass to CLIP, lower values uses less VRAM but higher values increases the image quality", ) + cut_pow = cut_pow.number_input( + "Cut power.", + value=defaults["cut_pow"], + min_value=0.0001, + step=0.1, + help="Specify the power each cut will have.", + ) else: cutn = 32 + cut_pow = 1.0 + + + custom_step_size = st.sidebar.checkbox( + "Custom Step Size/Learing Rate", + value=defaults["custom_step_size"], + help="Customize the Step Size or Learning Rate value.", + ) + step_size = st.sidebar.empty() + if custom_step_size is True: + step_size = step_size.number_input( + "Custom Step Size or Learning Rate", + value=defaults["step_size"], + min_value=0.0001, + step=0.001, + help="Specify a custom Step Size or Learning Rate to use. Ref: https://en.wikipedia.org/wiki/Learning_rate", + format="%.5f", + ) + else: + step_size = 0.05 use_custom_starting_image = st.sidebar.checkbox( "Use starting image", @@ -499,12 +635,6 @@ def generate_image( else: image_prompts = [] - continue_prev_run = st.sidebar.checkbox( - "Continue previous run", - value=defaults["continue_prev_run"], - help="Use existing image and existing weights for the next run. If yes, ignores 'Use starting image'", - ) - use_mse_reg = st.sidebar.checkbox( "Use MSE regularization", value=defaults["use_mse_regularization"], @@ -552,9 +682,9 @@ def generate_image( "TV loss weight", value=defaults["tv_loss_weight"], min_value=0.0, - step=1e-4, - help="Set weights for TV loss regularization, which encourages spatial smoothness. Ref: https://github.com/jcjohnson/neural-style/issues/302", - format="%.1e", + step=0.0001, + help="Set weights for TV loss regularization, which encourages spatial smoothness, the lower the value the better the result. Some good values are are 0.000085, 0.0001 or 0.0002 Ref: https://github.com/jcjohnson/neural-style/issues/302", + format="%.6f", ) else: tv_loss_weight = 0 @@ -587,10 +717,11 @@ def generate_image( zoom_factor = zoom_factor_widget.number_input( "Zoom factor", value=1.0, - min_value=0.1, - max_value=10.0, - step=0.02, - format="%.2f", + min_value=-100.0, + max_value=100.0, + step=0.0001, + help="Factor to zoom in each frame, 1 is no zoom, less than 1 is zoom out, more than 1 is zoom in.", + format="%.4f", ) transform_interval = transform_interval_widget.number_input( "Iterations per frame", @@ -606,18 +737,12 @@ def generate_image( zoom_factor = 1 transform_interval = 1 - use_cutout_augmentations = st.sidebar.checkbox( - "Use cutout augmentations", - value=True, - help="Adds cutout augmentatinos in the image generation process. Uses up to additional 4 GiB of GPU memory. Greatly improves image quality. Toggled on by default.", - ) - submitted = st.form_submit_button("Run!") # End of form status_text = st.empty() status_text.text("Pending input prompt") - step_progress_bar = st.progress(0) + step_progress_bar = st.progress(0) im_display_slot = st.empty() vid_display_slot = st.empty() @@ -668,11 +793,13 @@ def generate_image( # Inputs text_input=text_input, vqgan_ckpt=radio, + clip_model=clip_model, num_steps=num_steps, image_x=int(image_x), image_y=int(image_y), seed=int(seed) if set_seed is True else None, cutn=int(cutn), + step_size=float(step_size), init_image=init_image, image_prompts=image_prompts, continue_prev_run=continue_prev_run, From 22587a3941ca339cae9a5649ba31aa07ec1ce75e Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Tue, 23 Aug 2022 10:52:41 -0700 Subject: [PATCH 04/11] Config files from previous run that I forgot to include, without them the new changes would throw an error or you would need to manually change these files for it to work. --- .streamlit/config.toml | 2 +- defaults.yaml | 18 ++++++++++++------ 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/.streamlit/config.toml b/.streamlit/config.toml index 011d1f8..5456a34 100644 --- a/.streamlit/config.toml +++ b/.streamlit/config.toml @@ -1,3 +1,3 @@ [server] # Default is 200 MB -maxUploadSize = 10 +maxUploadSize = 5000 diff --git a/defaults.yaml b/defaults.yaml index 9ccdb11..9c03707 100644 --- a/defaults.yaml +++ b/defaults.yaml @@ -1,11 +1,16 @@ # Modify for different systems, e.g. larger default xdim/ydim for more powerful GPUs -num_steps: 500 -Xdim: 511 -ydim: 511 +use_clip_model: false +clip_model: ViT-B/32 +num_steps: -1 +Xdim: 662 +ydim: 360 set_seed: false seed: 0 -use_cutn: true +use_cutn: false cutn: 32 +cut_pow: 1.0 +custom_step_size: false +step_size: 0.05 use_starting_image: false use_image_prompts: false continue_prev_run: false @@ -13,5 +18,6 @@ mse_weight: 0.5 mse_weight_decay: 0.1 mse_weight_decay_steps: 50 use_mse_regularization: false -use_tv_loss_regularization: true -tv_loss_weight: 1e-3 \ No newline at end of file +use_tv_loss_regularization: false +# best values for tv_loss_weight are 0.000085, 0.0001 or 0.0002 +tv_loss_weight: 0.000085 \ No newline at end of file From dbb5546f16aa8a4b1bf4a9e0297b1e112530434e Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Tue, 23 Aug 2022 10:55:20 -0700 Subject: [PATCH 05/11] Changes made to the logic.py file that were part of the changes made two commits ago that I forgot to upload. --- logic.py | 43 +++++++++++++++++++++++++++++++------------ 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/logic.py b/logic.py index 460d96c..a9b938a 100644 --- a/logic.py +++ b/logic.py @@ -63,6 +63,7 @@ def __init__( self, text_input: str = "the first day of the waters", vqgan_ckpt: str = "vqgan_imagenet_f16_16384", + clip_model: str = "ViT-B/32", num_steps: int = 300, image_x: int = 300, image_y: int = 300, @@ -71,10 +72,12 @@ def __init__( continue_prev_run: bool = False, seed: Optional[int] = None, cutn: int = 32, + cut_pow: float = 1.0, + step_size: float = 0.05, mse_weight=0.5, mse_weight_decay=0.1, mse_weight_decay_steps=50, - tv_loss_weight=1e-3, + tv_loss_weight=0.000085, use_cutout_augmentations: bool = True, # use_augs: bool = True, # noise_fac: float = 0.1, @@ -87,11 +90,12 @@ def __init__( rotation_angle: float = 0, zoom_factor: float = 1, transform_interval: int = 10, - device: Optional[torch.device] = None, + device: Optional[torch.device] = "cpu", ) -> None: super().__init__() self.text_input = text_input self.vqgan_ckpt = vqgan_ckpt + self.clip_model = clip_model self.num_steps = num_steps self.image_x = image_x self.image_y = image_y @@ -100,6 +104,9 @@ def __init__( self.continue_prev_run = continue_prev_run self.seed = seed self.cutn = cutn + self.cut_pow = cut_pow + self.step_size = step_size + self.device = device # Setup ------------------------------------------------------------------------------ # Split text by "|" symbol @@ -117,23 +124,25 @@ def __init__( init_image=init_image, init_weight=mse_weight, # clip.available_models() - # ['RN50', 'RN101', 'RN50x4', 'ViT-B/32'] + # ['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64', 'ViT-B/32', 'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px'] # Visual Transformer seems to be the smallest - clip_model="ViT-B/32", + clip_model=clip_model, vqgan_config=f"assets/{vqgan_ckpt}.yaml", vqgan_checkpoint=f"assets/{vqgan_ckpt}.ckpt", - step_size=0.05, - #cutn=32, cutn=cutn, - cut_pow=1.0, + cut_pow=cut_pow, + step_size=step_size, display_freq=50, seed=seed, + device=device, ) - if device is None: - self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + if device is None or device == "cpu": + #self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + device = torch.device("cpu") else: - self.device = device + #self.device = device + self.device = torch.device(f"{device}" if torch.cuda.is_available() else "cpu") print("Using device:", device) @@ -319,11 +328,14 @@ def _ascend_txt(self) -> List: result[f"prompt_loss_{count}"] = prompt(iii) return result - + def iterate(self) -> Tuple[List[float], Image.Image]: if not self.use_scrolling_zooming: # Forward prop self.opt.zero_grad() + #for param in self.model.parameters(): + #param.grad = None + losses = self._ascend_txt() # Grab an image @@ -386,6 +398,9 @@ def iterate(self) -> Tuple[List[float], Image.Image]: for _ in range(self.transform_interval): # Forward prop self.opt.zero_grad() + #for param in self.model.parameters(): + #param.grad = None + losses = self._ascend_txt() # Grab an image @@ -397,9 +412,13 @@ def iterate(self) -> Tuple[List[float], Image.Image]: self.opt.step() with torch.no_grad(): self.z.copy_(self.z.maximum(self.z_min).minimum(self.z_max)) - + # Advance iteration counter self.iterate_counter += 1 + + for param_group in self.opt.param_groups: + #print (param_group) + print (f"Learning Rate: {param_group['lr']}") print( f"Step {self.iterate_counter} losses: {[(i, j.item()) for i, j in losses.items()]}" From eb8577bd35d8a2edee138c3e88355694aec35ae5 Mon Sep 17 00:00:00 2001 From: ZeroCool Date: Tue, 23 Aug 2022 12:58:06 -0700 Subject: [PATCH 06/11] Update issue templates --- .github/ISSUE_TEMPLATE/bug_report.md | 32 +++++++++++++++++++++++ .github/ISSUE_TEMPLATE/feature_request.md | 20 ++++++++++++++ 2 files changed, 52 insertions(+) create mode 100644 .github/ISSUE_TEMPLATE/bug_report.md create mode 100644 .github/ISSUE_TEMPLATE/feature_request.md diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..d9b6166 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,32 @@ +--- +name: Bug report +about: Create a report to help us improve +title: '' +labels: '' +assignees: '' + +--- + +**Describe the bug** +A clear and concise description of what the bug is. + +**To Reproduce** +Steps to reproduce the behavior: +1. Go to '...' +2. Click on '....' +3. Scroll down to '....' +4. See error + +**Expected behavior** +A clear and concise description of what you expected to happen. + +**Screenshots** +If applicable, add screenshots to help explain your problem. + +**Desktop (please complete the following information):** + - OS: [e.g. Linux, Windows 10] + - Browser [e.g. Chrome, Safari, Brave] + - Version [e.g. 22] + +**Additional context** +Add any other context about the problem here. diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 0000000..bbcbbe7 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,20 @@ +--- +name: Feature request +about: Suggest an idea for this project +title: '' +labels: '' +assignees: '' + +--- + +**Is your feature request related to a problem? Please describe.** +A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] + +**Describe the solution you'd like** +A clear and concise description of what you want to happen. + +**Describe alternatives you've considered** +A clear and concise description of any alternative solutions or features you've considered. + +**Additional context** +Add any other context or screenshots about the feature request here. From dffc0fdb93bee7d7ba80fcfc96061ad6e7d51c31 Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Tue, 23 Aug 2022 21:53:07 -0700 Subject: [PATCH 07/11] Changed the percent sign to be shown after the numbers as it should be and limited the float to 2 decimal places so it looks better. This is only shown when using a fixed number of steps and the percent will not show if the number of steps is -1 as when using -1 its supposed to run until you stop it so there is no point in showing the progress bar and the percent completed. --- app.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/app.py b/app.py index ede29e7..333c455 100644 --- a/app.py +++ b/app.py @@ -181,7 +181,7 @@ def generate_image( if num_steps > 0: total_number_steps = f"/{num_steps}" - percent = f"%{100 * float(step_counter)/float(num_steps)}" + percent = f"{100 * float(step_counter)/float(num_steps):.2f}%" else: total_number_steps = "" percent = "" @@ -585,7 +585,7 @@ def generate_image( step_size = step_size.number_input( "Custom Step Size or Learning Rate", value=defaults["step_size"], - min_value=0.0001, + min_value=0.00001, step=0.001, help="Specify a custom Step Size or Learning Rate to use. Ref: https://en.wikipedia.org/wiki/Learning_rate", format="%.5f", From 9da97e82e930fa46a622dd62612b0392ca5b8db4 Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Tue, 23 Aug 2022 22:03:24 -0700 Subject: [PATCH 08/11] Added the assets and samples folder to the gitignore, we do not need to upload these folders to the repository. --- .gitignore | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.gitignore b/.gitignore index 95674ac..c69a763 100644 --- a/.gitignore +++ b/.gitignore @@ -5,9 +5,13 @@ # VQGAN weights assets/*.ckpt assets/*.yaml +assets/* # Outputs output* +# samples +samples/* + # Test data test-samples/ From 7adecea32a64dd08dd344ee3bbabcbb0efe3a87b Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Tue, 23 Aug 2022 22:05:12 -0700 Subject: [PATCH 09/11] Changed st.script_runner.StopException to st.StopException as this is the way to access this exception on the newest version of the library. --- diffusion_app.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diffusion_app.py b/diffusion_app.py index 795f881..c3f1866 100644 --- a/diffusion_app.py +++ b/diffusion_app.py @@ -100,7 +100,7 @@ def generate_image( frames = [] try: - # Try block catches st.script_runner.StopExecution, no need of a dedicated stop button + # Try block catches st.StopExecution, no need of a dedicated stop button # Reason is st.form is meant to be self-contained either within sidebar, or in main body # The way the form is implemented in this app splits the form across both regions # This is intended to prevent the model settings from crowding the main body @@ -181,7 +181,7 @@ def generate_image( status_text.text("Done!") # End of run - except st.script_runner.StopException as e: + except st.StopException as e: # Dump output to dashboard print(f"Received Streamlit StopException") status_text.text("Execution interruped, dumping outputs ...") From 92148369aaaae1a406ca5af5f5490f672f0f2089 Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Wed, 24 Aug 2022 20:24:21 -0700 Subject: [PATCH 10/11] Added an option on the sidebar to customize the optimizer we want to use. By default it will be using the Adam optimizer, just like before, but a few others have been added to the drop menu under this option. --- app.py | 66 ++++++++++++++++++++++++++++++++++++++++++--------- defaults.yaml | 2 ++ logic.py | 61 +++++++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 116 insertions(+), 13 deletions(-) diff --git a/app.py b/app.py index 333c455..e164561 100644 --- a/app.py +++ b/app.py @@ -27,6 +27,7 @@ from omegaconf import OmegaConf import imageio import numpy as np +from retry import retry # Catch import issue, introduced in version 1.1 # Deprecate in a few minor versions @@ -47,6 +48,12 @@ #os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:30" +@retry((PermissionError, OSError), tries=5, delay=2) +def save_preview_image(im): + """Save the preview image and in case we get any error try again.""" + # Save prev_im.png + im.save(f"prev_image.png", format='PNG') + def generate_image( text_input: str = "the first day of the waters", vqgan_ckpt: str = "vqgan_imagenet_f16_16384", @@ -61,6 +68,7 @@ def generate_image( cutn: int = 32, cut_pow: float = 1.0, step_size: float = 0.05, + opt_name: str = "Adam", mse_weight: float = 0, mse_weight_decay: float = 0, mse_weight_decay_steps: int = 0, @@ -87,6 +95,7 @@ def generate_image( cutn=cutn, cut_pow=cut_pow, step_size=step_size, + opt_name=opt_name, init_image=init_image, image_prompts=image_prompts, continue_prev_run=continue_prev_run, @@ -192,13 +201,10 @@ def generate_image( im_display_slot.image(im, caption="Output image", output_format="PNG") st.session_state["prev_im"] = im - try: - # Save prev_im.png - im.save(f"prev_image.png", format='PNG') - except (PermissionError,OSError): - # Save prev_im.png - im.save(f"prev_image.png", format='PNG') - + # We save the preview image using the save_preview_image() function + # so we can try multiple times in case there is an error opening or saving the image. + save_preview_image(im) + # ref: https://stackoverflow.com/a/33117447/13095028 #im_byte_arr = io.BytesIO() #im.save(im_byte_arr, format="JPEG") @@ -262,8 +268,9 @@ def generate_image( "prev_run_id": prev_run_id, "seed": run.seed, "cutn": cutn, - "cut_pow":cut_pow, + "cut_pow": cut_pow, "step_size": step_size, + "opt_name": opt_name, "Xdim": image_x, "ydim": image_y, "vqgan_ckpt": vqgan_ckpt, @@ -302,6 +309,8 @@ def generate_image( json.dump(details, f, indent=4) status_text.text("Done!") # End of run + + vid_display_slot.video("temp.mp4") except st.StopException as e: # Dump output to dashboard @@ -364,6 +373,7 @@ def generate_image( "cutn": cutn, "cut_pow": cut_pow, "step_size": step_size, + "opt_name": opt_name, "Xdim": image_x, "ydim": image_y, "vqgan_ckpt": vqgan_ckpt, @@ -401,9 +411,9 @@ def generate_image( with open(runoutputdir / "details.json", "w") as f: json.dump(details, f, indent=4) - status_text.text("Done!") # End of run - + status_text.text("Done!") # End of run + vid_display_slot.video("temp.mp4") if __name__ == "__main__": @@ -634,6 +644,39 @@ def generate_image( image_prompts = [Image.open(i).convert("RGB") for i in image_prompts] else: image_prompts = [] + + + use_custom_opt = st.sidebar.checkbox( + "Custom Optimizer", + value=defaults["use_custom_opt"], + help="Use a custom optimizer.", + ) + + opt_name = st.sidebar.empty() + if use_custom_opt is True: + optimizer_list = ["Adam","Adadelta","Adagrad","AdamW","Adamax","ASGD","NAdam","RAdam","RMSprop","Rprop","SGD"] + opt_name = st.sidebar.selectbox( + "Optimizer:", optimizer_list, index=optimizer_list.index(defaults["opt_name"]), + help=""" + Defautl = Adam + + List of optimizers: + + Adadelta: Implements Adadelta algorithm. \n + Adagrad: Implements Adagrad algorithm.\n + Adam: Implements Adam algorithm.\n + AdamW: Implements AdamW algorithm.\n + Adamax: Implements Adamax algorithm (a variant of Adam based on infinity norm).\n + ASGD: Implements Averaged Stochastic Gradient Descent.\n + NAdam: Implements NAdam algorithm.\n + RAdam: Implements RAdam algorithm.\n + RMSprop: Implements RMSprop algorithm.\n + Rprop: Implements the resilient backpropagation algorithm.\n + SGD: Implements stochastic gradient descent (optionally with momentum).\n + """, + ) + else: + opt_name = "Adam" use_mse_reg = st.sidebar.checkbox( "Use MSE regularization", @@ -800,6 +843,7 @@ def generate_image( seed=int(seed) if set_seed is True else None, cutn=int(cutn), step_size=float(step_size), + opt_name=opt_name, init_image=init_image, image_prompts=image_prompts, continue_prev_run=continue_prev_run, @@ -816,5 +860,5 @@ def generate_image( device=device, ) - vid_display_slot.video("temp.mp4") + # debug_slot.write(st.session_state) # DEBUG diff --git a/defaults.yaml b/defaults.yaml index 9c03707..24636f5 100644 --- a/defaults.yaml +++ b/defaults.yaml @@ -11,6 +11,8 @@ cutn: 32 cut_pow: 1.0 custom_step_size: false step_size: 0.05 +use_custom_opt: false +opt_name: Adam use_starting_image: false use_image_prompts: false continue_prev_run: false diff --git a/logic.py b/logic.py index a9b938a..3c112e5 100644 --- a/logic.py +++ b/logic.py @@ -21,8 +21,57 @@ import cv2 import numpy as np import kornia.augmentation as K +from torch_optimizer import DiffGrad, AdamP, RAdam +# Set the optimiser +def get_opt(opt_name, z, opt_lr): + """ + List of optimizers + Adadelta: Implements Adadelta algorithm. + Adagrad: Implements Adagrad algorithm. + Adam: Implements Adam algorithm. + AdamW: Implements AdamW algorithm. + Adamax: Implements Adamax algorithm (a variant of Adam based on infinity norm). + ASGD: Implements Averaged Stochastic Gradient Descent. + NAdam: Implements NAdam algorithm. + RAdam: Implements RAdam algorithm. + RMSprop: Implements RMSprop algorithm. + Rprop: Implements the resilient backpropagation algorithm. + SGD: Implements stochastic gradient descent (optionally with momentum).""" + + if opt_name == "Adam": + opt = optim.Adam([z], lr=opt_lr) + elif opt_name == "AdamW": + opt = optim.AdamW([z], lr=opt_lr) + elif opt_name == "Adagrad": + opt = optim.Adagrad([z], lr=opt_lr) + elif opt_name == "Adamax": + opt = optim.Adamax([z], lr=opt_lr) + elif opt_name == "AdamP": + opt = AdamP([z], lr=opt_lr) + elif opt_name == "Adadelta": + opt = optim.Adadelta([z], lr=opt_lr, eps=1e-9, weight_decay=1e-9) + elif opt_name == "ASGD": + opt = optim.ASGD([z], lr=opt_lr) + elif opt_name == "DiffGrad": + opt = DiffGrad([z], lr=opt_lr, eps=1e-9, weight_decay=1e-9) + elif opt_name == "NAdam": + opt = optim.NAdam([z], lr=opt_lr) + elif opt_name == "RAdam": + opt = RAdam([z], lr=opt_lr) + elif opt_name == "RMSprop": + opt = optim.RMSprop([z], lr=opt_lr) + elif opt_name == "Rprop": + opt = optim.Rprop([z], lr=opt_lr) + elif opt_name == "SGD": + opt = optim.SGD([z], lr=opt_lr) + + else: + print(f"Unknown optimiser: {opt_name} | Are choices broken?") + opt = optim.Adam([z], lr=opt_lr) + return opt + class Run: """ Subclass this to house your own implementation of CLIP-based image generation @@ -74,6 +123,7 @@ def __init__( cutn: int = 32, cut_pow: float = 1.0, step_size: float = 0.05, + opt_name: str = "Adam", mse_weight=0.5, mse_weight_decay=0.1, mse_weight_decay_steps=50, @@ -106,6 +156,7 @@ def __init__( self.cutn = cutn self.cut_pow = cut_pow self.step_size = step_size + self.opt_name = opt_name self.device = device # Setup ------------------------------------------------------------------------------ @@ -132,6 +183,7 @@ def __init__( cutn=cutn, cut_pow=cut_pow, step_size=step_size, + opt_name=opt_name, display_freq=50, seed=seed, device=device, @@ -167,6 +219,8 @@ def __init__( self.rotation_angle = rotation_angle self.zoom_factor = zoom_factor self.transform_interval = transform_interval + + def load_model( self, prev_model: nn.Module = None, prev_perceptor: nn.Module = None @@ -252,7 +306,8 @@ def model_init(self, init_image: Image.Image = None) -> None: self.z = self.z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2) self.z_orig = self.z.clone() self.z.requires_grad_(True) - self.opt = optim.Adam([self.z], lr=self.args.step_size) + #self.opt = optim.Adam([self.z], lr=self.args.step_size) + self.opt = get_opt(self.opt_name, self.z, self.args.step_size) self.normalize = transforms.Normalize( mean=[0.48145466, 0.4578275, 0.40821073], @@ -393,7 +448,9 @@ def iterate(self) -> Tuple[List[float], Image.Image]: TF.to_tensor(transformed_im).to(self.device).unsqueeze(0) * 2 - 1 ) self.z.requires_grad_(True) - self.opt = optim.Adam([self.z], lr=self.args.step_size) + + #self.opt = optim.Adam([self.z], lr=self.args.step_size) + self.opt = get_opt(self.opt_name, self.z, self.args.step_size) for _ in range(self.transform_interval): # Forward prop From 29e90b2fd3bc470ef417d920719a729a7a38d874 Mon Sep 17 00:00:00 2001 From: ZeroCool940711 Date: Wed, 24 Aug 2022 20:28:05 -0700 Subject: [PATCH 11/11] Added torch-optimizer and retry to the list of pip dependencies, torch-optimizer is used for a few options on the list of customizable optimizers from the previous commit while retry is used to avoid a small issue when saving the preview image to disk where some times it could make the whole app crash if the file was been written to fast to disk and an external image viewer had the image open, this would raise a PermissionError or OSError, with the retry library we give it a couple of tries without crashing the app. --- environment.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/environment.yml b/environment.yml index 08cc1b5..615b7eb 100644 --- a/environment.yml +++ b/environment.yml @@ -32,6 +32,8 @@ dependencies: # - imgtag - einops - transformers + - torch-optimizer + - retry - git+https://github.com/openai/CLIP # For guided diffusion - lpips