Update edge_controlnet_name path for FLUX model in ScribbleColorEditModel to reflect correct versioning format.

Api call

install.py

@@ -1,20 +1,41 @@
 import os
 import subprocess
-from huggingface_hub import snapshot_download
+from huggingface_hub import snapshot_download, hf_hub_download
 parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 models_dir = os.path.abspath(os.path.join(parent_dir, '..', "models"))
 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)
+
+hugging_face_repos = [
+    {
+        'repo_id': 'LiuZichen/MagicQuill-models',
+        'local_dir': models_dir
+    },
+    {
+        'repo_id': 'InstantX/FLUX.1-dev-Controlnet-Union',
+        'local_dir': os.path.join(models_dir, 'controlnet', 'FLUX1'),
+        'filename': 'diffusion_pytorch_model.safetensors',
+        'target_filename': 'FLUX.1-dev-Controlnet-Union.safetensors'
+    }
+]
+
+for repo in hugging_face_repos:
+    if 'filename' in repo:
+       filename = hf_hub_download(repo_id=repo['repo_id'], local_dir=repo['local_dir'], filename=repo['filename'])
+       current_path = os.path.join(repo['local_dir'], filename)
+       target_path = os.path.join(repo['local_dir'], repo['target_filename'])
+       if os.path.exists(current_path):
+         if current_path != target_path:
+             os.rename(current_path, target_path)
+    else:
+        snapshot_download(repo_id=repo['repo_id'], local_dir=repo['local_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)
+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)

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
@@ -89,25 +90,92 @@ 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()
 
-    return web.json_response({"prompt": res, "error": False})
+    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):
@@ -194,6 +262,8 @@ async def run_magic_quill(request):
         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)
             
@@ -207,12 +277,21 @@ async def run_magic_quill(request):
         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", 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")
 
+        if "FLUX" in checkpoint_name:
+            base_model_version = "FLUX"
+        elif "SDXL" in checkpoint_name:
+            base_model_version = "SDXL"
+        else:
+            base_model_version = "SD1.5"
+
+        print(f"Base model version: {base_model_version} checkpoint_name: {checkpoint_name}")
         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
@@ -309,7 +388,7 @@ class MagicQuill(object):
                 "clip": ("CLIP",),
                 "vae": ("VAE",),
                 
-                "base_model_version": (['SD1.5'], {"default": "SD1.5"}),
+                "base_model_version": (['SD1.5', 'FLUX'], {"default": "SD1.5"}),
                 "positive_prompt": ("STRING", {"default": ""}),
                 "negative_prompt": ("STRING", {"default": ""}),
                 "dtype": (['float16', 'bfloat16', 'float32', 'float64'], {"default": "float16"}),

requirements.txt

@@ -1,4 +1,3 @@
-webcolors==1.13
 opencv-python
 diffusers
 torchsde==0.2.6
@@ -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

@@ -16,7 +16,11 @@ 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
 
-from nodes import ControlNetLoader, ControlNetApplyAdvanced, CLIPTextEncode, KSampler, VAEDecode
+from comfy_extras.nodes_controlnet import SetUnionControlNetType
+from comfy_extras.nodes_differential_diffusion import DifferentialDiffusion
+from comfy_extras.nodes_flux import FluxGuidance
+
+from nodes import ControlNetLoader, ControlNetApplyAdvanced, CLIPTextEncode, KSampler, VAEDecode, InpaintModelConditioning
 from nodes_mask import GrowMask
 
 class ScribbleColorEditModel():
@@ -36,17 +40,36 @@ class ScribbleColorEditModel():
         # self.load_models('SD1.5', 'float16')
 
     def load_models(self, base_model_version, dtype):
+        print(f"Loading models for base_model_version: {base_model_version}")
         if base_model_version == "SD1.5":
             edge_controlnet_name = "control_v11p_sd15_scribble.safetensors"
             color_controlnet_name = "color_finetune.safetensors"
             brushnet_name = os.path.join("brushnet", "random_mask_brushnet_ckpt", "diffusion_pytorch_model.safetensors")
