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

Load llava on demand

.github/workflows/main.yml

@@ -0,0 +1,20 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+    paths:
+      - "pyproject.toml"
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} ## Add your own personal access token to your Github Repository secrets and reference it here.

README.md

@@ -4,6 +4,19 @@
 <a href="https://huggingface.co/spaces/AI4Editing/MagicQuill"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)"></a>
 <a href="https://creativecommons.org/licenses/by-sa/4.0/"><img src="https://img.shields.io/badge/License-CC%20BY--SA%204.0-lightgrey.svg"></a>
 
+
+# Fixed
+- [x] SD1.5
+
+Todo
+- [ ] SDXL
+- [ ] FLUX 
+
+
+Note:
+[Brushnet](https://github.com/nullquant/ComfyUI-BrushNet) and [ComfyUI ControlNet Aux](https://github.com/Fannovel16/comfyui_controlnet_aux) are required.
+
+
 https://github.com/user-attachments/assets/8ee9663a-fef2-484a-a0b7-8427ab590424
 
 There is an HD video on [Youtube](https://www.youtube.com/watch?v=5DiKfONMnE4).

install.py

@@ -8,13 +8,3 @@ print(models_dir)
 # Download the model from the Hugging Face Hub
 repo_id = "LiuZichen/MagicQuill-models"  # or any other model/dataset ID
 snapshot_download(repo_id=repo_id, local_dir=models_dir)
-
-repos = {
-    'ComfyUI_BrushNet': 'https://github.com/nullquant/ComfyUI-BrushNet',
-    'comfyui_controlnet_aux': 'https://github.com/Fannovel16/comfyui_controlnet_aux'
-}
-
-# for name, url in repos.items():
-#     target_dir = os.path.join(parent_dir, name)
-#     if not os.path.exists(target_dir):
-#         subprocess.run(['git', 'clone', url, target_dir], check=True)

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

@@ -1,6 +1,7 @@
 import hashlib
 import os
 import json
+import random
 from server import PromptServer
 from PIL import Image, ImageOps
 import torch
@@ -9,12 +10,20 @@ import folder_paths
 from aiohttp import web
 import io
 import base64
+import time
 
 import comfy.samplers
 from .scribble_color_edit import ScribbleColorEditModel
 from .llava_new import LLaVAModel
+import torch.nn.functional as F
 
 def tensor_to_base64(tensor):
+    if isinstance(tensor, dict) and "samples" in tensor:
+        # Handle dictionary with 'samples' key (latent)
+        # For latent, we'll just return an empty string since latents aren't viewable directly
+        return ""  # or implement specific handling for latent samples
+    
+    # For tensor input, process normally
     tensor = tensor.squeeze(0) * 255.
     pil_image = Image.fromarray(tensor.cpu().byte().numpy())
     buffered = io.BytesIO()
@@ -81,25 +90,268 @@ async def process_background_img(request):
 @PromptServer.instance.routes.post("/magic_quill/guess_prompt")
 async def guess_prompt_handler(request):
     json_data = await request.json()
-    add_color_image = json_data.get("add_color_image", None)
-    original_image = json_data.get("original_image", None)
-    add_edge_image = json_data.get("add_edge_image", None)
+    add_color_image_data = json_data.get("add_color_image", None)
+    original_image_data = json_data.get("original_image", None)
+    add_edge_image_data = json_data.get("add_edge_image", None)
 
-    original_image_path = folder_paths.get_annotated_filepath(original_image)
-    original_image_tensor = load_and_preprocess_image(original_image_path)
-    
-    if add_color_image:
-        add_color_image_path = folder_paths.get_annotated_filepath(add_color_image)
-        add_color_image_tensor = load_and_preprocess_image(add_color_image_path)
-    else:
-        add_color_image_tensor = original_image_tensor
-    
-    width, height = original_image_tensor.shape[1], original_image_tensor.shape[2]
-    add_edge_mask = create_alpha_mask(folder_paths.get_annotated_filepath(add_edge_image)) if add_edge_image else torch.zeros((1, height, width), dtype=torch.float32, device="cpu")
+    if not original_image_data:
+        return web.json_response({"error": "Original image is required."}, status=400)
 
-    res = MagicQuill.guess_prompt(original_image_tensor, add_color_image_tensor, add_edge_mask)
+    temp_files_to_clean = []
+    input_dir = folder_paths.get_input_directory()
+
+    def is_base64(s):
+        return isinstance(s, str) and s.startswith("data:image/")
+
+    def handle_image_input(image_data, filename_base):
+        if is_base64(image_data):
+            try:
+                image = read_base64_image(image_data)
+                timestamp = int(time.time())
+                # Use a short hash to minimize collision chance but keep filename reasonable
+                hash_part = hashlib.sha1(str(timestamp).encode() + image_data.encode()).hexdigest()[:8]
+                filename = f"{filename_base}_{timestamp}_{hash_part}.png"
+                filepath = os.path.join(input_dir, filename)
+                image.save(filepath)
+                print(f"Saved temporary image to {filepath}")
+                temp_files_to_clean.append(filepath)
+                return filepath
+            except Exception as e:
+                print(f"Error processing base64 image for {filename_base}: {e}")
+                # Raise or return error? For now, let it raise to signal failure.
+                raise ValueError(f"Invalid base64 data for {filename_base}") from e
+        elif isinstance(image_data, str): # Assume filename
+            return folder_paths.get_annotated_filepath(image_data)
+        else: # Handle None or other invalid types
+            return None
+
+    try:
+        original_image_path = handle_image_input(original_image_data, "guess_original")
+        if not original_image_path or not os.path.exists(original_image_path):
+             return web.json_response({"error": f"Original image not found or invalid: {original_image_data}"}, status=400)
+             
+        original_image_tensor = load_and_preprocess_image(original_image_path)
+        
+        add_color_image_path = None
+        if add_color_image_data:
+            add_color_image_path = handle_image_input(add_color_image_data, "guess_add_color")
+            if not add_color_image_path or not os.path.exists(add_color_image_path):
+                print(f"Warning: Add color image specified but not found or invalid: {add_color_image_data}. Using original image.")
+                add_color_image_path = None # Fallback
+
+        if add_color_image_path:
+            add_color_image_tensor = load_and_preprocess_image(add_color_image_path)
+        else:
+            add_color_image_tensor = original_image_tensor # Fallback to original
+        
+        add_edge_image_path = None
+        if add_edge_image_data:
+            add_edge_image_path = handle_image_input(add_edge_image_data, "guess_add_edge")
+            if not add_edge_image_path or not os.path.exists(add_edge_image_path):
+                 print(f"Warning: Add edge image specified but not found or invalid: {add_edge_image_data}. Ignoring add edge.")
+                 add_edge_image_path = None # Ignore if invalid
+
+        width, height = original_image_tensor.shape[2], original_image_tensor.shape[1] # Corrected order W, H
+        add_edge_mask = create_alpha_mask(add_edge_image_path) if add_edge_image_path else torch.zeros((1, height, width), dtype=torch.float32, device="cpu")
+
+        # Ensure mask dimensions match original image
+        if add_edge_mask.shape[1] != height or add_edge_mask.shape[2] != width:
+             add_edge_mask = F.interpolate(add_edge_mask.unsqueeze(0), size=(height, width), mode='nearest').squeeze(0)
+
+        res = MagicQuill.guess_prompt(original_image_tensor, add_color_image_tensor, add_edge_mask)
+
+        return web.json_response({"prompt": res, "error": False})
+
+    except Exception as e:
+        import traceback
+        traceback_str = traceback.format_exc()
+        print(f"Error in guess_prompt_handler: {str(e)}\nTraceback: {traceback_str}")
+        return web.json_response({"error": str(e), "traceback": traceback_str}, status=500)
+    finally:
+        # Cleanup temporary files
+        for filepath in temp_files_to_clean:
+            try:
+                if os.path.exists(filepath):
+                    os.remove(filepath)
+                    print(f"Cleaned up temporary file: {filepath}")
+            except Exception as e:
+                print(f"Warning: Error cleaning up temporary file {filepath}: {str(e)}")
+
+@PromptServer.instance.routes.post("/magic_quill/run")
+async def run_magic_quill(request):
+    try:
+        post = await request.json()
+        base64_image = post.get("image")
+        if not base64_image:
+            return web.json_response({"error": "No image provided"}, status=400)
+
+        # Generate a unique prompt ID for this request
+        prompt_id = str(hashlib.sha256(str(time.time()).encode()).hexdigest()[:8])
+        PromptServer.instance.last_prompt_id = prompt_id
+        PromptServer.instance.last_node_id = "magic_quill"
+        
+        # Get the input directory path
+        input_dir = folder_paths.get_input_directory()
+        
+        # Generate unique filenames for this request
+        timestamp = int(time.time())
+        if post.get("dynamic_filenames", False):
+            main_image_filename = f"api_image_{timestamp}.png"
+            original_image_filename = f"api_original_{timestamp}.png"
+            add_color_image_filename = f"api_add_color_{timestamp}.png"
+            add_edge_image_filename = f"api_add_edge_{timestamp}.png"
+            remove_edge_image_filename = f"api_remove_edge_{timestamp}.png"
+        else:
+            main_image_filename = f"clipspace-mask-MagicQuill_-1.png"
+            original_image_filename = f"original_MagicQuill_-1.png"
+            add_color_image_filename = f"add_color_MagicQuill_-1.png"
+            add_edge_image_filename = f"add_edge_MagicQuill_-1.png"
+            remove_edge_image_filename = f"remove_edge_MagicQuill_-1.png"   
+        
+        # Save main image
+        main_image_path = os.path.join(input_dir, main_image_filename)
+        image = read_base64_image(base64_image)
+        image.save(main_image_path)
+        print(f"Saved main image to {main_image_path}")
+        
+        # Save original image if provided, otherwise use main image
+        base64_original = post.get("original_image")
+        if base64_original:
+            original_image_path = os.path.join(input_dir, original_image_filename)
+            original_image = read_base64_image(base64_original)
+            original_image.save(original_image_path)
+            original_image_file = original_image_filename
+        else:
+            original_image_file = main_image_filename
+        
+        # Process and save optional images if provided
+        add_color_image_file = None
+        add_edge_image_file = None
+        remove_edge_image_file = None
+        
+        # Save add_color_image if provided
+        base64_add_color = post.get("add_color_image")
+        if base64_add_color:
+            add_color_image_path = os.path.join(input_dir, add_color_image_filename)
+            add_color_image = read_base64_image(base64_add_color)
+            add_color_image.save(add_color_image_path)
+            add_color_image_file = add_color_image_filename
+            
+        # Save add_edge_image if provided
+        base64_add_edge = post.get("add_edge_image")
+        if base64_add_edge:
+            add_edge_image_path = os.path.join(input_dir, add_edge_image_filename)
+            add_edge_image = read_base64_image(base64_add_edge)
+            add_edge_image.save(add_edge_image_path)
+            add_edge_image_file = add_edge_image_filename
+            
+        # Save remove_edge_image if provided
+        base64_remove_edge = post.get("remove_edge_image")
+        if base64_remove_edge:
+            remove_edge_image_path = os.path.join(input_dir, remove_edge_image_filename)
+            remove_edge_image = read_base64_image(base64_remove_edge)
+            remove_edge_image.save(remove_edge_image_path)
+            remove_edge_image_file = remove_edge_image_filename
+        
+        # Get other parameters from the request
+        checkpoint_name = post.get("checkpoint_name", "SD1.5/DreamShaper.safetensors")
+        ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint_name)
+        out = comfy.sd.load_checkpoint_guess_config(ckpt_path)
+        
+        model = out[0]
+        clip = out[1]
+        vae = out[2]
+        
+        random_seed = random.randint(0, 0xffffffffffffffff)
+        
+        if not model or not vae or not clip:
+            return web.json_response({"error": "Missing required model objects"}, status=400)
+            
+        base_model_version = post.get("base_model_version", "SD1.5")
+        positive_prompt = post.get("positive_prompt", "")
+        negative_prompt = post.get("negative_prompt", "")
+        dtype = post.get("dtype", "float16")
+        grow_size = post.get("grow_size", 15)
+        stroke_as_edge = post.get("stroke_as_edge", "enable")
+        fine_edge = post.get("fine_edge", "disable")
+        edge_strength = post.get("edge_strength", 0.5)
+        color_strength = post.get("color_strength", 0.5)
+        inpaint_strength = post.get("inpaint_strength", 1.0)
+        
+        seed = post.get("seed", random_seed) #Random seed
+        steps = post.get("steps", 20)
+        cfg = post.get("cfg", 4.0)
+        sampler_name = post.get("sampler_name", "euler_ancestral")
+        scheduler = post.get("scheduler", "exponential")
+
+        print(f"Using files - Main: {main_image_filename}, Original: {original_image_file}, Add Color: {add_color_image_file}, Add Edge: {add_edge_image_file}, Remove Edge: {remove_edge_image_file}")
+
+        # Call painter_execute with file paths instead of tensors
+        result = MagicQuill.painter_execute(
+            image=main_image_filename,
+            original_image=original_image_file,
+            add_color_image=add_color_image_file,
+            add_edge_image=add_edge_image_file,
+            remove_edge_image=remove_edge_image_file,
+            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,
+        )
+
+        # Convert the result tensors to base64
+        latent, image, edge_map, color_palette = result
+        
+        # Send progress update
+        PromptServer.instance.send_sync(
+            "progress", {"value": 1, "max": 1, "prompt_id": prompt_id, "node": "magic_quill"}
+        )
+        
+        # Clean up temporary files (optional)
+        if post.get("cleanup", False):
+            try:
+                if os.path.exists(main_image_path):
+                    os.remove(main_image_path)
+                if base64_original and os.path.exists(os.path.join(input_dir, original_image_filename)):
+                    os.remove(os.path.join(input_dir, original_image_filename))
+                if base64_add_color and os.path.exists(os.path.join(input_dir, add_color_image_filename)):
+                    os.remove(os.path.join(input_dir, add_color_image_filename))
+                if base64_add_edge and os.path.exists(os.path.join(input_dir, add_edge_image_filename)):
+                    os.remove(os.path.join(input_dir, add_edge_image_filename))
+                if base64_remove_edge and os.path.exists(os.path.join(input_dir, remove_edge_image_filename)):
+                    os.remove(os.path.join(input_dir, remove_edge_image_filename))
+            except Exception as e:
+                print(f"Warning: Error cleaning up temporary files: {str(e)}")
+        
+        return web.json_response({
+            "status": "success",
+            "result": {
+                "latent": tensor_to_base64(latent),
+                "image": tensor_to_base64(image),
+                "edge_map": tensor_to_base64(edge_map),
+                "color_palette": tensor_to_base64(color_palette)
+            }
+        })
+    except Exception as e:
+        import traceback
+        traceback_str = traceback.format_exc()
+        print(f"Error: {str(e)}\nTraceback: {traceback_str}")
+        return web.json_response({"error": str(e), "traceback": traceback_str}, status=500)
 
-    return web.json_response({"prompt": res, "error": False})
 
 class MagicQuill(object):
     scribbleColorEditModel = ScribbleColorEditModel()
@@ -145,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",),
             }
         }
 
