fix broken generate and add progress bars

README.md

@@ -42,15 +42,24 @@ Simplified start script: `automatic.sh`
 
 Start in default mode with optimizations enabled  
 
-    SD server: optimized
-    Version: 56f779a9 Sat Feb 25 14:04:19 2023 -0500
-    Repository: https://github.com/vladmandic/automatic
-    Last Merge: Sun Feb 19 10:11:25 2023 -0500 Merge pull request #37 from AUTOMATIC1111/master
-    System
-    - Platform: Ubuntu 22.04.1 LTS 5.15.90.1-microsoft-standard-WSL2 x86_64
-    - nVIDIA: NVIDIA GeForce RTX 3060, 528.49
-    - Python: 3.10.6 Torch: 2.0.0.dev20230224+cu118 CUDA: 11.8 cuDNN: 8700 GPU: NVIDIA GeForce RTX 3060 Arch: (8, 6)
-    Launching Web UI
+      Stable Diffusion server: optimized
+      Version: a4d00060 Sun Mar 26 10:28:05 2023 -0400
+      Repository: https://github.com/vladmandic/automatic
+      Platform: Ubuntu 22.04.2 LTS 5.15.90.1-microsoft-standard-WSL2 x86_64
+      Installing requirements for Web UI
+      Launching Web UI with arguments: --cors-allow-origins=http://127.0.0.1:7860 --ckpt models/v1-5-pruned-emaonly.safetensors
+      Torch 2.0.0+cu118 CUDA 11.8 cuDNN 8700
+      GPU NVIDIA GeForce RTX 3060 VRAM 12288 Arch (8, 6) Cores 28
+      Running on local URL:  http://127.0.0.1:7860
+      Loading weights: models/v1-5-pruned-emaonly.safetensors ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/4.3 GB -:--:--
+      Creating model from config: /home/vlado/dev/automatic/configs/v1-inference.yaml
+      DiffusionWrapper has 859.52 M params.
+      Loading weights: models/VAE/vae-ft-mse-840000-ema-pruned.ckpt ━━━━━━━━━━━━━━━━━━━━━━━ 0.0/334.7 MB -:--:--
+      Applying scaled dot product cross attention optimization.
+      Textual inversion embeddings loaded(2): ti-mia, ti-vlado
+      Model loaded in 1.6s (load weights: 0.1s, create model: 0.3s, apply weights: 0.4s, load vae: 0.3s, device move: 0.5s).
+      Startup time: 11.8s (import torch: 1.7s, import libraries: 1.0s, list models: 1.9s, load scripts: 1.0s, create ui: 4.4s, load checkpoint: 1.7s).
+      Progress 6.55it/s ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00 0:00:04
 
 > ./automatic.sh clean  
 

modules/generation_parameters_copypaste.py

