hidream add llm info to metadata

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

CHANGELOG.md

@@ -5,7 +5,7 @@
 - **Features**
   - [Nunchaku](https://github.com/mit-han-lab/nunchaku) inference engine with custom **SVDQuant** 4-bit execution  
     highly experimental and with limited support, but when it works, its magic: **Flux.1 at 5.90 it/s** *(not sec/it)*!  
-    see [Nunchaku Wiki](https://github.com/vladmandic/sdnext/wiki/Nunchaku) for details  
+    see [Nunchaku Wiki](https://github.com/vladmandic/sdnext/wiki/Nunchaku) for installation guide and list of supported models & features  
   - [CFG-Zero](https://github.com/WeichenFan/CFG-Zero-star) new guidance method optimized for flow-matching models  
     implemented for **FLUX.1, HiDream-I1, SD3.x, CogView4, HunyuanVideo, WanAI**  
     enable and configure in *settings -> pipeline modifiers -> cfg zero*  
@@ -14,10 +14,12 @@
   - **HiDream** optimized offloading and prompt-encode caching  
     it now works in 12GB VRAM / 26GB RAM!  
   - **CogView3** and **CogView4** model loader optimizations  
+  - **Sana** model loader optimizations
   - add explicit offload after encode prompt  
     configure in *settings -> text encoder -> offload*  
 - **Other**
   - **HiDream** add HF gated access auth check  
+  - **HiDream** add LLM into to metadata  
   - add **UniPC FlowMatch** scheduler  
   - add **LCM FlowMatch** scheduler  
   - networks: set which networks to skip when scanning civitai  

modules/model_hidream.py

@@ -54,7 +54,8 @@ def load_text_encoders(repo_id, diffusers_load_config={}):
         sd_models.move_model(text_encoder_3, devices.cpu)
 
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='LLM', device_map=True)
-    shared.log.debug(f'Load model: type=HiDream te4="{shared.opts.model_h1_llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
+    llama_repo = shared.opts.model_h1_llama_repo if shared.opts.model_h1_llama_repo != 'Default' else 'meta-llama/Meta-Llama-3.1-8B-Instruct'
+    shared.log.debug(f'Load model: type=HiDream te4="{llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
 
     text_encoder_4 = transformers.LlamaForCausalLM.from_pretrained(
         shared.opts.model_h1_llama_repo,

modules/model_quant.py

@@ -494,8 +494,8 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al
         del config['safety_checker']
     if 'requires_safety_checker' in config:
         del config['requires_safety_checker']
-    if 'variant' in config:
-        del config['variant']
+    # if 'variant' in config:
+    #     del config['variant']
     if device_map:
         if shared.opts.device_map == 'cpu':
             config['device_map'] = 'cpu'

modules/model_sana.py

@@ -6,15 +6,21 @@ from modules import shared, sd_models, devices, modelloader, model_quant
 
 
 def load_quants(kwargs, repo_id, cache_dir):
-    quant_args = {}
-    quant_args = model_quant.create_config()
-    if not quant_args:
-        return kwargs
-    load_args = kwargs.copy()
-    if 'transformer' not in kwargs and (('Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization) or ('Transformer' in shared.opts.bnb_quantization or 'Transformer' in shared.opts.torchao_quantization or 'Transformer' in shared.opts.quanto_quantization)):
-        kwargs['transformer'] = diffusers.models.SanaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, **load_args, **quant_args)
-    if 'text_encoder' not in kwargs and ('TE' in shared.opts.bnb_quantization or 'TE' in shared.opts.torchao_quantization or 'TE' in shared.opts.quanto_quantization):
-        kwargs['text_encoder'] = transformers.AutoModelForCausalLM.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=cache_dir, **load_args, **quant_args)
+    kwargs_copy = kwargs.copy()
+    if model_quant.check_nunchaku('Transformer') and 'Sana_1600M' in repo_id: # only sana-1600m
+        import nunchaku
+        nunchaku_precision = nunchaku.utils.get_precision()
+        nunchaku_repo = f"mit-han-lab/svdq-{nunchaku_precision}-sana-1600m"
+        shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} attention={shared.opts.nunchaku_attention}')
+        kwargs['transformer'] = nunchaku.NunchakuSanaTransformer2DModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype)
+    elif model_quant.check_quant('Transformer'):
+        load_args, quant_args = model_quant.get_dit_args(kwargs_copy, module='Transformer')
+        if quant_args:
+            kwargs['transformer'] = diffusers.SanaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, **load_args, **quant_args)
+    if model_quant.check_quant('TE'):
+        load_args, quant_args = model_quant.get_dit_args(kwargs_copy, module='TE')
+        if quant_args:
+            kwargs['text_encoder'] = transformers.AutoModelForCausalLM.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=cache_dir, **load_args, **quant_args)
     return kwargs
 
 
@@ -28,9 +34,9 @@ def load_sana(checkpoint_info, kwargs={}):
     kwargs.pop('requires_safety_checker', None)
     kwargs.pop('torch_dtype', None)
 
+    # set variant since hf repos are a mess
     if not repo_id.endswith('_diffusers'):
         repo_id = f'{repo_id}_diffusers'
-
     if 'Sana_1600M' in repo_id:
         if devices.dtype == torch.bfloat16 or 'BF16' in repo_id:
             if 'BF16' not in repo_id:
@@ -45,6 +51,7 @@ 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
     if 'Sprint' in repo_id:
@@ -56,21 +63,31 @@ def load_sana(checkpoint_info, kwargs={}):
         cache_dir=shared.opts.diffusers_dir,
         **kwargs,
     )
-    if devices.dtype == torch.bfloat16 or devices.dtype == torch.float32:
-        if 'transformer' not in kwargs:
-            pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
-        if 'text_encoder' not in kwargs:
-            pipe.text_encoder = pipe.text_encoder.to(dtype=devices.dtype)
-        pipe.vae = pipe.vae.to(dtype=devices.dtype)
-    if devices.dtype == torch.float16:
-        if 'transformer' not in kwargs:
-            pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
-        if 'text_encoder' not in kwargs:
-            pipe.text_encoder = pipe.text_encoder.to(dtype=torch.float32) # gemma2 does not support fp16
-        pipe.vae = pipe.vae.to(dtype=torch.float32) # dc-ae often overflows in fp16
-    if shared.opts.diffusers_eval:
-        pipe.text_encoder.eval()
-        pipe.transformer.eval()
+
+    # only cast if not quant-loaded
+    try:
+        if devices.dtype == torch.bfloat16 or devices.dtype == torch.float32:
+            if 'transformer' not in kwargs:
+                pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
+            if 'text_encoder' not in kwargs:
+                pipe.text_encoder = pipe.text_encoder.to(dtype=devices.dtype)
+            pipe.vae = pipe.vae.to(dtype=devices.dtype)
+        if devices.dtype == torch.float16:
+            if 'transformer' not in kwargs:
+                pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
+            if 'text_encoder' not in kwargs:
+                pipe.text_encoder = pipe.text_encoder.to(dtype=torch.float32) # gemma2 does not support fp16
+            pipe.vae = pipe.vae.to(dtype=torch.float32) # dc-ae often overflows in fp16
+    except Exception as e:
+        shared.log.error(f'Load model: type=Sana {e}')
+
+    try:
+        if shared.opts.diffusers_eval:
+            pipe.text_encoder.eval()
+            pipe.transformer.eval()
+    except Exception:
+        pass
+
     t1 = time.time()
     shared.log.debug(f'Load model: type=Sana target={devices.dtype} te={pipe.text_encoder.dtype} transformer={pipe.transformer.dtype} vae={pipe.vae.dtype} time={t1-t0:.2f}')
     devices.torch_gc(force=True)

modules/processing_diffusers.py

@@ -147,9 +147,6 @@ def process_base(p: processing.StableDiffusionProcessing):
         hidiffusion.unapply()
         sd_models_compile.check_deepcache(enable=False)
 
-    if hasattr(shared.sd_model, 'embedding_db') and len(shared.sd_model.embedding_db.embeddings_used) > 0: # register used embeddings
-        p.extra_generation_params['Embeddings'] = ', '.join(shared.sd_model.embedding_db.embeddings_used)
-
     shared.state.nextjob()
     return output
 

modules/processing_info.py

@@ -7,10 +7,6 @@ from modules.processing_class import StableDiffusionProcessing
 args = {} # maintain history
 infotext = '' # maintain history
 debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
-if not shared.native:
-    from modules import sd_hijack
-else:
-    sd_hijack = None
 
 
 def get_last_args():
@@ -62,11 +58,9 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No
         "Refiner prompt": p.refiner_prompt if len(p.refiner_prompt) > 0 else None,
         "Refiner negative": p.refiner_negative if len(p.refiner_negative) > 0 else None,
         "Styles": "; ".join(p.styles) if p.styles is not None and len(p.styles) > 0 else None,
-        # sdnext
         "App": 'SD.Next',
         "Version": git_commit,
         "Backend": 'Legacy' if not shared.native else None,
-        "Pipeline": 'LDM' if not shared.native else None,
         "Parser": shared.opts.prompt_attention if shared.opts.prompt_attention != 'native' else None,
         "Comment": comment,
         "Operations": '; '.join(ops).replace('"', '') if len(p.ops) > 0 else 'none',
@@ -77,9 +71,9 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No
         args["VAE"] = 'TAESD'
     elif p.vae_type == 'Remote':
         args["VAE"] = 'Remote'
-    if shared.opts.add_model_name_to_info and getattr(shared.sd_model, 'sd_checkpoint_info', None) is not None:
+    if getattr(shared.sd_model, 'sd_checkpoint_info', None) is not None:
         args["Model"] = shared.sd_model.sd_checkpoint_info.model_name.replace(',', '').replace(':', '')
-    if shared.opts.add_model_hash_to_info and getattr(shared.sd_model, 'sd_model_hash', None) is not None:
+    if getattr(shared.sd_model, 'sd_model_hash', None) is not None:
         args["Model hash"] = shared.sd_model.sd_model_hash
     # native
     if grid is None and (p.n_iter > 1 or p.batch_size > 1) and index >= 0:
@@ -88,8 +82,10 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No
         args['Grid'] = grid
     if shared.native:
         args['Pipeline'] = shared.sd_model.__class__.__name__
-        args['TE'] = None if (not shared.opts.add_model_name_to_info or shared.opts.sd_text_encoder is None or shared.opts.sd_text_encoder == 'Default') else shared.opts.sd_text_encoder
-        args['UNet'] = None if (not shared.opts.add_model_name_to_info or shared.opts.sd_unet is None or shared.opts.sd_unet == 'Default') else shared.opts.sd_unet
+        args['TE'] = None if (shared.opts.sd_text_encoder is None or shared.opts.sd_text_encoder == 'Default') else shared.opts.sd_text_encoder
+        args['UNet'] = None if (shared.opts.sd_unet is None or shared.opts.sd_unet == 'Default') else shared.opts.sd_unet
+    else:
+        args['Pipeline'] = 'LDM'
     if 'txt2img' in p.ops:
         args["Variation seed"] = all_subseeds[index] if p.subseed_strength > 0 else None
         args["Variation strength"] = p.subseed_strength if p.subseed_strength > 0 else None
@@ -155,11 +151,14 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No
         args["Detailer negative"] = p.detailer_negative if len(p.detailer_negative) > 0 else None
     if 'color' in p.ops:
         args["Color correction"] = True
-    # embeddings
-    if sd_hijack is not None and hasattr(sd_hijack.model_hijack, 'embedding_db') and len(sd_hijack.model_hijack.embedding_db.embeddings_used) > 0: # this is for original hijaacked models only, diffusers are handled separately
-        args["Embeddings"] = ', '.join(sd_hijack.model_hijack.embedding_db.embeddings_used)
-    # samplers
+    if shared.opts.token_merging_method == 'ToMe': # tome/todo
+        args['ToMe'] = shared.opts.tome_ratio if shared.opts.tome_ratio != 0 else None
+    else:
+        args['ToDo'] = shared.opts.todo_ratio if shared.opts.todo_ratio != 0 else None
+    if hasattr(shared.sd_model, 'embedding_db') and len(shared.sd_model.embedding_db.embeddings_used) > 0: # register used embeddings
+        args['Embeddings'] = ', '.join(shared.sd_model.embedding_db.embeddings_used)
 
+    # samplers
     if getattr(p, 'sampler_name', None) is not None and p.sampler_name.lower() != 'default':
         args["Sampler eta delta"] = shared.opts.eta_noise_seed_delta if shared.opts.eta_noise_seed_delta != 0 and sd_samplers_common.is_sampler_using_eta_noise_seed_delta(p) else None
         args["Sampler eta multiplier"] = p.initial_noise_multiplier if getattr(p, 'initial_noise_multiplier', 1.0) != 1.0 else None
@@ -177,11 +176,10 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No
         args['Sampler range'] = shared.opts.schedulers_timesteps_range if shared.opts.schedulers_timesteps_range != shared.opts.data_labels.get('schedulers_timesteps_range').default else None
         args['Sampler shift'] = shared.opts.schedulers_shift if shared.opts.schedulers_shift != shared.opts.data_labels.get('schedulers_shift').default else None
         args['Sampler dynamic shift'] = shared.opts.schedulers_dynamic_shift if shared.opts.schedulers_dynamic_shift != shared.opts.data_labels.get('schedulers_dynamic_shift').default else None
-    # tome/todo
-    if shared.opts.token_merging_method == 'ToMe':
-        args['ToMe'] = shared.opts.tome_ratio if shared.opts.tome_ratio != 0 else None
-    else:
-        args['ToDo'] = shared.opts.todo_ratio if shared.opts.todo_ratio != 0 else None
+
+    # model specific
+    if shared.sd_model_type == 'h1':
+        args['LLM'] =  None if shared.opts.model_h1_llama_repo == 'Default' else shared.opts.model_h1_llama_repo
 
     args.update(p.extra_generation_params)
     for k, v in args.copy().items():

modules/shared.py

@@ -416,7 +416,7 @@ options_templates.update(options_section(('sd', "Models & Loading"), {
 
 options_templates.update(options_section(('model_options', "Models Options"), {
     "model_sd3_disable_te5": OptionInfo(False, "StableDiffusion3: T5 disable encoder"),
-    "model_h1_llama_repo": OptionInfo("meta-llama/Meta-Llama-3.1-8B-Instruct", "HiDream: LLama repo", gr.Textbox),
+    "model_h1_llama_repo": OptionInfo("Default", "HiDream: LLama repo", gr.Textbox),
 }))
 
 options_templates.update(options_section(('vae_encoder', "Variable Auto Encoder"), {