@@ -159,11 +420,11 @@ 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)
-        
-        width, height = image_tensor.shape[1], image_tensor.shape[2]
-        
+        height, width = image_tensor.shape[1], image_tensor.shape[2]
         total_mask = create_alpha_mask(image_path)
         
         original_image_path = folder_paths.get_annotated_filepath(original_image)
@@ -175,27 +436,60 @@ class MagicQuill(object):
         else:
             add_color_image_tensor = original_image_tensor
         
-        add_edge_mask = create_alpha_mask(folder_paths.get_annotated_filepath(add_edge_image)) if add_edge_image else torch.zeros_like(total_mask)
+        add_edge_mask = create_alpha_mask(folder_paths.get_annotated_filepath(add_edge_image)) if add_edge_image else torch.zeros((1, height, width), dtype=torch.float32, device="cpu")
+        remove_edge_mask = create_alpha_mask(folder_paths.get_annotated_filepath(remove_edge_image)) if remove_edge_image else torch.zeros((1, height, width), dtype=torch.float32, device="cpu")
         
-        remove_edge_mask = create_alpha_mask(folder_paths.get_annotated_filepath(remove_edge_image)) if remove_edge_image else torch.zeros_like(total_mask)
+        # Ensure all tensors have correct dimensions
+        if add_edge_mask.shape[1] != height or add_edge_mask.shape[2] != width:
+            add_edge_mask = F.interpolate(add_edge_mask.unsqueeze(0), size=(height, width), mode='nearest').squeeze(0)
+            
+        if remove_edge_mask.shape[1] != height or remove_edge_mask.shape[2] != width:
+            remove_edge_mask = F.interpolate(remove_edge_mask.unsqueeze(0), size=(height, width), mode='nearest').squeeze(0)
+            
+        if total_mask.shape[1] != height or total_mask.shape[2] != width:
+            total_mask = F.interpolate(total_mask.unsqueeze(0), size=(height, width), mode='nearest').squeeze(0)
         