@@ -325,7 +325,6 @@ def create_override_settings_dict(text_pairs):
     """
 
     res = {}
-
     params = {}
     for pair in text_pairs:
         k, v = pair.split(":", maxsplit=1)
@@ -417,5 +416,3 @@ def connect_paste(button, paste_fields, input_comp, override_settings_component,
         inputs=[],
         outputs=[],
     )
-
-

modules/img2img.py

@@ -77,7 +77,7 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args):
                 processed_image.save(os.path.join(output_dir, filename))
 
 
-def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, override_settings_texts, *args):
+def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, height: int, width: int, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, override_settings_texts, *args):
     override_settings = create_override_settings_dict(override_settings_texts)
 
     is_batch = mode == 5
@@ -127,7 +127,7 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
         subseed_strength=subseed_strength,
         seed_resize_from_h=seed_resize_from_h,
         seed_resize_from_w=seed_resize_from_w,
-        seed_enable_extras=seed_enable_extras,
+        seed_enable_extras=True,
         sampler_name=sd_samplers.samplers_for_img2img[sampler_index].name,
         batch_size=batch_size,
         n_iter=n_iter,

modules/models/diffusion/uni_pc/uni_pc.py

@@ -1,7 +1,8 @@
 import torch
 import torch.nn.functional as F
 import math
-from tqdm.auto import trange
+import time
+from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
 
 
 class NoiseScheduleVP:
@@ -753,44 +754,50 @@ class UniPC:
             timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
             #print(f"Running UniPC Sampling with {timesteps.shape[0]} timesteps, order {order}")
             assert timesteps.shape[0] - 1 == steps
-            with torch.no_grad():
-                vec_t = timesteps[0].expand((x.shape[0]))
-                model_prev_list = [self.model_fn(x, vec_t)]
-                t_prev_list = [vec_t]
-                # Init the first `order` values by lower order multistep DPM-Solver.
-                for init_order in range(1, order):
-                    vec_t = timesteps[init_order].expand(x.shape[0])
-                    x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
-                    if model_x is None:
-                        model_x = self.model_fn(x, vec_t)
-                    if self.after_update is not None:
-                        self.after_update(x, model_x)
-                    model_prev_list.append(model_x)
-                    t_prev_list.append(vec_t)
-                for step in trange(order, steps + 1):
-                    vec_t = timesteps[step].expand(x.shape[0])
-                    if lower_order_final:
-                        step_order = min(order, steps + 1 - step)
-                    else:
-                        step_order = order
-                    #print('this step order:', step_order)
-                    if step == steps:
-                        #print('do not run corrector at the last step')
-                        use_corrector = False
-                    else:
-                        use_corrector = True
-                    x, model_x =  self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, step_order, use_corrector=use_corrector)
-                    if self.after_update is not None:
-                        self.after_update(x, model_x)
-                    for i in range(order - 1):
-                        t_prev_list[i] = t_prev_list[i + 1]
-                        model_prev_list[i] = model_prev_list[i + 1]
-                    t_prev_list[-1] = vec_t
-                    # We do not need to evaluate the final model value.
-                    if step < steps:
+            with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn()) as progress:
+                task = progress.add_task(description="Initializing", total=steps)
+                t = time.time()
+                with torch.no_grad():
+                    vec_t = timesteps[0].expand((x.shape[0]))
+                    model_prev_list = [self.model_fn(x, vec_t)]
+                    t_prev_list = [vec_t]
+                    # Init the first `order` values by lower order multistep DPM-Solver.
+                    for init_order in range(1, order):
+                        vec_t = timesteps[init_order].expand(x.shape[0])
+                        x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
                         if model_x is None:
                             model_x = self.model_fn(x, vec_t)
-                        model_prev_list[-1] = model_x
+                        if self.after_update is not None:
+                            self.after_update(x, model_x)
+                        model_prev_list.append(model_x)
+                        t_prev_list.append(vec_t)                      
+                        progress.update(task, advance=1, description=f"Progress {round(len(vec_t) * init_order / (time.time() - t), 2)}it/s")
+                    # for step in trange(order, steps + 1):
+                    for step in range(order, steps + 1):
+                        vec_t = timesteps[step].expand(x.shape[0])
+                        if lower_order_final:
+                            step_order = min(order, steps + 1 - step)
+                        else:
+                            step_order = order
+                        #print('this step order:', step_order)
+                        if step == steps:
+                            #print('do not run corrector at the last step')
+                            use_corrector = False
+                        else:
+                            use_corrector = True
+                        x, model_x =  self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, step_order, use_corrector=use_corrector)
+                        if self.after_update is not None:
+                            self.after_update(x, model_x)
+                        for i in range(order - 1):
+                            t_prev_list[i] = t_prev_list[i + 1]
+                            model_prev_list[i] = model_prev_list[i + 1]
+                        t_prev_list[-1] = vec_t
+                        # We do not need to evaluate the final model value.
+                        if step < steps:
+                            if model_x is None:
+                                model_x = self.model_fn(x, vec_t)
+                            model_prev_list[-1] = model_x
+                        progress.update(task, advance=1, description=f"Progress {round(len(vec_t) * step / (time.time() - t), 2)}it/s")
         else:
             raise NotImplementedError()
         if denoise_to_zero:

modules/sd_models.py

@@ -18,8 +18,7 @@ from modules.paths import models_path
 from modules.sd_hijack_inpainting import do_inpainting_hijack
 from modules.timer import Timer
 
-import rich
-from rich import print
+from rich import print, progress
 
 model_dir = "Stable-diffusion"
 model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
@@ -239,19 +238,21 @@ def read_metadata_from_safetensors(filename):
 
 
 def read_state_dict(checkpoint_file):
-    if 'v1-5-pruned-emaonly.safetensors' in checkpoint_file:
-        pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu')
-    else:
+    with progress.open(checkpoint_file, 'rb', description=f'Loading weights: [cyan]{checkpoint_file}', auto_refresh=True) as f:
         _, extension = os.path.splitext(checkpoint_file)
-        with rich.progress.open(checkpoint_file, 'rb') as f:
+        if 'v1-5-pruned-emaonly.safetensors' or 'vae-ft-mse-840000-ema-pruned.ckpt' in checkpoint_file:
+            if extension.lower() == ".safetensors":
+                pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu')
+            else:
+                pl_sd = torch.load(checkpoint_file, map_location='cpu')
+        else:
             if extension.lower() == ".safetensors":
                 buffer = f.read()
                 pl_sd = safetensors.torch.load(buffer)
             else:
                 buffer = io.BytesIO(f.read())
                 pl_sd = torch.load(buffer, map_location='cpu')
-
-    sd = get_state_dict_from_checkpoint(pl_sd)
+        sd = get_state_dict_from_checkpoint(pl_sd)
     return sd
 
 
@@ -264,7 +265,6 @@ def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer):
         print(f"Loading weights from cache")
         return checkpoints_loaded[checkpoint_info]
 
-    print(f"Loading weights from {checkpoint_info.filename}")
     res = read_state_dict(checkpoint_info.filename)
     timer.record("load weights")
 

modules/sd_vae.py

@@ -139,7 +139,6 @@ def load_vae(model, vae_file=None, vae_source="from unknown source"):
             _load_vae_dict(model, checkpoints_loaded[vae_file])
         else:
             assert os.path.isfile(vae_file), f"VAE {vae_source} doesn't exist: {vae_file}"
-            print(f"Loading VAE weights {vae_source}: {vae_file}")
             store_base_vae(model)
 
             vae_dict_1 = load_vae_dict(vae_file)

modules/txt2img.py

@@ -9,9 +9,8 @@ import modules.processing as processing
 from modules.ui import plaintext_to_html
 
 
-def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, override_settings_texts, *args):
+def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, override_settings_texts, *args):
     override_settings = create_override_settings_dict(override_settings_texts)
-
     p = StableDiffusionProcessingTxt2Img(
         sd_model=shared.sd_model,
         outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,
@@ -24,7 +23,7 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step
         subseed_strength=subseed_strength,
         seed_resize_from_h=seed_resize_from_h,
         seed_resize_from_w=seed_resize_from_w,
-        seed_enable_extras=seed_enable_extras,
+        seed_enable_extras=True,
         sampler_name=sd_samplers.samplers[sampler_index].name,
         batch_size=batch_size,
         n_iter=n_iter,

modules/ui.py

@@ -999,10 +999,7 @@ def create_ui():
         return interp_descriptions[value]
 
     with gr.Blocks(analytics_enabled=False) as train_interface:
-        with gr.Row().style(equal_height=False):
-            gr.HTML(value="<p style='margin-bottom: 0.7em'>See <b><a href=\"https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Textual-Inversion\">wiki</a></b> for detailed explanation.</p>")
-
-        with gr.Row(variant="compact").style(equal_height=False):
+        with gr.Column(elem_id='ti_train_container'):
             with gr.Tabs(elem_id="train_tabs"):
 
                 with gr.Tab(label="Merge models") as modelmerger_interface:

modules/ui_common.py

@@ -135,7 +135,7 @@ Requested path was: {f}
         generation_info = None
         with gr.Column():
             with gr.Row(elem_id=f"image_buttons_{tabname}", elem_classes="image-buttons"):
-                open_folder_button = gr.Button(folder_symbol, visible=not shared.cmd_opts.hide_ui_dir_config)
+                open_folder_button = gr.Button('Load', visible=not shared.cmd_opts.hide_ui_dir_config)
 
                 if tabname != "extras":
                     save = gr.Button('Save', elem_id=f'save_{tabname}')

user.css

@@ -39,7 +39,7 @@ div.gradio-container.dark > div.w-full.flex.flex-col.min-h-screen > div { backgr
 .eta-bar { display: none !important }
 .extra-network-cards .card .actions .name { font-weight: 400; font-size: 1.2rem; }
 .extra-network-cards .card { width: 10em; height: 10em; box-shadow: none; }
-.extra-network-cards .card:hover { transform: scale(2); transition: all 0.3s ease; z-index: 99; box-shadow: none; }
+.extra-network-cards .card:hover { transform: scale(1.3); transition: all 0.3s ease; z-index: 99; box-shadow: none; }
 .feather .feather-image { display: none }
 .gap-2 { padding-top: 8px; }
 .gr-box > div > div > input.gr-text-input { right: 0; width: 4em; padding: 0; top: -12px; border: none; max-height: 20px; }
@@ -67,8 +67,8 @@ div.gradio-container.dark > div.w-full.flex.flex-col.min-h-screen > div { backgr
 /* gradio elements overrides */
 #div.gradio-container.dark { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-#img2img_neg_prompt > label > textarea { font-size: 1.2rem; }
-#img2img_prompt > label > textarea { font-size: 1.2rem; }
+#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt { background-color: black; box-shadow: 4px 4px 4px 0px #333333  !important; }
+#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea { font-size: 1.2rem; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: #111111; padding-top: 16px; }
 #interrogate, #deepbooru { margin: 16px 0px 16px 0px; max-width: 100px; max-height: 74px; font-weight: normal; font-size: 14px; }
 #lightboxModal { background-color: rgba(20, 20, 20, 0.8) }
@@ -85,10 +85,7 @@ div.gradio-container.dark > div.w-full.flex.flex-col.min-h-screen > div { backgr
 #txt2img_gallery, #img2img_gallery, #extras_gallery { background: black !important; padding: 0; margin: 0; object-fit: contain; box-shadow: none; min-height: 0; }
 #txt2img_generate, #img2img_generate { margin-left: 10px; height: 36px; border: none; border-radius: 0; min-height: 36px; padding: 0; }
 #txt2img_interrupt, #img2img_interrupt, #txt2img_skip, #img2img_skip { margin-left: 10px; height: 36px; min-width: 116px; max-width: 116px; border: none; border-radius: 0; background-color: var(--inactive-color); margin-top: 46px; display: block !important; padding: 0; }
-#txt2img_neg_prompt > label > textarea { font-size: 1.2rem; }
 #txt2img_progress_row > div { min-width: var(--left-column); max-width: var(--left-column); }
-#txt2img_prompt > label > textarea { font-size: 1.2rem; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt { background-color: black; }
 #txt2img_results, #img2img_results, #extras_results { background-color: black; padding: 0; }
 #txt2img_seed_row { padding: 0; margin-top: 8px; }
 #txt2img_settings { min-width: var(--left-column); max-width: var(--left-column); background-color: #111111; padding-top: 16px; }
@@ -101,13 +98,14 @@ div.gradio-container.dark > div.w-full.flex.flex-col.min-h-screen > div { backgr
 /* custom elements overrides */
 #steps-animation, #controlnet { border-width: 0; }
 
+/* gradio built-in theme */
 .dark {
   --body-background-fill: var(--background-fill-primary);
   --body-text-color: var(--neutral-100);
   --color-accent-soft: var(--neutral-700);
-  --background-fill-primary: var(--neutral-950);
-  --background-fill-secondary: var(--neutral-900);
-  --border-color-accent: var(--neutral-600);
+  --background-fill-primary: none;
+  --background-fill-secondary: none;
+  --border-color-accent: black;
   --border-color-primary: black;
   --link-text-color-active: var(--secondary-500);
   --link-text-color: var(--secondary-500);
@@ -166,8 +164,8 @@ div.gradio-container.dark > div.w-full.flex.flex-col.min-h-screen > div { backgr
   --slider_color: None;
   --stat-background-fill: linear-gradient(to right, var(--primary-400), var(--primary-600));
   --table-border-color: var(--neutral-700);
-  --table-even-background-fill: var(--neutral-950);
-  --table-odd-background-fill: var(--neutral-900);
+  --table-even-background-fill: #222222;
+  --table-odd-background-fill: #333333;
   --table-row-focus: var(--color-accent-soft);
   --button-border-width: var(--input-border-width);
   --button-cancel-background-fill: linear-gradient(to bottom right, #dc2626, #b91c1c);
@@ -176,13 +174,13 @@ div.gradio-container.dark > div.w-full.flex.flex-col.min-h-screen > div { backgr
   --button-cancel-border-color-hover: var(--button-cancel-border-color);
   --button-cancel-text-color: white;
   --button-cancel-text-color-hover: var(--button-cancel-text-color);
-  --button-primary-background-fill: linear-gradient(to bottom right, var(--primary-500), var(--primary-600));
+  --button-primary-background-fill: linear-gradient(to bottom right, var(--primary-500), var(--primary-800));
   --button-primary-background-fill-hover: linear-gradient(to bottom right, var(--primary-500), var(--primary-500));
   --button-primary-border-color: var(--primary-500);
   --button-primary-border-color-hover: var(--button-primary-border-color);
   --button-primary-text-color: white;
   --button-primary-text-color-hover: var(--button-primary-text-color);
-  --button-secondary-background-fill: linear-gradient(to bottom right, var(--neutral-600), var(--neutral-700));
+  --button-secondary-background-fill: linear-gradient(to bottom right, var(--neutral-600), var(--neutral-800));
   --button-secondary-background-fill-hover: linear-gradient(to bottom right, var(--neutral-600), var(--neutral-600));
   --button-secondary-border-color: var(--neutral-600);
   --button-secondary-border-color-hover: var(--button-secondary-border-color);
@@ -248,7 +246,7 @@ div.gradio-container.dark > div.w-full.flex.flex-col.min-h-screen > div { backgr
   --body-text-weight: 400;
   --embed-radius: var(--radius-lg);
   --color-accent: var(--primary-500);
-  --shadow-drop: rgba(0,0,0,0.05) 0px 1px 2px 0px;
+  --shadow-drop: 0;
   --shadow-drop-lg: 0 1px 3px 0 rgb(0 0 0 / 0.1), 0 1px 2px -1px rgb(0 0 0 / 0.1);
   --shadow-inset: rgba(0,0,0,0.05) 0px 2px 4px 0px inset;
   --block-border-width: 1px;
@@ -302,8 +300,8 @@ div.gradio-container.dark > div.w-full.flex.flex-col.min-h-screen > div { backgr
   --button-large-padding: var(--spacing-lg) calc(2 * var(--spacing-lg));
   --button-large-radius: var(--radius-lg);
   --button-large-text-size: var(--text-lg);
-  --button-large-text-weight: 600;
-  --button-shadow: var(--shadow-drop);
+  --button-large-text-weight: 400;
+  --button-shadow: 4px 4px 4px 0px #333333;
   --button-shadow-active: var(--shadow-inset);
   --button-shadow-hover: var(--shadow-drop-lg);
   --button-small-padding: var(--spacing-sm) calc(2 * var(--spacing-sm));