Merge branch 'master' into directml

TODO.md

@@ -18,6 +18,7 @@ Stuff to be added...
 - Redo Extensions tab: see <https://vladmandic.github.io/sd-extension-manager/pages/extensions.html>
 - Stream-load models as option for slow storage
 - Autodetect nVidia and AMD: `nvidia-smi` vs `rocm-smi`
+- [Temporal Weighing](https://github.com/comfyanonymous/ComfyUI/discussions/473)
 
 ## Investigate
 
@@ -56,9 +57,4 @@ Tech that can be integrated as part of the core workflow...
 
 ### Pending Code Updates
 
-- use samples format for live preview
-- identify race condition where generate locks up while fetching preview
-- use **Approx NN** for live preview
-- create default `styles.csv`
-- fix setup not installing `tensorflow` dependencies
-- update default git flags to reduce number of warnings
+- update GPU utility search paths for better GPU type detection

javascript/progressbar.js

@@ -106,7 +106,7 @@ function requestProgress(id_task, progressbarContainer, gallery, atEnd = null, o
               removeProgressBar()
               return
           }
-          if (elapsedFromStart > 10 && !res.queued && res.progress == prevProgress) {
+          if (elapsedFromStart > 30 && !res.queued && res.progress == prevProgress) {
                 removeProgressBar()
                 return
             }

modules/cmd_args.py

@@ -44,8 +44,6 @@ 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)
 
 
 def compatibility_args(opts, args):
@@ -67,6 +65,9 @@ def compatibility_args(opts, args):
     parser.add_argument("--opt-channelslast", help=argparse.SUPPRESS, default=opts.opt_channelslast)
     parser.add_argument("--xformers", default = (opts.cross_attention_optimization == "xFormers"), action='store_true', help=argparse.SUPPRESS)
     parser.add_argument("--disable-nan-check", help=argparse.SUPPRESS, default=opts.disable_nan_check)
+    parser.add_argument("--token-merging", help=argparse.SUPPRESS, default=opts.token_merging)
+    parser.add_argument("--rollback-vae", help=argparse.SUPPRESS, default=opts.rollback_vae)
+
     parser.add_argument("--no-half", help=argparse.SUPPRESS, default=opts.no_half)
     parser.add_argument("--no-half-vae", help=argparse.SUPPRESS, default=opts.no_half_vae)
     parser.add_argument("--precision", help=argparse.SUPPRESS, default=opts.precision)

modules/sd_models.py

@@ -10,13 +10,12 @@ from rich import print, progress # pylint: disable=redefined-builtin
 import torch
 import safetensors.torch
 from omegaconf import OmegaConf
+import tomesd
 import ldm.modules.midas as midas
 from ldm.util import instantiate_from_config
-
 from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config
 from modules.sd_hijack_inpainting import do_inpainting_hijack
 from modules.timer import Timer
-import tomesd
 
 
 model_dir = "Stable-diffusion"
@@ -231,7 +230,7 @@ def read_metadata_from_safetensors(filename):
         return res
 
 
-def read_state_dict(checkpoint_file):
+def read_state_dict(checkpoint_file, map_location=None): # pylint: disable=unused-argument
     try:
         with progress.open(checkpoint_file, 'rb', description=f'Loading weights: [cyan]{checkpoint_file}', auto_refresh=True) as f:
             _, extension = os.path.splitext(checkpoint_file)

modules/sd_samplers.py

@@ -31,11 +31,11 @@ def create_sampler(name, model):
 def set_samplers():
     global samplers, samplers_for_img2img
 
-    hidden = set(shared.opts.hide_samplers)
-    hidden_img2img = set(shared.opts.hide_samplers + ['PLMS', 'UniPC'])
+    shown_img2img = set(shared.opts.show_samplers)
+    shown = set(shared.opts.show_samplers + ['PLMS', 'UniPC'])
 
-    samplers = [x for x in all_samplers if x.name not in hidden]
-    samplers_for_img2img = [x for x in all_samplers if x.name not in hidden_img2img]
+    samplers = [x for x in all_samplers if x.name in shown]
+    samplers_for_img2img = [x for x in all_samplers if x.name in shown_img2img]
 
     samplers_map.clear()
     for sampler in all_samplers:

modules/sd_vae.py

@@ -3,6 +3,7 @@ import collections
 import glob
 from copy import deepcopy
 from rich import print # pylint: disable=redefined-builtin
+import torch
 from modules import paths, shared, devices, script_callbacks, sd_models
 
 

modules/shared.py

@@ -259,7 +259,6 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "sub_quad_kv_chunk_size": OptionInfo(512, "Sub-quadratic cross-attentionkv chunk size for the sub-quadratic cross-attention layer optimization to use", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8}),
     "sub_quad_chunk_threshold": OptionInfo(80, "Sub-quadratic cross-attention percentage of VRAM chunking threshold", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}),
     "always_batch_cond_uncond": OptionInfo(False, "Disables cond/uncond batching that is enabled to save memory with --medvram or --lowvram"),