-        return add_color_image_tensor, original_image_tensor, total_mask, add_edge_mask, remove_edge_mask
+        # Ensure all tensor operations detach gradients
+        return add_color_image_tensor.detach(), original_image_tensor.detach(), total_mask.detach(), add_edge_mask.detach(), remove_edge_mask.detach()
 
     @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):
-        print(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)
-        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)
+    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}")
+        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 == "":
@@ -208,9 +502,18 @@ class MagicQuill(object):
         print("positive prompt: ", positive_prompt)
         latent_samples, final_image, lineart_output, color_output = cls.scribbleColorEditModel.process(model, vae, clip, original_image, add_color_image, base_model_version, positive_prompt, negative_prompt, dtype, total_mask, add_edge_mask, remove_edge_mask, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler)
 
+        # Ensure all data is serializable before sending via JSON
+        # Convert tensor to base64 string
         final_image_base64 = tensor_to_base64(final_image)
+        
+        # Get the string representation of the image path if it's not a tensor
+        image_name = image
+        if isinstance(image, torch.Tensor):
+            image_name = "generated_image"
+            
+        # Send the serializable data
         PromptServer.instance.send_sync(
-            "magic_quill/final_image", {"image": final_image_base64, "image_name": image}
+            "magic_quill/final_image", {"image": final_image_base64, "image_name": image_name}
         )
         
         return (latent_samples, final_image, lineart_output, color_output)
