Merge pull request #5 from brantje/load-llava-on-demand

Load llava on demand

js/extension.js

@@ -28,7 +28,7 @@ app.registerExtension({
         for (const w of this.widgets) {
           if (["collapse_setting", "clear_canvas"].includes(w.name)) {
             // always show these widgets
-          } else if (["image", "original_image", "add_color_image", "add_edge_image", "remove_edge_image", "positive_prompt"].includes(w.name)) {
+          } else if (["image", "original_image", "add_color_image", "add_edge_image", "remove_edge_image"].includes(w.name)) {
             w.type = "hidden";
             w.value = null;
             w.computeSize = () => [0, -4];

js/magicquill_widget.js

@@ -266,7 +266,7 @@ export function MaigcQuillWidget(node, inputName, inputData, app) {
   };
 
   node.onDrawBackground = (ctx) => {
-    const setting_widgets = ["base_model_version", "negative_prompt", "dtype", "grow_size", "stroke_as_edge", "fine_edge", "edge_strength", "color_strength", "inpaint_strength", "seed", "steps", "cfg", "sampler_name", "scheduler"]
+    const setting_widgets = ["base_model_version","positive_prompt", "negative_prompt", "dtype", "grow_size", "stroke_as_edge", "fine_edge", "edge_strength", "color_strength", "inpaint_strength", "seed", "steps", "cfg", "sampler_name", "scheduler",  "optional_original_image_name", "optional_add_color_image_name", "optional_add_edge_image_name", "optional_remove_edge_image_name"]
     if (!this.flags.setting_collapsed) {
       for (const w of this.widgets) {
         if (setting_widgets.includes(w.name)) {

llava_new.py

@@ -23,15 +23,30 @@ import re
 class LLaVAModel:
     def __init__(self):
         # replace the model_path with correct path folder
-        base_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-        models_dir = os.path.join(base_path, "models")
-        model_path = os.path.join(models_dir, "llava-v1.5-7b-finetune-clean")
+        self.base_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+        self.models_dir = os.path.join(self.base_path, "models")
+        self.model_path = os.path.join(self.models_dir, "llava-v1.5-7b-finetune-clean")
+        self.tokenizer = None
+        self.model = None
+        self.image_processor = None
+        self.context_len = None
+    def load_model(self):
         self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
-            model_path=model_path,
+            model_path=self.model_path,
             model_base=None,
-            model_name=get_model_name_from_path(model_path),
+            model_name=get_model_name_from_path(self.model_path),
             
-        )
+        )   
+
+    def unload_model(self):
+        """Unload the model and clear GPU memory."""
+        if self.model is not None:
+            self.model.cpu()
+            del self.model
+            torch.cuda.empty_cache()
+        self.tokenizer = None
+        self.image_processor = None
+        self.context_len = None
 
     def generate_description(self, images, question):
         qs = question

magic_quill.py

@@ -310,7 +310,7 @@ async def run_magic_quill(request):
             steps=steps,
             cfg=cfg,
             sampler_name=sampler_name,
-            scheduler=scheduler
+            scheduler=scheduler,
         )
 
         # Convert the result tensors to base64
@@ -397,9 +397,18 @@ class MagicQuill(object):
                 "cfg": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "display": "slider"}),
                 "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler_ancestral"}),
                 "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "exponential"}),
+                # "optional_original_image_name": ("STRING", {"default": ""}),
+                # "optional_add_color_image_name": ("STRING", {"default": ""}),
+                # "optional_add_edge_image_name": ("STRING", {"default": ""}),
+                # "optional_remove_edge_image_name": ("STRING", {"default": ""}),
             },
             "optional": {
-             
+                "optional_image": ("IMAGE",),
+                "optional_image_mask": ("MASK",),
+                "optional_original_image": ("IMAGE",),
+                "optional_add_color_image": ("IMAGE",),
+                "optional_add_edge_mask": ("MASK",),
+                "optional_remove_edge_mask": ("MASK",),
             }
         }
 
@@ -412,6 +421,7 @@ class MagicQuill(object):
     @classmethod
     def prepare_images_and_masks(cls, image, original_image, add_color_image, add_edge_image, remove_edge_image):
         # Handle file path inputs
+        print(f"image: {image} original_image: {original_image} add_color_image: {add_color_image} add_edge_image: {add_edge_image} remove_edge_image: {remove_edge_image}")
         image_path = folder_paths.get_annotated_filepath(image)
         image_tensor = load_and_preprocess_image(image_path)
         height, width = image_tensor.shape[1], image_tensor.shape[2]
