From 78d8bfeba79932755c6f1f6af6bcd19b358d45ba Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Fri, 11 Apr 2025 21:27:48 -0400
Subject: [PATCH] hidream allow custom llama

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 CHANGELOG.md             |  1 +
 modules/model_hidream.py | 16 ++++++++++------
 modules/shared.py        |  1 +
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 48615ddd4..3ddcea45d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@
     and 4 (!?) text-encoders: *clip-l, clip-g, t5-1.1-xxl, llama-3.1-8b-instruct* for total of **8.3B** parameters  
     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  
+    custom `llama` model can be set in *settings -> model options*
     difference between variants is recommended number of steps: *fast=16, dev=28, full=50*  
     hidream supportability: *offloading, quantization, taesd live-preview, remote-vae*  
     hidream compatibility: *flowmatching samplers*  
diff --git a/modules/model_hidream.py b/modules/model_hidream.py
index a345f809d..82cc38218 100644
--- a/modules/model_hidream.py
+++ b/modules/model_hidream.py
@@ -5,9 +5,6 @@ import diffusers
 from modules import shared, devices, sd_models, timer, model_quant, modelloader
 
 
-llama_repo = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-
-
 def hijack_encode_prompt(*args, **kwargs):
     t0 = time.time()
     res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
@@ -57,6 +54,8 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
         **load_args,
         **quant_args,
     )
+    if shared.opts.diffusers_offload_mode != 'none':
+        transformer = transformer.to(devices.cpu)
 
     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}')
@@ -67,22 +66,27 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
         **load_args,
         **quant_args,
     )
+    if shared.opts.diffusers_offload_mode != 'none':
+        text_encoder_3 = text_encoder_3.to(devices.cpu)
 
     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}')
+    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}')
     tokenizer_4 = transformers.PreTrainedTokenizerFast.from_pretrained(
-        llama_repo,
+        shared.opts.model_h1_llama_repo,
         cache_dir=shared.opts.hfcache_dir,
         **load_args,
     )
     text_encoder_4 = transformers.LlamaForCausalLM.from_pretrained(
-        llama_repo,
+        shared.opts.model_h1_llama_repo,
         output_hidden_states=True,
         output_attentions=True,
         cache_dir=shared.opts.hfcache_dir,
         **load_args,
         **quant_args,
     )
+    if shared.opts.diffusers_offload_mode != 'none':
+        text_encoder_4 = text_encoder_4.to(devices.cpu)
+
     load_args, quant_args = get_args(diffusers_load_config, module='Model')
     pipe = diffusers.HiDreamImagePipeline.from_pretrained(
         repo_id,
diff --git a/modules/shared.py b/modules/shared.py
index 04a5e822f..72b2e0994 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -416,6 +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, "SD3 disable T5 encoder"),
+    "model_h1_llama_repo": OptionInfo("meta-llama/Meta-Llama-3.1-8B-Instruct", "LLama repo", gr.Textbox),
 }))
 
 options_templates.update(options_section(('vae_encoder', "Variable Auto Encoder"), {