Prompt cache support for HiDream

modules/processing_args.py

@@ -143,7 +143,8 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
     if (prompt_attention != 'fixed') and ('Onnx' not in model.__class__.__name__) and ('prompt' not in p.task_args) and (
         'StableDiffusion' in model.__class__.__name__ or
         'StableCascade' in model.__class__.__name__ or
-        'Flux' in model.__class__.__name__
+        'Flux' in model.__class__.__name__ or
+        'HiDreamImage' in model.__class__.__name__
     ):
         try:
             prompt_parser_diffusers.embedder = prompt_parser_diffusers.PromptEmbedder(prompts, negative_prompts, steps, clip_skip, p)
@@ -162,25 +163,31 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
             prompts = [p.replace('|image|', '<|image_1|>') for p in prompts]
         if hasattr(model, 'text_encoder') and hasattr(model, 'tokenizer') and 'prompt_embeds' in possible and prompt_parser_diffusers.embedder is not None:
             args['prompt_embeds'] = prompt_parser_diffusers.embedder('prompt_embeds')
-            if 'StableCascade' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('positive_pooleds').unsqueeze(0)
-            elif 'XL' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
-            elif 'StableDiffusion3' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
-            elif 'Flux' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+            if prompt_parser_diffusers.embedder is not None:
+                if 'StableCascade' in model.__class__.__name__:
+                    args['prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('positive_pooleds').unsqueeze(0)
+                elif 'XL' in model.__class__.__name__:
+                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+                elif 'StableDiffusion3' in model.__class__.__name__:
+                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+                elif 'Flux' in model.__class__.__name__:
+                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+                elif 'HiDreamImage' in model.__class__.__name__:
+                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
         else:
             args['prompt'] = prompts
     if 'negative_prompt' in possible:
         if hasattr(model, 'text_encoder') and hasattr(model, 'tokenizer') and 'negative_prompt_embeds' in possible and prompt_parser_diffusers.embedder is not None:
             args['negative_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_prompt_embeds')
-            if 'StableCascade' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['negative_prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('negative_pooleds').unsqueeze(0)
-            if 'XL' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
-            if 'StableDiffusion3' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
+            if prompt_parser_diffusers.embedder is not None:
+                if 'StableCascade' in model.__class__.__name__:
+                    args['negative_prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('negative_pooleds').unsqueeze(0)
+                elif 'XL' in model.__class__.__name__:
+                    args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
+                elif 'StableDiffusion3' in model.__class__.__name__:
+                    args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
+                elif 'HiDreamImage' in model.__class__.__name__:
+                    args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
         else:
             if 'PixArtSigmaPipeline' in model.__class__.__name__: # pixart-sigma pipeline throws list-of-list for negative prompt
                 args['negative_prompt'] = negative_prompts[0]

modules/prompt_parser_diffusers.py

@@ -26,7 +26,8 @@ def prompt_compatible(pipe = None):
         'StableDiffusion' not in pipe.__class__.__name__ and
         'DemoFusion' not in pipe.__class__.__name__ and
         'StableCascade' not in pipe.__class__.__name__ and
-        'Flux' not in pipe.__class__.__name__
+        'Flux' not in pipe.__class__.__name__ and
+        'HiDreamImage' not in pipe.__class__.__name__
     ):
         shared.log.warning(f"Prompt parser not supported: {pipe.__class__.__name__}")
         return False
@@ -190,14 +191,32 @@ class PromptEmbedder:
     def __call__(self, key, step=0):
         batch = getattr(self, key)
         res = []
-        for i in range(self.batchsize):
-            if len(batch[i]) == 0:  # if asking for a null key, ie pooled on SD1.5
-                return None
-            try:
-                res.append(batch[i][step])
-            except IndexError:
-                res.append(batch[i][0])  # if not scheduled, return default
-        return torch.cat(res)
+        if isinstance(batch[0][0], list) and len(batch[0][0]) == 2 and isinstance(batch[0][0][1], torch.Tensor) and batch[0][0][1].shape[0] == 32:
+            # hidream uses a list of t5 + llama prompt embeds: [t5_embeds, llama_embeds]
+            # t5_embeds shape: [batch_size, seq_len, dim]
+            # llama_embeds shape: [number_of_hidden_states, batch_size, seq_len, dim]
+            res2 = []
+            for i in range(self.batchsize):
+                if len(batch[i]) == 0:  # if asking for a null key, ie pooled on SD1.5
+                    return None
+                try:
+                    res.append(batch[i][step][0])
+                    res2.append(batch[i][step][1])
+                except IndexError:
+                    # if not scheduled, return default
+                    res.append(batch[i][0][0])
+                    res2.append(batch[i][0][1])
+            res = [torch.cat(res, dim=0), torch.cat(res2, dim=1)]
+            return res
+        else:
+            for i in range(self.batchsize):
+                if len(batch[i]) == 0:  # if asking for a null key, ie pooled on SD1.5
+                    return None
+                try:
+                    res.append(batch[i][step])
+                except IndexError:
+                    res.append(batch[i][0])  # if not scheduled, return default
+            return torch.cat(res)
 
 
 def compel_hijack(self, token_ids: torch.Tensor, attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor:
@@ -457,21 +476,43 @@ def split_prompts(pipe, prompt, SD3 = False):
     else:
         prompt3 = prompt
 
+    if prompt.find("TE4:") != -1:
+        prompt, prompt4 = prompt.split("TE4:")
+    elif prompt2.find("TE4:") != -1:
+        prompt2, prompt4 = prompt2.split("TE4:")
+    elif prompt3.find("TE4:") != -1:
+        prompt3, prompt4 = prompt3.split("TE4:")
+    else:
+        prompt4 = prompt
+
     prompt = prompt.strip()
     prompt2 = " " if prompt2.strip() == "" else prompt2.strip()
     prompt3 = " " if prompt3.strip() == "" else prompt3.strip()
+    prompt4 = " " if prompt4.strip() == "" else prompt4.strip()
 
     if SD3 and prompt3 != " ":
         ps, _ws = get_prompts_with_weights(pipe, prompt3)
         prompt3 = " ".join(ps)
-    return prompt, prompt2, prompt3
+    return prompt, prompt2, prompt3, prompt4
 
 
 def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", clip_skip: int = None):
     device = devices.device
-    SD3 = hasattr(pipe, 'text_encoder_3')
-    prompt, prompt_2, prompt_3 = split_prompts(pipe, prompt, SD3)
-    neg_prompt, neg_prompt_2, neg_prompt_3 = split_prompts(pipe, neg_prompt, SD3)
+    SD3 = bool(hasattr(pipe, 'text_encoder_3') and not hasattr(pipe, 'text_encoder_4'))
+    prompt, prompt_2, prompt_3, prompt_4 = split_prompts(pipe, prompt, SD3)
+    neg_prompt, neg_prompt_2, neg_prompt_3, neg_prompt_4 = split_prompts(pipe, neg_prompt, SD3)
+
+    if "Flux" in pipe.__class__.__name__: # clip is only used for the pooled embeds
+        prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt(prompt=prompt, prompt_2=prompt_2, device=device, num_images_per_prompt=1)
+        return prompt_embeds, pooled_prompt_embeds, None, None # no negative support
+
+    if "HiDreamImage" in pipe.__class__.__name__: # clip is only used for the pooled embeds
+        prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
+            prompt=prompt, prompt_2=prompt_2, prompt_3=prompt_3, prompt_4=prompt_4,
+            negative_prompt=neg_prompt, negative_prompt_2=neg_prompt_2, negative_prompt_3=neg_prompt_3, negative_prompt_4=neg_prompt_4,
+            device=device, num_images_per_prompt=1,
+        )
+        return prompt_embeds, pooled_prompt_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds
 
     if prompt != prompt_2:
         ps = [get_prompts_with_weights(pipe, p) for p in [prompt, prompt_2]]
@@ -488,10 +529,6 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c
         negatives.pop(0)
         negative_weights.pop(0)
 
-    if "Flux" in pipe.__class__.__name__: # clip is only used for the pooled embeds
-        prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt(prompt=prompt, prompt_2=prompt_2, device=device, num_images_per_prompt=1)
-        return prompt_embeds, pooled_prompt_embeds, None, None # no negative support
-
     embedding_providers = prepare_embedding_providers(pipe, clip_skip)
     empty_embedding_providers = None
     if 'StableCascade' in pipe.__class__.__name__:
@@ -593,8 +630,8 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c
 
 def get_xhinker_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", clip_skip: int = None):
     is_sd3 = hasattr(pipe, 'text_encoder_3')
-    prompt, prompt_2, _prompt_3 = split_prompts(pipe, prompt, is_sd3)
-    neg_prompt, neg_prompt_2, _neg_prompt_3 = split_prompts(pipe, neg_prompt, is_sd3)
+    prompt, prompt_2, _prompt_3, _ = split_prompts(pipe, prompt, is_sd3)
+    neg_prompt, neg_prompt_2, _neg_prompt_3, _ = split_prompts(pipe, neg_prompt, is_sd3)
     try:
         prompt = pipe.maybe_convert_prompt(prompt, pipe.tokenizer)
         neg_prompt = pipe.maybe_convert_prompt(neg_prompt, pipe.tokenizer)

modules/sd_hijack_te.py

@@ -15,6 +15,8 @@ def hijack_encode_prompt(*args, **kwargs):
         res = None
     t1 = time.time()
     timer.process.add('te', t1-t0)
+    if hasattr(shared.sd_model, "maybe_free_model_hooks"):
+        shared.sd_model.maybe_free_model_hooks()
     shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
     return res