add tunable ops

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/control/proc/dwpose/__init__.py

@@ -4,6 +4,7 @@
 # 3rd Edited by ControlNet
 # 4th Edited by ControlNet (added face and correct hands)
 
+from typing import Type, Optional, Union, List
 import os
 os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
 import cv2
@@ -16,37 +17,61 @@ checked_ok = False
 busy = False
 
 
+def _register_module(self, module: Type, module_name: Optional[Union[str, List[str]]] = None, force: bool = False) -> None:
+    if not callable(module):
+        raise TypeError(f'module must be Callable, but got {type(module)}')
+    if module_name is None:
+        module_name = module.__name__
+    if isinstance(module_name, str):
+        module_name = [module_name]
+    for name in module_name:
+        if not force and name in self._module_dict: # pylint: disable=protected-access
+            pass # patch for 'Adafactor is already registered in optimizer at torch.optim'
+        self._module_dict[name] = module # pylint: disable=protected-access
+
+
 def check_dependencies():
     global checked_ok, busy # pylint: disable=global-statement
+    busy = True
     debug = log.trace if os.environ.get('SD_DWPOSE_DEBUG', None) is not None else lambda *args, **kwargs: None
+    # pip install --upgrade --no-deps --force-reinstall termcolor xtcocotools terminaltables pycocotools munkres shapely openmim==0.3.9 mmengine==0.10.5 mmcv==2.2.0 mmpose==1.3.2 mmdet==3.3.0
     packages = [
         'termcolor',
+        'xtcocotools',
+        'terminaltables',
+        'pycocotools',
+        'munkres',
+        'shapely',
         'openmim==0.3.9',
-        'mmengine==0.10.4',
-        'mmcv==2.1.0',
-        'mmpose==1.3.1',
+        'mmengine==0.10.5',
+        'mmcv==2.2.0',
+        'mmpose==1.3.2',
         'mmdet==3.3.0',
     ]
-    status = [installed(p, reload=False, quiet=False) for p in packages]
+    status = [installed(p, reload=False, quiet=True) for p in packages]
     debug(f'DWPose required={packages} status={status}')
     if not all(status):
-        log.info(f'Installing DWPose dependencies: {[packages]}')
+        log.info(f'Installing DWPose dependencies: {packages}')
         cmd = 'install --upgrade --no-deps --force-reinstall '
         pkgs = ' '.join(packages)
-        res = pip(cmd + pkgs, ignore=False, quiet=False)
-        debug(f'DWPose pip install: {res}')
+        pip(cmd + pkgs, ignore=False, quiet=True, uv=False)
     try:
         import pkg_resources
         import imp # pylint: disable=deprecated-module
         imp.reload(pkg_resources)
         import mmcv # pylint: disable=unused-import
         import mmengine # pylint: disable=unused-import
+        from mmengine.registry import Registry
+        Registry._register_module = _register_module # pylint: disable=protected-access
         import mmpose # pylint: disable=unused-import
         import mmdet # pylint: disable=unused-import
         debug('DWPose import ok')
         checked_ok = True
     except Exception as e:
         log.error(f'DWPose: {e}')
+        # from modules import errors
+        # errors.display(e, 'DWPose')
+    busy = False
     return checked_ok
 
 

modules/devices.py

@@ -288,6 +288,20 @@ def set_cuda_memory_limit():
         log.warning(f'Torch CUDA memory limit: fraction={opts.cuda_mem_fraction:.2f} {e}')
 
 
+def set_cuda_tunable():
+    if not cuda_ok:
+        return
+    try:
+        if opts.torch_tunable_ops != 'default':
+            torch.cuda.tunable.enable(opts.torch_tunable_ops == 'true')
+            torch.cuda.tunable.tuning_enable(opts.torch_tunable_ops == 'true')
+            # torch.cuda.tunable.set_max_tuning_duration(100)
+            torch.cuda.tunable.set_max_tuning_iterations(opts.torch_tunable_limit)
+            # log.debug(f'Torce tunable: enabled={torch.cuda.tunable.is_enabled()} tuning={torch.cuda.tunable.tuning_is_enabled()} iterations={torch.cuda.tunable.get_max_tuning_iterations()} duration={torch.cuda.tunable.get_max_tuning_duration()}')
+    except Exception:
+        pass
+
+
 def test_fp16():
     global fp16_ok # pylint: disable=global-statement
     if fp16_ok is not None:
@@ -484,6 +498,7 @@ def set_dtype():
 def set_cuda_params():
     override_ipex_math()
     set_cuda_memory_limit()
+    set_cuda_tunable()
     set_cudnn_params()
     set_sdpa_params()
     set_dtype()
@@ -492,7 +507,12 @@ def set_cuda_params():
         device_name = get_raw_openvino_device()
     else:
         device_name = torch.device(get_optimal_device_name())
-    log.info(f'Torch parameters: backend={backend} device={device_name} config={opts.cuda_dtype} dtype={dtype} context={inference_context.__name__} nohalf={opts.no_half} nohalfvae={opts.no_half_vae} upcast={opts.upcast_sampling} deterministic={opts.cudnn_deterministic} fp16={"pass" if fp16_ok else "fail"} bf16={"pass" if bf16_ok else "fail"} optimization="{opts.cross_attention_optimization}"')
+    try:
+        # tunable = torch._C._jit_get_tunable_op_enabled() # pylint: disable=protected-access
+        tunable = [torch.cuda.tunable.is_enabled(), torch.cuda.tunable.tuning_is_enabled()]
+    except Exception:
+        tunable = [False, False]
+    log.info(f'Torch parameters: backend={backend} device={device_name} config={opts.cuda_dtype} dtype={dtype} context={inference_context.__name__} nohalf={opts.no_half} nohalfvae={opts.no_half_vae} upcast={opts.upcast_sampling} deterministic={opts.cudnn_deterministic} tunable={tunable} fp16={"pass" if fp16_ok else "fail"} bf16={"pass" if bf16_ok else "fail"} optimization="{opts.cross_attention_optimization}"')
 
 
 def cond_cast_unet(tensor):

modules/shared.py

@@ -528,6 +528,8 @@ options_templates.update(options_section(('backends', "Backend Settings"), {
     "cudnn_benchmark": OptionInfo(False, "Full-depth cuDNN benchmark"),
     "diffusers_fuse_projections": OptionInfo(False, "Fused projections"),
     "torch_expandable_segments": OptionInfo(False, "Expandable segments"),
+    "torch_tunable_ops": OptionInfo("default", "Tunable ops", gr.Radio, {"choices": ["default", "true", "false"]}),
+    "torch_tunable_limit": OptionInfo(30, "Tunable ops limit", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),
     "cuda_mem_fraction": OptionInfo(0.0, "Memory limit", gr.Slider, {"minimum": 0, "maximum": 2.0, "step": 0.05}),
     "torch_gc_threshold": OptionInfo(70, "GC threshold", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}),
     "inference_mode": OptionInfo("no-grad", "Inference mode", gr.Radio, {"choices": ["no-grad", "inference-mode", "none"]}),