From 0439e5652d013133f799685e76e20a41a9a4c55d Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 11 Apr 2025 14:53:41 -0400 Subject: [PATCH] add shared.opts.device_map Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 19 ++++++--------- installer.py | 2 +- modules/model_hidream.py | 46 +++++++++++++++++++---------------- modules/processing_helpers.py | 2 +- modules/sd_vae_remote.py | 16 ++++++++++-- modules/shared.py | 1 + 6 files changed, 49 insertions(+), 37 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ce4da4cf7..48615ddd4 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/installer.py b/installer.py index 26ba40229..1750ba745 100644 --- a/installer.py +++ b/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 '' diff --git a/modules/model_hidream.py b/modules/model_hidream.py index 075c4bce9..d58d87d39 100644 --- a/modules/model_hidream.py +++ b/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, diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 8fffa8313..71e31f663 100644 --- a/modules/processing_helpers.py +++ b/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 diff --git a/modules/sd_vae_remote.py b/modules/sd_vae_remote.py index 3c1846ea8..35c9c6d76 100644 --- a/modules/sd_vae_remote.py +++ b/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: diff --git a/modules/shared.py b/modules/shared.py index a7bb10e13..04a5e822f 100644 --- a/modules/shared.py +++ b/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 }),