add shared.opts.device_map

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

CHANGELOG.md

@@ -1,14 +1,6 @@
 # Change Log for SD.Next
 
-## TODO for 2025-04-10
-
-- HiDream requires sidebranch:  
-  > `pip install git+https://github.com/hlky/diffusers@hidream`  
-  > `./webui.sh --debug --experimental`  
-- HiDream make `LLama` optional  
-- `FlashAttn` wiki  
-
-## Update for 2025-04-10
+## Update for 2025-04-11
 
 - **Models**  
   - [HiDream-I1](https://huggingface.co/HiDream-ai/HiDream-I1-Full) in fast, dev and full variants!  
@@ -17,9 +9,8 @@
     simply select from *networks -> models -> reference*  
     due to size (over 25B params in 58GB), offloading and on-the-fly quantization are pretty much a necessity  
     difference between variants is recommended number of steps: *fast=16, dev=28, full=50*  
-    hidream is compatible with flowmatching samplers and with taesd live-preview  
-    *note* HiDream-I1 requires `flash-attn` to be installed  
-    > REQUIRES SIDE-BRANCH
+    hidream supportability: *offloading, quantization, taesd live-preview, remote-vae*  
+    hidream compatibility: *flowmatching samplers*  
 - **Features**  
   - Custom model loader  
       can be used to load any known diffusion model with default or custom model components  
@@ -47,6 +38,10 @@
   - diag: add get-server-status to UI generate context menu  
   - diag: memory monitor detect gpu swapping  
   - use [hf-xet](https://huggingface.co/blog/xet-on-the-hub) for huggingface downloads where possible  
+  - quant: update & fix `optimum-quanto` for transformers  
+  - quant: update & fix `torchao`  
+  - model load: new setting for model load initial device map  
+    can be used to force gpu vs cpu when loading model to avoid oom before model offloading is even activated after load  
 - **Changes**  
   - params: Reset default guidance-rescale from 0.7 to 0.0  
   - progress: add additional fields to progress API  

installer.py

@@ -538,7 +538,7 @@ def check_diffusers():
     t_start = time.time()
     if args.skip_all or args.skip_git or args.experimental:
         return
-    sha = 'd1387ecee5262e75386ce8948ddcf9a4de0ebbfa' # diffusers commit hash
+    sha = '0ef29355c9d65b78eabb6a4ac5bee73aa685e9a6' # diffusers commit hash
     pkg = pkg_resources.working_set.by_key.get('diffusers', None)
     minor = int(pkg.version.split('.')[1] if pkg is not None else 0)
     cur = opts.get('diffusers_version', '') if minor > 0 else ''

modules/model_hidream.py

@@ -1,7 +1,7 @@
 import time
 import transformers
 import diffusers
-from modules import shared, devices, sd_models, timer
+from modules import shared, devices, sd_models, timer, model_quant, modelloader
 
 
 llama_repo = "meta-llama/Meta-Llama-3.1-8B-Instruct"
@@ -17,31 +17,34 @@ def hijack_encode_prompt(*args, **kwargs):
     return res
 
 
-def get_args(diffusers_load_config={}, module=None):
-    from modules import model_quant, modelloader
+def get_args(load_config:dict={}, module:str=None, device_map:bool=False):
+    config = load_config.copy()
     modelloader.hf_login()
-    if 'torch_dtype' not in diffusers_load_config:
-        diffusers_load_config['torch_dtype'] = devices.dtype
-    if 'low_cpu_mem_usage' in diffusers_load_config:
-        del diffusers_load_config['low_cpu_mem_usage']
-    if 'load_connected_pipeline' in diffusers_load_config:
-        del diffusers_load_config['load_connected_pipeline']
-    if 'safety_checker' in diffusers_load_config:
-        del diffusers_load_config['safety_checker']
-    if 'requires_safety_checker' in diffusers_load_config:
-        del diffusers_load_config['requires_safety_checker']
+    if 'torch_dtype' not in config:
+        config['torch_dtype'] = devices.dtype
+    if 'low_cpu_mem_usage' in config:
+        del config['low_cpu_mem_usage']
+    if 'load_connected_pipeline' in config:
+        del config['load_connected_pipeline']
+    if 'safety_checker' in config:
+        del config['safety_checker']
+    if 'requires_safety_checker' in config:
+        del config['requires_safety_checker']
+    if device_map:
+        if shared.opts.device_map == 'cpu':
+            config['device_map'] = 'cpu'
+        if shared.opts.device_map == 'gpu':
+            config['device_map'] = devices.device
     quant_args = model_quant.create_config(module=module)
-    quant_type = model_quant.get_quant_type(quant_args)
-    if quant_type:
-        shared.log.debug(f'Load model: type=HiDream quantization module="{module}" {quant_type}')
-    return diffusers_load_config, quant_args
+    return config, quant_args
 
 
 def load_hidream(checkpoint_info, diffusers_load_config={}):
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
     shared.log.debug(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}')
 
-    load_args, quant_args = get_args(diffusers_load_config, module='Transformer')
+    load_args, quant_args = get_args(diffusers_load_config, module='Transformer', device_map=True)
+    shared.log.debug(f'Load model: type=HiDream transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
     transformer = diffusers.HiDreamImageTransformer2DModel.from_pretrained(
         repo_id,
         subfolder="transformer",
@@ -50,7 +53,8 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
         **quant_args,
     )
 
-    load_args, quant_args = get_args(diffusers_load_config, module='TE')
+    load_args, quant_args = get_args(diffusers_load_config, module='TE', device_map=True)
+    shared.log.debug(f'Load model: type=HiDream te3="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
     text_encoder_3 = transformers.T5EncoderModel.from_pretrained(
         repo_id,
         subfolder="text_encoder_3",
@@ -59,7 +63,8 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
         **quant_args,
     )
 
-    load_args, quant_args = get_args(diffusers_load_config, module='LLM')
+    load_args, quant_args = get_args(diffusers_load_config, module='LLM', device_map=True)
+    shared.log.debug(f'Load model: type=HiDream te4="{llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
     tokenizer_4 = transformers.PreTrainedTokenizerFast.from_pretrained(
         llama_repo,
         cache_dir=shared.opts.hfcache_dir,
@@ -73,7 +78,6 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
         **load_args,
         **quant_args,
     )
-
     load_args, quant_args = get_args(diffusers_load_config, module='Model')
     pipe = diffusers.HiDreamImagePipeline.from_pretrained(
         repo_id,

modules/processing_helpers.py

@@ -371,7 +371,7 @@ def validate_sample(tensor):
         shared.log.error(f'Decode: sample={sample.shape} invalid={nans} dtype={dtype} vae={vae} upcast={upcast} failed to validate')
         if upcast is not None and not upcast:
             setattr(shared.sd_model.vae.config, 'force_upcast', True) # noqa: B010
-            shared.log.warning('Decode: upcast=True set, retry operation')
+            shared.log.info('Decode: set upcast=True and attempt to retry operation')
     t1 = time.time()
     timer.process.add('validate', t1 - t0)
     return cast

modules/sd_vae_remote.py

@@ -12,6 +12,7 @@ hf_decode_endpoints = {
     'sd': 'https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud',
     'sdxl': 'https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud',
     'f1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud',
+    'h1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud',
     'hunyuanvideo': 'https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud',
 }
 hf_encode_endpoints = {
@@ -27,6 +28,13 @@ dtypes = {
 }
 
 
+def h1_pack_latents(latents, _batch_size, _num_channels_latents, _height, _width): # TODO hidream: pack latents for remote vae
+    # latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+    # latents = latents.permute(0, 2, 4, 1, 3, 5)
+    # latents = latents.reshape(batch_size, (height // 2) * (width // 2) // (num_channels_latents * 4), num_channels_latents * 4)
+    return latents
+
+
 def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_type: str = None) -> Image.Image:
     from modules import devices, shared, errors, modelloader
     tensors = []
@@ -44,10 +52,14 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
         latent_copy = latent_copy.unsqueeze(0)
 
     for i in range(latent_copy.shape[0]):
+        params = {}
         try:
             latent = latent_copy[i]
-            if model_type != 'f1':
+            if model_type != 'f1' and model_type != 'h1':
                 latent = latent.unsqueeze(0)
+            if model_type == 'h1':
+                num_channels_latents = shared.sd_model.transformer.config.in_channels
+                latent = h1_pack_latents(latent, 1, num_channels_latents, height, width) # pylint: disable=protected-access
             params = {
                 "input_tensor_type": "binary",
                 "shape": list(latent.shape),
@@ -72,7 +84,7 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
                 params["output_type"] = "pt"
                 params["output_tensor_type"] = "binary"
                 headers["Accept"] = "tensor/binary"
-            if (model_type == 'f1') and (width > 0) and (height > 0):
+            if (model_type == 'f1' or model_type == 'h1') and (width > 0) and (height > 0):
                 params['width'] = width
                 params['height'] = height
             if shared.sd_model.vae is not None and shared.sd_model.vae.config is not None:

modules/shared.py

@@ -408,6 +408,7 @@ options_templates.update(options_section(('sd', "Models & Loading"), {
     "stream_load": OptionInfo(False, "Model load using streams", gr.Checkbox),
     "diffusers_eval": OptionInfo(True, "Force model eval", gr.Checkbox, {"visible": False }),
     "diffusers_to_gpu": OptionInfo(False, "Model load model direct to GPU"),
+    "device_map": OptionInfo('default', "Model load device map", gr.Radio, {"choices": ['default', 'gpu', 'cpu'] }),
     "disable_accelerate": OptionInfo(False, "Disable accelerate", gr.Checkbox, {"visible": False }),
     "sd_model_dict": OptionInfo('None', "Use separate base dict", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_titles(), "visible": False}, refresh=refresh_checkpoints),
     "sd_checkpoint_cache": OptionInfo(0, "Cached models", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1, "visible": not native }),