@@ -444,20 +454,42 @@ class MagicQuill(object):
 
     @classmethod
     def guess_prompt(cls, original_image_tensor, add_color_image_tensor, add_edge_mask):
+        cls.llavaModel.load_model()
         description, ans1, ans2 = cls.llavaModel.process(original_image_tensor, add_color_image_tensor, add_edge_mask)
         ans_list = []
         if ans1 and ans1 != "":
             ans_list.append(ans1)
         if ans2 and ans2 != "":
             ans_list.append(ans2)
-
+        cls.llavaModel.unload_model()
         return ", ".join(ans_list)
 
     @classmethod
-    def painter_execute(cls, image, original_image, add_color_image, add_edge_image, remove_edge_image, model, vae, clip, base_model_version, positive_prompt, negative_prompt, dtype, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
+    def painter_execute(cls, image, original_image, add_color_image, add_edge_image, remove_edge_image, model, vae, clip, base_model_version, positive_prompt, negative_prompt, dtype, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler, optional_image = None, optional_image_mask = None, optional_original_image = None, optional_original_image_mask = None, optional_add_color_image = None, optional_add_color_image_mask = None, optional_add_edge_mask = None, optional_add_edge_mask_mask = None, optional_remove_edge_mask = None, optional_remove_edge_mask_mask = None):
         print(f"model: {model} vae: {vae} clip: {clip} base_model_version: {base_model_version} positive_prompt: {positive_prompt} negative_prompt: {negative_prompt} dtype: {dtype} grow_size: {grow_size} stroke_as_edge: {stroke_as_edge} fine_edge: {fine_edge} edge_strength: {edge_strength} color_strength: {color_strength} inpaint_strength: {inpaint_strength} seed: {seed} steps: {steps} cfg: {cfg} sampler_name: {sampler_name} scheduler: {scheduler}")
         print(f"original_image: {original_image} add_color_image: {add_color_image} add_edge_image: {add_edge_image} remove_edge_image: {remove_edge_image}")
-        add_color_image, original_image, total_mask, add_edge_mask, remove_edge_mask = cls.prepare_images_and_masks(image, original_image, add_color_image, add_edge_image, remove_edge_image)
+        print(f"optional_image: {optional_image} optional_image_mask: {optional_image_mask} optional_original_image: {optional_original_image} optional_original_image_mask: {optional_original_image_mask} optional_add_color_image: {optional_add_color_image} optional_add_color_image_mask: {optional_add_color_image_mask} optional_add_edge_mask: {optional_add_edge_mask} optional_add_edge_mask_mask: {optional_add_edge_mask_mask} optional_remove_edge_mask: {optional_remove_edge_mask} optional_remove_edge_mask_mask: {optional_remove_edge_mask_mask}")
+        # check if optional_original_image is tensor
+        if isinstance(optional_image, torch.Tensor):
+            image = optional_image
+        if isinstance(optional_original_image, torch.Tensor):
+            original_image = optional_original_image
+        if isinstance(optional_add_color_image, torch.Tensor):
+            add_color_image = optional_add_color_image
+        if isinstance(optional_add_edge_mask, torch.Tensor):
+            add_edge_mask = optional_add_edge_mask
+        if isinstance(optional_remove_edge_mask, torch.Tensor):
+            remove_edge_mask = optional_remove_edge_mask       
+          
+        if isinstance(optional_image, torch.Tensor) and isinstance(optional_image_mask, torch.Tensor):
+            #if if not the same size, resize the mask
+            if optional_image_mask.shape[1] != optional_image.shape[1] or optional_image_mask.shape[2] != optional_image.shape[2]:
+                print("resizing mask")
+                optional_image_mask = F.interpolate(optional_image_mask.unsqueeze(0), size=(optional_image.shape[1], optional_image.shape[2]), mode='nearest').squeeze(0)
+            total_mask = optional_image_mask
+          
+        if not isinstance(image, torch.Tensor) and not isinstance(original_image, torch.Tensor) and not isinstance(add_color_image, torch.Tensor) and not isinstance(add_edge_image, torch.Tensor) and not isinstance(remove_edge_image, torch.Tensor):
+            add_color_image, original_image, total_mask, add_edge_mask, remove_edge_mask = cls.prepare_images_and_masks(image, original_image, add_color_image, add_edge_image, remove_edge_image)
 
         if torch.sum(remove_edge_mask).item() > 0 and torch.sum(add_edge_mask).item() == 0:
             if positive_prompt == "":
@@ -496,9 +528,5 @@ class MagicQuill(object):
         return m.digest().hex()
 
     @classmethod
-    def VALIDATE_INPUTS(self, image, original_image, add_color_image, add_edge_image, remove_edge_image, model, vae, clip, base_model_version, positive_prompt, negative_prompt, dtype, grow_size, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
-        if not folder_paths.exists_annotated_filepath(image):
-            print(image)
-            return "Invalid image file: {}".format(image)
-
+    def VALIDATE_INPUTS(self, image, original_image, add_color_image, add_edge_image, remove_edge_image, model, vae, clip, base_model_version, positive_prompt, negative_prompt, dtype, grow_size, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler, optional_image = None):
         return True

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ComfyUI-MagicQuill"
-version = "1.0.0"
+version = "1.0.4"
 description = "Fixed version of the original MagicQuill node."
 authors = ["brantje <brantje@gmail.com>"]
 license = { text = "MIT License" }
@@ -18,7 +18,7 @@ build-backend = "poetry.core.masonry.api"
 [project]
 name = "comfyui_magicquill_fixed" # Unique identifier for your node. Immutable after creation..
 description = "Fixed version of the original MagicQuill node. Required nodes: ComfyUI-Brushnet and ComfyUI Controlnet AUX"
-version = "1.0.0" # Custom Node version. Must be semantically versioned.
+version = "1.0.4" # Custom Node version. Must be semantically versioned.
 dependencies  = [
     'opencv-python',
     'diffusers',

requirements.txt

@@ -2,7 +2,7 @@ opencv-python
 diffusers
 torchsde==0.2.6
 protobuf==4.25.4
-transformers==4.37.2 
+transformers==4.38.0 
 tokenizers==0.15.1 
 sentencepiece==0.2.0
 shortuuid

scribble_color_edit.py

@@ -4,14 +4,46 @@ import torch
 import sys
 import torch.utils._pytree as pytree
 import numpy as np
-
+import subprocess
 current_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(current_dir)
-sys.path.append(os.path.abspath(os.path.join(current_dir, '..')))
+custom_nodes_dir = os.path.abspath(os.path.join(current_dir, '..'))
+sys.path.append(custom_nodes_dir)
 sys.path.append(os.path.abspath(os.path.join(current_dir, '..', '..', 'comfy_extras')))
-print(sys.path)
 
-from ComfyUI_BrushNet.brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint, get_files_with_extension
+
+brushnet_hyphen_dir = os.path.join(custom_nodes_dir, 'comfyui-brushnet')
+brushnet_underscore_dir = os.path.join(custom_nodes_dir, 'comfyui_brushnet')
+
+if not os.path.exists(brushnet_underscore_dir):
+    print(f"Creating symlink from {brushnet_hyphen_dir} to {brushnet_underscore_dir}")
+    
+    # Create the symlink - use different methods based on OS
+    if os.name == 'nt':  # Windows
+        # Requires admin privileges or developer mode
+        subprocess.run(['mklink', '/D', brushnet_underscore_dir, brushnet_hyphen_dir], shell=True)
+    else:  # Unix/Linux/Mac
+        os.symlink(brushnet_hyphen_dir, brushnet_underscore_dir)
+    
+    print(f"Symlink created: {os.path.exists(brushnet_underscore_dir)}")
+
+# Now try importing from the symlinked directory
+try:
+    # Add to path
+    sys.path.append(custom_nodes_dir)
+    
+    # Import from symlinked directory
+    from comfyui_brushnet.brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint, get_files_with_extension
+    print("Successfully imported from symlinked directory")
+except ImportError as e:
+    print(f"Import from symlink failed: {e}")
+    try:
+        from ComfyUI_BrushNet.brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint, get_files_with_extension
+    except ImportError as e:
+        print(f"Import from ComfyUI_BrushNet failed: {e}")
+        raise ImportError("Failed to import even with ComfyUI_BrushNet. Please check file permissions and structure.")
+
+
 from comfyui_controlnet_aux.node_wrappers.lineart import LineArt_Preprocessor
 from comfyui_controlnet_aux.node_wrappers.pidinet import PIDINET_Preprocessor
 from comfyui_controlnet_aux.node_wrappers.color import Color_Preprocessor