@@ -225,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,35 +1,77 @@
+[tool.poetry]
+name = "ComfyUI-MagicQuill"
+version = "1.0.4"
+description = "Fixed version of the original MagicQuill node."
+authors = ["brantje <brantje@gmail.com>"]
+license = { text = "MIT License" }
+readme = "README.md"
+
+[tool.poetry.dependencies]
+
+pynvml = "^11.4.0"
+
+
 [build-system]
-requires = ["setuptools>=61.0"]
-build-backend = "setuptools.build_meta"
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
 
 [project]
-name = "llava"
-version = "1.2.2.post1"
-description = "Towards GPT-4 like large language and visual assistant."
-readme = "README.md"
-requires-python = ">=3.8"
-classifiers = [
-    "Programming Language :: Python :: 3",
-    "License :: OSI Approved :: Apache Software License",
-]
-dependencies = [
-    "transformers==4.37.2", "tokenizers==0.15.1", "sentencepiece==0.2.0", "shortuuid",
-    "accelerate==0.33.0", "peft", "bitsandbytes",
-    "pydantic", "markdown2[all]", "numpy", "scikit-learn==1.2.2",
-    "gradio==5.4.0", "gradio_client==1.4.2", "requests", "httpx==0.24.1", "uvicorn", "fastapi",
-    "einops==0.6.1", "einops-exts==0.0.4", "timm==0.6.13",
-]
-
-[project.optional-dependencies]
-train = ["deepspeed==0.12.6", "ninja", "wandb"]
-build = ["build", "twine"]
+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.4" # Custom Node version. Must be semantically versioned.
+dependencies  = [
+    'opencv-python',
+    'diffusers',
+    'torchsde',
+    'protobuf',
+    'transformers',
+    'tokenizers',
+    'sentencepiece',
+    'shortuuid',
+    'accelerate',
+    'peft',
+    'bitsandbytes',
+    'pydantic',
+    'markdown2',
+    'scikit-learn',
+    'requests',
+    'httpx',
+    'uvicorn',
+    'fastapi',
+    'einops',
+    'einops-exts',
+    'timm',
+    'webcolors',
+    'torch',
+    'importlib_metadata',
+    'huggingface_hub',
+    'scipy',
+    'opencv-python',
+    'filelock',
+    'numpy',
+    'Pillow',
+    'einops',
+    'torchvision',
+    'pyyaml',
+    'scikit-image',
+    'python-dateutil',
+    'mediapipe',
+    'svglib',
+    'fvcore',
+    'yapf',
+    'omegaconf',
+    'ftfy',
+    'addict',
+    'yacs',
+    'trimesh[easy]',
+    'albumentations',
+    'scikit-learn',
+    'matplotlib'
+] # Filled in from requirements.txt
 
 [project.urls]
