Merge branch 'master' of https://github.com/vladmandic/automatic

modules/cmd_args.py

@@ -17,6 +17,7 @@ parser.add_argument("--lowram", action='store_true', help="Load checkpoint weigh
 parser.add_argument("--ckpt", type=str, default=sd_model_file, help="Path to checkpoint of stable diffusion model to load immediately",)
 parser.add_argument('--vae', type=str, help='Path to checkpoint of stable diffusion VAE model to load immediately', default=None)
 parser.add_argument("--data-dir", type=str, default=os.path.dirname(os.path.dirname(os.path.realpath(__file__))), help="Base path where all user data is stored")
+parser.add_argument("--models-dir", type=str, default="models", help="Nase path where all models are stored",)
 
 parser.add_argument("--allow-code", action='store_true', help="Allow custom script execution")
 parser.add_argument("--share", action='store_true', help="Enable to make the UI accessible through Gradio site")
@@ -43,6 +44,7 @@ parser.add_argument("--no-hashing", action='store_true', help="Disable sha256 ha
 parser.add_argument("--no-download-sd-model", action='store_true', help="Disable download of default model even if no model is found", default=False)
 parser.add_argument("--profile", action='store_true', help="Run profiler, default: %(default)s")
 parser.add_argument("--disable-queue", action='store_true', help="Disable Gradio queues and force use of HTTP instead of WebSockets, default: %(default)s")
+parser.add_argument("--rollback-vae", action='store_true', help="trying to roll back vae when produced nan image, need to enable nan check", default=False)
 parser.add_argument("--token-merging", action='store_true', help="Provides speed and memory improvements by merging redundant tokens. This has a more pronounced effect on higher resolutions.", default=False)
 
 

modules/generation_parameters_copypaste.py

@@ -3,10 +3,10 @@ import io
 import os
 import re
 
+from PIL import Image
 import gradio as gr
 from modules.paths import data_path
 from modules import shared, ui_tempdir, script_callbacks
-from PIL import Image
 
 re_param_code = r'\s*([\w ]+):\s*("(?:\\"[^,]|\\"|\\|[^\"])+"|[^,]*)(?:,|$)'
 re_param = re.compile(re_param_code)
@@ -251,7 +251,7 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model
         lines.append(lastline)
         lastline = ''
 
-    for i, line in enumerate(lines):
+    for _i, line in enumerate(lines):
         line = line.strip()
         if line.startswith("Negative prompt:"):
             done_with_prompt = True
@@ -382,6 +382,8 @@ def connect_paste(button, paste_fields, input_comp, override_settings_component,
             if os.path.exists(filename):
                 with open(filename, "r", encoding="utf8") as file:
                     prompt = file.read()
+            else:
+                prompt = ''
 
         params = parse_generation_parameters(prompt)
         script_callbacks.infotext_pasted_callback(prompt, params)

modules/paths_internal.py

@@ -15,7 +15,6 @@ parser_pre.add_argument("--data-dir", type=str, default=os.path.dirname(os.path.
 parser_pre.add_argument("--models-dir", type=str, default="models", help="base path where all models are stored",)
 cmd_opts_pre = parser_pre.parse_known_args()[0]
 data_path = cmd_opts_pre.data_dir
-
 models_path = os.path.join(data_path, cmd_opts_pre.models_dir)
 extensions_dir = os.path.join(data_path, "extensions")
 extensions_builtin_dir = os.path.join(script_path, "extensions-builtin")

modules/processing.py

@@ -442,9 +442,8 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
 
 
 def decode_first_stage(model, x):
-    with devices.autocast(disable=x.dtype == devices.dtype_vae):
+    with devices.autocast(disable = x.dtype==devices.dtype_vae):
         x = model.decode_first_stage(x)
-
     return x
 
 
@@ -690,8 +689,20 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, prompts=prompts)
 
             x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]
-            for x in x_samples_ddim:
-                devices.test_for_nans(x, "vae")
+            try:
+                for x in x_samples_ddim:
+                    devices.test_for_nans(x, "vae")
+            except devices.NansException as e:
+                if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae and shared.cmd_opts.rollback_vae:
+                    print('\nA tensor with all NaNs was produced in VAE, try converting to bf16.')
+                    devices.dtype_vae = torch.bfloat16
+                    vae_file, vae_source = sd_vae.resolve_vae(p.sd_model.sd_model_checkpoint)
+                    sd_vae.load_vae(p.sd_model, vae_file, vae_source)
+                    x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]
+                    for x in x_samples_ddim:
+                        devices.test_for_nans(x, "vae")
+                else:
+                    raise e
 
             x_samples_ddim = torch.stack(x_samples_ddim).float()
             x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)

modules/progress.py

@@ -1,11 +1,7 @@
 import base64
 import io
 import time
-
-from pydantic import BaseModel, Field
-
-from modules.shared import opts
-
+from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
 import modules.shared as shared
 
 
@@ -15,18 +11,15 @@ finished_tasks = []
 
 
 def start_task(id_task):
-    global current_task
-
+    global current_task # pylint: disable=global-statement
     current_task = id_task
     pending_tasks.pop(id_task, None)
 
 
 def finish_task(id_task):
-    global current_task
-
+    global current_task # pylint: disable=global-statement
     if current_task == id_task:
         current_task = None
-
     finished_tasks.append(id_task)
     if len(finished_tasks) > 16:
         finished_tasks.pop(0)
@@ -60,39 +53,31 @@ def progressapi(req: ProgressRequest):
     active = req.id_task == current_task
     queued = req.id_task in pending_tasks
     completed = req.id_task in finished_tasks
-
     if not active:
         return ProgressResponse(active=active, queued=queued, completed=completed, id_live_preview=-1, textinfo="In queue..." if queued else "Waiting...")
-
     progress = 0
-
     job_count, job_no = shared.state.job_count, shared.state.job_no
     sampling_steps, sampling_step = shared.state.sampling_steps, shared.state.sampling_step
-
     if job_count > 0:
         progress += job_no / job_count
     if sampling_steps > 0 and job_count > 0:
         progress += 1 / job_count * sampling_step / sampling_steps
-
     progress = min(progress, 1)
-
     elapsed_since_start = time.time() - shared.state.time_start
     predicted_duration = elapsed_since_start / progress if progress > 0 else None
     eta = predicted_duration - elapsed_since_start if predicted_duration is not None else None
-
     id_live_preview = req.id_live_preview
     shared.state.set_current_image()
-    if opts.live_previews_enable and shared.state.id_live_preview != req.id_live_preview:
+    if shared.opts.live_previews_enable and shared.state.id_live_preview != req.id_live_preview:
         image = shared.state.current_image
         if image is not None:
             buffered = io.BytesIO()
-            image.save(buffered, format="png")
-            live_preview = 'data:image/png;base64,' + base64.b64encode(buffered.getvalue()).decode("ascii")
+            fmt = 'jpeg' if shared.opts.samples_format == 'jpg' else shared.opts.samples_format
+            image.save(buffered, format=fmt)
+            live_preview = f'data:image/{fmt};base64,{base64.b64encode(buffered.getvalue()).decode("ascii")}'
             id_live_preview = shared.state.id_live_preview
         else:
             live_preview = None
     else:
         live_preview = None
-
     return ProgressResponse(active=active, queued=queued, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo)
-

modules/sd_samplers_common.py

@@ -28,14 +28,12 @@ approximation_indexes = {"Full": 0, "Approx NN": 1, "Approx cheap": 2}
 def single_sample_to_image(sample, approximation=None):
     if approximation is None:
         approximation = approximation_indexes.get(opts.show_progress_type, 0)
-
     if approximation == 2:
         x_sample = sd_vae_approx.cheap_approximation(sample)
     elif approximation == 1:
         x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
     else:
         x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0]
