Add Omost DenseDiffusion regional prompt support

omost_nodes.py

@@ -10,6 +10,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 import comfy.model_management
+from comfy.model_patcher import ModelPatcher
 from comfy.sd import CLIP
 from nodes import CLIPTextEncode, ConditioningSetMask
 from .lib_omost.canvas import (
@@ -40,6 +41,9 @@ def create_logger(level=logging.INFO):
 
 logger = create_logger(level=logging.INFO)
 
+# Canvas size used in original Omost repo.
+CANVAS_SIZE = 90
+
 
 # Type definitions.
 class OmostConversationItem(TypedDict):
@@ -354,6 +358,71 @@ def encode(tokens: CLIPTokens) -> EncoderOutput:
         ]
 
 
+class OmostDenseDiffusionLayoutNode:
+    """Apply Omost layout with Omost's area condition system. This is the regional
+    prompt system implemented in the original Omost repo.
+
+    You need to install https://github.com/huchenlei/ComfyUI_densediffusion to use this node.
+    """
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "canvas_conds": ("OMOST_CANVAS_CONDITIONING",),
+                "clip": ("CLIP",),
+            },
+        }
+
+    RETURN_TYPES = ("MODEL", "CONDITIONING")
+    FUNCTION = "layout_cond"
+    CATEGORY = "omost"
+
+    def __init__(self):
+        try:
+            from custom_nodes.ComfyUI_densediffusion.densediffusion_node import (
+                DenseDiffusionApplyNode,
+                DenseDiffusionAddCondNode,
+            )
+        except Exception as e:
+            logger.error(
+                "Failed to import ComfyUI_densediffusion. Make sure it's installed."
+                "https://github.com/huchenlei/ComfyUI_densediffusion"
+            )
+            raise e
+
+        self.dense_diffusion_apply_node = DenseDiffusionApplyNode()
+        self.dense_diffusion_add_cond_node = DenseDiffusionAddCondNode()
+
+    def layout_cond(
+        self,
+        model: ModelPatcher,
+        canvas_conds: list[OmostCanvasCondition],
+        clip: CLIP,
+    ) -> tuple[ModelPatcher, ComfyUIConditioning]:
+        """Layout conditioning"""
+        work_model: ModelPatcher = model.clone()
+
+        for canvas_cond in canvas_conds:
+            mask = torch.ones([CANVAS_SIZE, CANVAS_SIZE], dtype=torch.float32)
+            a, b, c, d = canvas_cond["rect"]
+            mask[a:b, c:d] = 1.0
+
+            cond: ComfyUIConditioning = PromptEncoding.encode_bag_of_subprompts_greedy(
+                clip, canvas_cond["prefixes"], canvas_cond["suffixes"]
+            )
+            # Set area cond
+            work_model = self.dense_diffusion_add_cond_node.append(
+                work_model,
+                conditioning=cond,
+                mask=mask,
+                strength=1.0,
+            )[0]
+
+        return self.dense_diffusion_apply_node.apply(work_model)
+
+
 class OmostComfyLayoutNode:
     """Apply Omost layout with ComfyUI's area condition system."""
 
@@ -402,7 +471,6 @@ def calc_cond_mask(
         method: AreaOverlapMethod = AreaOverlapMethod.OVERLAY,
     ) -> list[OmostCanvasCondition]:
         """Calculate canvas cond mask."""
-        CANVAS_SIZE = 90
         assert len(canvas_conds) > 0
         canvas_conds = canvas_conds.copy()
 
@@ -507,6 +575,7 @@ def load_canvas(self, json_str: str) -> Tuple[list[OmostCanvasCondition]]:
     "OmostLLMLoaderNode": OmostLLMLoaderNode,
     "OmostLLMChatNode": OmostLLMChatNode,
     "OmostLayoutCondNode": OmostComfyLayoutNode,
+    "OmostDenseDiffusionLayoutNode": OmostDenseDiffusionLayoutNode,
     "OmostLoadCanvasConditioningNode": OmostLoadCanvasConditioningNode,
     "OmostRenderCanvasConditioningNode": OmostRenderCanvasConditioningNode,
 }
@@ -515,6 +584,7 @@ def load_canvas(self, json_str: str) -> Tuple[list[OmostCanvasCondition]]:
     "OmostLLMLoaderNode": "Omost LLM Loader",
     "OmostLLMChatNode": "Omost LLM Chat",
     "OmostLayoutCondNode": "Omost Layout Cond (ComfyUI-Area)",
+    "OmostDenseDiffusionLayoutNode": "Omost Layout Cond (OmostDenseDiffusion)",
     "OmostLoadCanvasConditioningNode": "Omost Load Canvas Conditioning",
     "OmostRenderCanvasConditioningNode": "Omost Render Canvas Conditioning",
 }