-"Homepage" = "https://llava-vl.github.io"
-"Bug Tracker" = "https://github.com/haotian-liu/LLaVA/issues"
+Repository = "https://github.com/brantje/ComfyUI_MagicQuill"
 
-[tool.setuptools.packages.find]
-exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]
-
-[tool.wheel]
-exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]
+[tool.comfy]
+PublisherId = "brantje"
+DisplayName = "ComfyUI MagicQuill (fixed)"

requirements.txt

@@ -1,9 +1,8 @@
-webcolors==1.13
 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
@@ -20,7 +19,7 @@ fastapi
 einops==0.6.1
 einops-exts==0.0.4
 timm==0.6.13
-
+webcolors==24.11.1
 torch
 importlib_metadata
 huggingface_hub

scribble_color_edit.py

@@ -2,14 +2,48 @@ import os
 import torch.nn.functional as F
 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
@@ -46,6 +80,32 @@ class ScribbleColorEditModel():
         print("self.brushnet_loader.inpaint_files: ", get_files_with_extension('inpaint'))
         self.brushnet = self.brushnet_loader.brushnet_loading(brushnet_name, dtype)[0]
 
+    def safe_vae_decode(self, vae, latent_samples):
+        """Safe VAE decoding that handles inference tensors correctly."""
+        # First, ensure the latent samples are on CPU and detached
+        samples = latent_samples["samples"].to(device="cpu").detach().clone()
+        
+        # Convert to standard float format
+        samples = samples.float()
+        
+        # Disable gradient tracking for this operation
+        with torch.no_grad():
+            # Create a fresh dictionary with the cloned tensor
+            latent_dict = {"samples": samples}
+            try:
+                # Decode using the VAE
+                return self.vae_decoder.decode(vae, latent_dict)
+            except RuntimeError as e:
+                # If we still encounter an error, try a deeper copy approach
+                print(f"First VAE decode attempt failed: {str(e)}")
+                
+                # Create completely fresh tensors by serializing and deserializing
+                serialized = pytree.tree_map(lambda x: x.detach().cpu().numpy() if isinstance(x, torch.Tensor) else x, latent_dict)
+                deserialized = pytree.tree_map(lambda x: torch.tensor(x) if isinstance(x, np.ndarray) else x, serialized)
+                
+                # Try decoding again with the completely new tensors
+                return self.vae_decoder.decode(vae, deserialized)
+
     def process(self, model, vae, clip, image, colored_image, base_model_version, positive_prompt, negative_prompt, dtype, mask, add_mask, remove_mask, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
         print("mask.shape", mask.shape)
         print("image.shape", image.shape)
@@ -115,8 +175,47 @@ class ScribbleColorEditModel():
             latent_image=latent,
         )[0]
 