-
     x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
     x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
     x_sample = x_sample.astype(np.uint8)

modules/sd_vae.py

@@ -197,6 +197,8 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
         sd_model.to(devices.cpu)
 
     sd_hijack.model_hijack.undo_hijack(sd_model)
+    if shared.cmd_opts.rollback_vae and devices.dtype_vae == torch.bfloat16:
+        devices.dtype_vae = torch.float16
 
     load_vae(sd_model, vae_file, vae_source)
 

modules/shared.py

@@ -112,7 +112,6 @@ class State:
             "sampling_step": self.sampling_step,
             "sampling_steps": self.sampling_steps,
         }
-
         return obj
 
     def begin(self):
@@ -142,20 +141,17 @@ class State:
         """sets self.current_image from self.current_latent if enough sampling steps have been made after the last call to this"""
         if not parallel_processing_allowed:
             return
-
         if self.sampling_step - self.current_image_sampling_step >= opts.show_progress_every_n_steps and opts.live_previews_enable and opts.show_progress_every_n_steps != -1:
             self.do_set_current_image()
 
     def do_set_current_image(self):
         if self.current_latent is None:
             return
-
         import modules.sd_samplers # pylint: disable=W0621
         if opts.show_progress_grid:
             self.assign_current_image(modules.sd_samplers.samples_to_image_grid(self.current_latent))
         else:
             self.assign_current_image(modules.sd_samplers.sample_to_image(self.current_latent))
-
         self.current_image_sampling_step = self.sampling_step
 
     def assign_current_image(self, image):
@@ -425,8 +421,8 @@ options_templates.update(options_section(('ui', "Live previews"), {
     "show_progressbar": OptionInfo(True, "Show progressbar"),
     "live_previews_enable": OptionInfo(True, "Show live previews of the created image"),
     "show_progress_grid": OptionInfo(True, "Show previews of all images generated in a batch as a grid"),
-    "show_progress_every_n_steps": OptionInfo(-1, "Show new live preview image every N sampling steps. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
-    "show_progress_type": OptionInfo("Full", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Approx NN", "Approx cheap"]}),
+    "show_progress_every_n_steps": OptionInfo(1, "Show new live preview image every N sampling steps. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
+    "show_progress_type": OptionInfo("Approx NN", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Approx NN", "Approx cheap"]}),
     "live_preview_content": OptionInfo("Combined", "Live preview subject", gr.Radio, {"choices": ["Combined", "Prompt", "Negative prompt"]}),
     "live_preview_refresh_period": OptionInfo(250, "Progressbar/preview update period, in milliseconds")
 }))

modules/styles.py

@@ -49,7 +49,8 @@ class StyleDatabase:
         self.styles.clear()
 
         if not os.path.exists(self.path):
-            return
+            print(f'Creating styles database: {self.path}')
+            self.save_styles(self.path)
 
         with open(self.path, "r", encoding="utf-8-sig", newline='') as file:
             reader = csv.DictReader(file)
@@ -79,9 +80,5 @@ class StyleDatabase:
             # and collections.NamedTuple has explicit documentation for accessing _fields. Same goes for _asdict()
             writer = csv.DictWriter(file, fieldnames=PromptStyle._fields)
             writer.writeheader()
-            writer.writerows(style._asdict() for k,     style in self.styles.items())
-
-        # Always keep a backup file around
-        if os.path.exists(path):
-            shutil.move(path, path + ".bak")
+            writer.writerows(style._asdict() for k, style in self.styles.items())
         shutil.move(temp_path, path)