add hunyuandit controlnets

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

CHANGELOG.md

@@ -9,6 +9,7 @@
   - [Nunchaku-Qwen-Image-Lightning](https://huggingface.co/nunchaku-tech/nunchaku-qwen-image)  
     if you have a compatible nVidia GPU, Nunchaku is the fastest quantization engine, currently available for Flux.1, SANA and Qwen-Image models  
     *note*: release version of `nunchaku==0.3.2` does NOT include support, so you need to build [nunchaku](https://nunchaku.tech/docs/nunchaku/installation/installation.html) from source  
+  - [HunyuanDiT ControlNet](https://huggingface.co/Tencent-Hunyuan/HYDiT-ControlNet-v1.2) Canny, Depth, Pose  
   - updated [SD.Next Model Samples Gallery](https://vladmandic.github.io/sd-samples/compare.html)  
 - **UI**  
   - improved image scaling in img2img and control interfaces  

modules/control/units/controlnet.py

@@ -105,6 +105,12 @@ predefined_sd3 = {
 predefined_qwen = {
     "InstantX Union Qwen": 'InstantX/Qwen-Image-ControlNet-Union',
 }
+predefined_hunyuandit = {
+    "HunyuanDiT Canny": 'Tencent-Hunyuan/HunyuanDiT-v1.2-ControlNet-Diffusers-Canny',
+    "HunyuanDiT Pose": 'Tencent-Hunyuan/HunyuanDiT-v1.2-ControlNet-Diffusers-Pose',
+    "HunyuanDiT Depth": 'Tencent-Hunyuan/HunyuanDiT-v1.2-ControlNet-Diffusers-Depth',
+}
+
 variants = {
     'NoobAI Canny XL': 'fp16',
     'NoobAI Lineart Anime XL': 'fp16',
@@ -120,6 +126,7 @@ all_models.update(predefined_sdxl)
 all_models.update(predefined_f1)
 all_models.update(predefined_sd3)
 all_models.update(predefined_qwen)
+all_models.update(predefined_hunyuandit)
 cache_dir = 'models/control/controlnet'
 load_lock = threading.Lock()
 
@@ -156,6 +163,8 @@ def api_list_models(model_type: str = None):
         model_list += list(predefined_sd3)
     if model_type == 'qwen' or model_type == 'all':
         model_list += list(predefined_qwen)
+    if model_type == 'hunyuandit' or model_type == 'all':
+        model_list += list(predefined_hunyuandit)
     model_list += sorted(find_models())
     return model_list
 
@@ -178,6 +187,8 @@ def list_models(refresh=False):
         models = ['None'] + list(predefined_sd3) + sorted(find_models())
     elif modules.shared.sd_model_type == 'qwen':
         models = ['None'] + list(predefined_qwen) + sorted(find_models())
+    elif modules.shared.sd_model_type == 'hunyuandit':
+        models = ['None'] + list(predefined_hunyuandit) + sorted(find_models())
     else:
         log.warning(f'Control {what} model list failed: unknown model type')
         models = ['None'] + sorted(predefined_sd15) + sorted(predefined_sdxl) + sorted(predefined_f1) + sorted(predefined_sd3) + sorted(find_models())
@@ -233,6 +244,9 @@ class ControlNet():
         elif shared.sd_model_type == 'qwen':
             from diffusers import QwenImageControlNetModel as cls
             config = 'InstantX/Qwen-Image-ControlNet-Union'
+        elif shared.sd_model_type == 'hunyuandit':
+            from diffusers import HunyuanDiT2DControlNetModel as cls
+            config = 'Tencent-Hunyuan/HunyuanDiT-v1.2-ControlNet-Diffusers-Canny'
         else:
             log.error(f'Control {what}: type={shared.sd_model_type} unsupported model')
             return None, None
@@ -444,6 +458,20 @@ class ControlNetPipeline():
                 scheduler=pipeline.scheduler,
                 controlnet=controlnets[0] if isinstance(controlnets, list) else controlnets, # can be a list
             )
+        elif detect.is_hunyuandit(pipeline) and len(controlnets) > 0:
+            from diffusers import HunyuanDiTControlNetPipeline
+            self.pipeline = HunyuanDiTControlNetPipeline(
+                vae=pipeline.vae,
+                text_encoder=pipeline.text_encoder,
+                tokenizer=pipeline.tokenizer,
+                text_encoder_2=pipeline.text_encoder_2,
+                tokenizer_2=pipeline.tokenizer_2,
+                transformer=pipeline.transformer,
+                scheduler=pipeline.scheduler,
+                safety_checker=None,
+                feature_extractor=None,
+                controlnet=controlnets[0] if isinstance(controlnets, list) else controlnets, # can be a list
+            )
         elif len(loras) > 0:
             self.pipeline = pipeline
             for lora in loras:
@@ -476,7 +504,7 @@ class ControlNetPipeline():
         debug_log(f'Control {what} pipeline: class={self.pipeline.__class__.__name__} time={t1-t0:.2f}')
 
     def restore(self):
-        if self.pipeline is not None:
+        if self.pipeline is not None and hasattr(self.pipeline, 'unload_lora_weights'):
             self.pipeline.unload_lora_weights()
         self.pipeline = None
         return self.orig_pipeline

modules/control/units/detect.py

@@ -24,3 +24,7 @@ def is_sd3(model):
 
 def is_qwen(model):
     return is_compatible(model, pattern='Qwen')
+
+
+def is_hunyuandit(model):
+    return is_compatible(model, pattern='HunyuanDiT')

modules/processing_args.py

@@ -358,8 +358,8 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
     task_kwargs = task_specific_kwargs(p, model)
     pipe_args = getattr(p, 'task_args', {})
     model_args = getattr(model, 'task_args', {})
-    task_kwargs.update(pipe_args)
-    task_kwargs.update(model_args)
+    task_kwargs.update(pipe_args or {})
+    task_kwargs.update(model_args or {})
     if debug_enabled:
         debug_log(f'Process task args: {task_kwargs}')
     for k, v in task_kwargs.items():

modules/ui_extra_networks_checkpoints.py

@@ -43,6 +43,7 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
                 "info": {},
                 "metadata": {},
                 "description": v.get('desc', ''),
+                "version": "ref",
             }
 
     def create_item(self, name):

pipelines/model_hunyuandit.py

@@ -1,6 +1,6 @@
 import transformers
 import diffusers
-from modules import shared, sd_models, devices, model_quant, sd_hijack_te
+from modules import shared, sd_models, devices, model_quant
 from pipelines import generic
 
 
@@ -8,6 +8,11 @@ def load_hunyuandit(checkpoint_info, diffusers_load_config={}):
     repo_id = sd_models.path_to_repo(checkpoint_info)
     sd_models.hf_auth_check(checkpoint_info)
 
+    # import torch # override for hunyuandit
+    # devices.dtype = torch.float16
+    # devices.dtype_vae = torch.float16
+    # devices.dtype_unet = torch.float16
+    # diffusers_load_config['torch_dtype'] = devices.dtype
     load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config)
     shared.log.debug(f'Load model: type=HunyuanDiT repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
 
@@ -19,13 +24,15 @@ def load_hunyuandit(checkpoint_info, diffusers_load_config={}):
         repo_id,
         transformer=transformer,
         text_encoder_2=text_encoder_2,
+        safety_checker=None,
+        feature_extractor=None,
         cache_dir=shared.opts.diffusers_dir,
         **load_args,
     )
 
     del text_encoder_2
     del transformer
-    sd_hijack_te.init_hijack(pipe)
+    # sd_hijack_te.init_hijack(pipe)
 
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_sana.py

@@ -1,4 +1,3 @@
-import time
 import torch
 import diffusers
 import transformers
@@ -47,7 +46,6 @@ def load_sana(checkpoint_info, kwargs={}):
 
     kwargs = load_quants(kwargs, repo_id, cache_dir=shared.opts.diffusers_dir)
     shared.log.debug(f'Load model: type=Sana repo="{repo_id}" args={list(kwargs)}')
-    t0 = time.time()
 
     if devices.dtype == torch.bfloat16 or devices.dtype == torch.float32:
         kwargs['torch_dtype'] = devices.dtype