-        final_image = self.vae_decoder.decode(vae, latent_samples)[0]
-        final_image = self.blender.blend_inpaint(final_image, image, mask, kernel=10, sigma=10.0)[0]
-
+        # Use the safe VAE decode method instead of direct decoding
+        final_image = self.safe_vae_decode(vae, latent_samples)[0]
+        
+        # Ensure image dimensions match before blending (handle RGB vs RGBA)
+        if final_image.shape[-1] != image.shape[-1]:
+            print(f"Dimension mismatch: final_image shape: {final_image.shape}, image shape: {image.shape}")
+            
+            # Handle different dimensions properly
+            # First, make sure we understand the tensor shapes
+            print(f"final_image.dim() = {final_image.dim()}, image.dim() = {image.dim()}")
+            
+            # Convert both to 3D (H,W,C) format if they're not already
+            if final_image.dim() == 4:
+                final_image = final_image.squeeze(0)  # Remove batch dimension if present
+            if image.dim() == 4:
+                image = image.squeeze(0)  # Remove batch dimension if present
+                
+            # Now handle channel differences
+            if final_image.shape[-1] == 3 and image.shape[-1] == 4:
+                # Add an alpha channel (fully opaque) to final_image
+                alpha_channel = torch.ones((final_image.shape[0], final_image.shape[1], 1), 
+                                          device=final_image.device, dtype=final_image.dtype)
+                final_image = torch.cat([final_image, alpha_channel], dim=-1)
+            elif final_image.shape[-1] == 4 and image.shape[-1] == 3:
+                # Use only RGB channels from final_image or add alpha to image
+                final_image = final_image[..., :3]
+            else:
+                # Just use the first 3 channels for both
+                final_image = final_image[..., :3]
+                if image.shape[-1] > 3:
+                    image = image[..., :3]
+        
+        # Print shape information before blending
+        print(f"Before blending - final_image shape: {final_image.shape}, image shape: {image.shape}, mask shape: {mask.shape}")
+        
+        # Make sure mask has the right dimensions for blending
+        if mask.dim() == 3 and mask.shape[0] == 1:  # If mask is [1, H, W]
+            mask_for_blend = mask.squeeze(0)        # Convert to [H, W]
+        else:
+            mask_for_blend = mask
+            
+        final_image = self.blender.blend_inpaint(final_image, image, mask_for_blend, kernel=10, sigma=10.0)[0]
 
         return (latent_samples, final_image, lineart_output, color_output)

utils.py

@@ -34,8 +34,8 @@ def get_colored_contour(img1, img2, threshold=10):
 
 def closest_colour(requested_colour):
     min_colours = {}
-    for key, name in webcolors.CSS3_HEX_TO_NAMES.items():
-        r_c, g_c, b_c = webcolors.hex_to_rgb(key)
+    for name in webcolors.names("css3"):
+        r_c, g_c, b_c = webcolors.name_to_rgb(name)
         rd = (r_c - requested_colour[0].item()) ** 2
         gd = (g_c - requested_colour[1].item()) ** 2
         bd = (b_c - requested_colour[2].item()) ** 2