-    "upcast_sampling": OptionInfo(False, "Enable upcast sampling. Usually produces similar results to --no-half with better performance while using less memory"),
 }))
 
 options_templates.update(options_section(('system-paths', "System Paths"), {
@@ -332,7 +331,9 @@ options_templates.update(options_section(('cuda', "CUDA Settings"), {
     "cuda_dtype": OptionInfo("FP16", "Device precision type", gr.Radio, lambda: {"choices": ["FP32", "FP16", "BF16"]}),
     "no_half": OptionInfo(True if is_device_dml else False, "Use full precision for model (--no-half)"),
     "no_half_vae": OptionInfo(True if is_device_dml else False, "Use full precision for VAE (--no-half-vae)"),
+    "upcast_sampling": OptionInfo(False, "Enable upcast sampling. Usually produces similar results to --no-half with better performance while using less memory"),
     "disable_nan_check": OptionInfo(True, "Do not check if produced images/latent spaces have NaN values"),
+    "rollback_vae": OptionInfo(False, "Attempt to roll back VAE when produced NaN values, requires NaN check (experimental)"),
     "opt_channelslast": OptionInfo(False, "Use channels last as torch memory format "),
     "cudnn_benchmark": OptionInfo(False, "Enable cuDNN benchmark feature"),
     "cuda_allow_tf32": OptionInfo(True, "Allow TF32 math ops"),
@@ -434,14 +435,14 @@ options_templates.update(options_section(('ui', "Live previews"), {
 }))
 
 options_templates.update(options_section(('sampler-params', "Sampler parameters"), {
-    "hide_samplers": OptionInfo(["Euler", "LMS", "Heun", "DPM2", "DPM2 a", "DPM++ 2M", "DPM fast", "DPM adaptive", "DPM++ 2S a Karras", "DPM++ 2S a", "DPM++ SDE Karras", "DPM2 a Karras", "LMS Karras"], "Hide samplers in user interface", gr.CheckboxGroup, lambda: {"choices": [x.name for x in list_samplers()]}),
-    "eta_ddim": OptionInfo(0.0, "eta (noise multiplier) for DDIM", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
-    "eta_ancestral": OptionInfo(1.0, "eta (noise multiplier) for ancestral samplers", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
-    "ddim_discretize": OptionInfo('uniform', "img2img DDIM discretize", gr.Radio, {"choices": ['uniform', 'quad']}),
+    "show_samplers": OptionInfo(["Euler a", "UniPC", "DDIM", "DPM++ SDE", "DPM++ SDE", "DPM2 Karras", "DPM++ 2M Karras"], "Show samplers in user interface", gr.CheckboxGroup, lambda: {"choices": [x.name for x in list_samplers()]}),
+    "eta_ancestral": OptionInfo(1.0, "Noise multiplier for ancestral samplers (eta)", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+    "eta_ddim": OptionInfo(0.0, "Noise multiplier for DDIM (eta)", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+    "ddim_discretize": OptionInfo('uniform', "DDIM discretize img2img", gr.Radio, {"choices": ['uniform', 'quad']}),
     's_churn': OptionInfo(0.0, "sigma churn", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     's_tmin':  OptionInfo(0.0, "sigma tmin",  gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     's_noise': OptionInfo(1.0, "sigma noise", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
-    'eta_noise_seed_delta': OptionInfo(0, "Eta noise seed delta", gr.Number, {"precision": 0}),
+    'eta_noise_seed_delta': OptionInfo(0, "Noise seed delta (eta)", gr.Number, {"precision": 0}),
     'always_discard_next_to_last_sigma': OptionInfo(False, "Always discard next-to-last sigma"),
     'uni_pc_variant': OptionInfo("bh1", "UniPC variant", gr.Radio, {"choices": ["bh1", "bh2", "vary_coeff"]}),
     'uni_pc_skip_type': OptionInfo("time_uniform", "UniPC skip type", gr.Radio, {"choices": ["time_uniform", "time_quadratic", "logSNR"]}),
@@ -449,6 +450,20 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
     'uni_pc_lower_order_final': OptionInfo(True, "UniPC lower order final"),
 }))
 
+options_templates.update(options_section(('token_merging', 'Token Merging'), {
+    "token_merging": OptionInfo(False, "Enable redundant token merging via tomesd. This can provide significant speed and memory improvements.", gr.Checkbox),
+    "token_merging_ratio": OptionInfo(0.5, "Merging Ratio", gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}),
+    "token_merging_hr_only": OptionInfo(True, "Apply only to high-res fix pass. Disabling can yield a ~20-35% speedup on contemporary resolutions.", gr.Checkbox),
+    "token_merging_ratio_hr": OptionInfo(0.5, "Merging Ratio (high-res pass) - If 'Apply only to high-res' is enabled, this will always be the ratio used.", gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}),
+    "token_merging_random": OptionInfo(False, "Use random perturbations - Can improve outputs for certain samplers. For others, it may cause visual artifacting.", gr.Checkbox),
+    "token_merging_merge_attention": OptionInfo(True, "Merge attention", gr.Checkbox),
+    "token_merging_merge_cross_attention": OptionInfo(False, "Merge cross attention", gr.Checkbox),
+    "token_merging_merge_mlp": OptionInfo(False, "Merge mlp", gr.Checkbox),
+    "token_merging_maximum_down_sampling": OptionInfo(1, "Maximum down sampling", gr.Dropdown, lambda: {"choices": ["1", "2", "4", "8"]}),
+    "token_merging_stride_x": OptionInfo(2, "Stride - X", gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}),
+    "token_merging_stride_y": OptionInfo(2, "Stride - Y", gr.Slider, {"minimum": 2, "maximum": 8, "step": 2})
+}))
+
 options_templates.update(options_section(('postprocessing', "Postprocessing"), {
     'postprocessing_enable_in_main_ui': OptionInfo([], "Enable postprocessing operations in txt2img and img2img tabs", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts()]}),
     'postprocessing_operation_order': OptionInfo([], "Postprocessing operation order", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts()]}),
@@ -461,54 +476,6 @@ options_templates.update(options_section((None, "Hidden options"), {
     "sd_checkpoint_hash": OptionInfo("", "SHA256 hash of the current checkpoint"),
 }))
 
-options_templates.update(options_section(('token_merging', 'Token Merging'), {
-    "token_merging": OptionInfo(
-        False, "Enable redundant token merging via tomesd. This can provide significant speed and memory improvements.",
-        gr.Checkbox
-    ),
-    "token_merging_ratio": OptionInfo(
-        0.5, "Merging Ratio",
-        gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}
-    ),
-    "token_merging_hr_only": OptionInfo(
-        True, "Apply only to high-res fix pass. Disabling can yield a ~20-35% speedup on contemporary resolutions.",
-        gr.Checkbox
-    ),
-    "token_merging_ratio_hr": OptionInfo(
-        0.5, "Merging Ratio (high-res pass) - If 'Apply only to high-res' is enabled, this will always be the ratio used.",
-        gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}
-    ),
-    # More advanced/niche settings:
-    "token_merging_random": OptionInfo(
-        False, "Use random perturbations - Can improve outputs for certain samplers. For others, it may cause visual artifacting.",
-        gr.Checkbox
-    ),
-    "token_merging_merge_attention": OptionInfo(
-        True, "Merge attention",
-        gr.Checkbox
-    ),
-     "token_merging_merge_cross_attention": OptionInfo(
-        False, "Merge cross attention",
-        gr.Checkbox
-    ),
-    "token_merging_merge_mlp": OptionInfo(
-        False, "Merge mlp",
-        gr.Checkbox
-    ),
-    "token_merging_maximum_down_sampling": OptionInfo(
-        1, "Maximum down sampling",
-        gr.Dropdown, lambda: {"choices": ["1", "2", "4", "8"]}
-    ),
-    "token_merging_stride_x": OptionInfo(
-        2, "Stride - X",
-        gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}
-    ),
-    "token_merging_stride_y": OptionInfo(
-        2, "Stride - Y",
-        gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}
-    )
-}))
-
 options_templates.update()
 
 

modules/ui.py

@@ -1514,10 +1514,6 @@ def create_ui():
                 with gr.TabItem(label, id=ifid, elem_id='tab_' + ifid):
                     interface.render()
 
-        footer = shared.html("footer.html")
-        footer = footer.format(versions=versions_html())
-        gr.HTML(footer, elem_id="footer")
-
         text_settings = gr.Textbox(elem_id="settings_json", value=lambda: opts.dumpjson(), visible=False)
         settings_submit.click(
             fn=wrap_gradio_call(run_settings, extra_outputs=[gr.update()]),

setup.py

@@ -141,7 +141,7 @@ def update(folder):
         git('checkout master', folder)
     else:
         log.warning(f'Unknown branch for: {folder}')
-    git('pull --autostash', folder)
+    git('pull --autostash --rebase', folder)
     branch = git('branch', folder)
 
 
@@ -179,11 +179,11 @@ def check_python():
 
 # check torch version
 def check_torch():
-    if shutil.which('nvidia-smi') is not None:
+    if shutil.which('nvidia-smi') is not None or os.path.exists(os.path.join(os.environ.get('SystemRoot') or r'C:\Windows', 'System32', 'nvidia-smi.exe')):
         log.info('nVidia toolkit detected')
         torch_command = os.environ.get('TORCH_COMMAND', 'torch torchaudio torchvision --index-url https://download.pytorch.org/whl/cu118')
         xformers_package = os.environ.get('XFORMERS_PACKAGE', 'xformers==0.0.17' if opts.get('cross_attention_optimization', '') == 'xFormers' else 'none')
-    elif shutil.which('rocminfo') is not None:
+    elif shutil.which('rocminfo') is not None or os.path.exists('/opt/rocm/bin/rocminfo'):
         log.info('AMD toolkit detected')
         os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', '10.3.0')
         torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.4.2')
@@ -232,7 +232,7 @@ def check_torch():
         log.debug(f'Cannot install xformers package: {e}')
     try:
         tensorflow_package = os.environ.get('TENSORFLOW_PACKAGE', 'tensorflow==2.12.0')
-        install(tensorflow_package, ignore=True)
+        install(tensorflow_package, 'tensorflow', ignore=True)
     except Exception as e:
         log.debug(f'Cannot install tensorflow package: {e}')
 
@@ -411,10 +411,10 @@ def check_version():
     if not os.path.exists('.git'):
         log.error('Not a git repository')
         exit(1)
-    status = git('status')
-    if 'branch' not in status:
-        log.error('Cannot get git repository status')
-        exit(1)
+    _status = git('status')
+    # if 'branch' not in status:
+    #    log.error('Cannot get git repository status')
+    #    exit(1)
     ver = git('log -1 --pretty=format:"%h %ad"')
     log.info(f'Version: {ver}')
     commit = git('rev-parse HEAD')

webui.py

@@ -66,11 +66,11 @@ else:
 
 def check_rollback_vae():
     if shared.cmd_opts.rollback_vae:
-        if version.parse(torch.__version__) < version.parse('2.1'):
-            print("If your PyTorch version is lower than PyTorch 2.1, Rollback VAE will not work.")
+        if not torch.__version__.startswith('2.1'):
+            print("Rollback VAE functionality requires Torch 2.1 or higher")
             shared.cmd_opts.rollback_vae = False
-        elif 0 < torch.cuda.get_device_capability()[0] < 8:
-            print('Rollback VAE will not work because your device does not support it.')
+        if 0 < torch.cuda.get_device_capability()[0] < 8:
+            print('Rollback VAE functionality device capabilities not met')
             shared.cmd_opts.rollback_vae = False