+            print(f"Loading models for base_model_version: {base_model_version}")
+        elif base_model_version == "SDXL":
+            edge_controlnet_name = os.path.join("SDXL", "sd_xl_base_1.0_controlnet.safetensors")
+            color_controlnet_name = None # TODO: add color controlnet for SDXL
+            brushnet_name = None # TODO: add brushnet for SDXL
+            print(f"Loading models for SDXL base_model_version: {base_model_version}")
+        elif base_model_version == "FLUX":
+            edge_controlnet_name = os.path.join("FLUX.1", "Shakker-Labs-ControlNet-Union-Pro","diffusion_pytorch_model.safetensors")
+            color_controlnet_name = None # TODO: add color controlnet for FLUX
+            brushnet_name = None # TODO: add brushnet for FLUX
+            print(f"Loading models for Flux base_model_version: {base_model_version}")
+
         else:
             raise ValueError("Invalid base_model_version, not supported yet!!!: {}".format(base_model_version))
         self.edge_controlnet = self.controlnet_loader.load_controlnet(edge_controlnet_name)[0]
-        self.color_controlnet = self.controlnet_loader.load_controlnet(color_controlnet_name)[0]
-        self.brushnet_loader.inpaint_files = get_files_with_extension('inpaint')
-        print("self.brushnet_loader.inpaint_files: ", get_files_with_extension('inpaint'))
-        self.brushnet = self.brushnet_loader.brushnet_loading(brushnet_name, dtype)[0]
+        if color_controlnet_name:
+            self.color_controlnet = self.controlnet_loader.load_controlnet(color_controlnet_name)[0]
+        else:
+            self.color_controlnet = None
+        if brushnet_name:
+            self.brushnet_loader.inpaint_files = get_files_with_extension('inpaint')
+            print("self.brushnet_loader.inpaint_files: ", get_files_with_extension('inpaint'))
+            self.brushnet = self.brushnet_loader.brushnet_loading(brushnet_name, dtype)[0]
+        else:
+            self.brushnet = None
 
     def safe_vae_decode(self, vae, latent_samples):
         """Safe VAE decoding that handles inference tensors correctly."""
@@ -92,6 +115,8 @@ class ScribbleColorEditModel():
                 image_copy[bool_add_mask] = 1.0
 
         if not torch.equal(image, colored_image):
+            if base_model_version == "FLUX":
+                raise ValueError('Not implemented.')            
             print("Apply color controlnet")
             color_output = self.color_processor.execute(colored_image, resolution=2048)[0]
             lineart_output = self.lineart_processor.execute(image, resolution=512, coarse=False)[0]
@@ -116,20 +141,39 @@ class ScribbleColorEditModel():
                 lineart_output[bool_add_mask_resized] = 1.0
             else:
                 lineart_output[bool_remove_mask_resized & ~bool_add_mask_resized] = 0.0
-            positive, negative = self.controlnet_apply.apply_controlnet(positive, negative, self.edge_controlnet, lineart_output, edge_strength, 0.0, 1.0)
+           
+            if base_model_version == "FLUX":
+
+                self.edge_controlnet = (SetUnionControlNetType().set_controlnet_type(self.edge_controlnet, 2))[0] # set union type to hed/pidi/scribble/ted
+                model = (DifferentialDiffusion().apply(model))[0] # apply Differential Diffusion
+                positive = (FluxGuidance().append(positive, 30))[0] # apply flux guidence
+            positive, negative = self.controlnet_apply.apply_controlnet(positive, negative, self.edge_controlnet, lineart_output, edge_strength, 0.0, 1.0, vae)
+            
+            if base_model_version == "FLUX":
+                positive, negative, latent = InpaintModelConditioning().encode(positive, negative, image, vae, mask, False) # apply inpaint
+           
+            # if base_model_version == "SD1.5":
+            #     positive, negative = self.controlnet_apply.apply_controlnet(positive, negative, self.edge_controlnet, lineart_output, edge_strength, 0.0, 1.0)
+            # else:
+            #     positive, negative = self.controlnet_apply.apply_controlnet(positive, negative, self.edge_controlnet, lineart_output, edge_strength, vae, 0.0, 1.0)
+        
+        
+            if base_model_version == "FLUX":
+                positive, negative, latent = InpaintModelConditioning().encode(positive, negative, image, vae, mask, False) # apply inpaint
+            else:
+                model, positive, negative, latent = self.brushnet_node.model_update(
+                    model=model,
+                    vae=vae,
+                    image=image,
+                    mask=mask,
+                    brushnet=self.brushnet,
+                    positive=positive,
+                    negative=negative,
+                    scale=inpaint_strength,
+                    start_at=0,
+                    end_at=10000
+                )
 
-        model, positive, negative, latent = self.brushnet_node.model_update(
-            model=model,
-            vae=vae,
-            image=image,
-            mask=mask,
-            brushnet=self.brushnet,
-            positive=positive,
-            negative=negative,
-            scale=inpaint_strength,
-            start_at=0,
-            end_at=10000
-        )
 
         latent_samples = self.ksampler.sample(
             model=model, 

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