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

.github/workflows/main.yml

@@ -1,20 +0,0 @@
-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,19 +4,6 @@
 <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

@@ -1,10 +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)

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"].includes(w.name)) {
+          } else if (["image", "original_image", "add_color_image", "add_edge_image", "remove_edge_image", "positive_prompt"].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","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"]
+    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"]
     if (!this.flags.setting_collapsed) {
       for (const w of this.widgets) {
         if (setting_widgets.includes(w.name)) {

llava_new.py

@@ -23,30 +23,15 @@ import re
 class LLaVAModel:
     def __init__(self):
         # replace the model_path with correct path folder
-        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):
+        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.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
-            model_path=self.model_path,
+            model_path=model_path,
             model_base=None,
-            model_name=get_model_name_from_path(self.model_path),
+            model_name=get_model_name_from_path(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

@@ -284,6 +284,14 @@ async def run_magic_quill(request):
         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
@@ -310,7 +318,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
@@ -380,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"}),
@@ -397,18 +405,9 @@ 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",),
+             
             }
         }
 
@@ -421,7 +420,6 @@ 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]
@@ -454,42 +452,20 @@ 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, 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):
+    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(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)
+        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 == "":
@@ -528,5 +504,9 @@ 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, optional_image = None):
+    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)
+
         return True

pyproject.toml

@@ -1,77 +1,35 @@
-[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 = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
 
 [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.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
+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"]
 
 [project.urls]
-Repository = "https://github.com/brantje/ComfyUI_MagicQuill"
+"Homepage" = "https://llava-vl.github.io"
+"Bug Tracker" = "https://github.com/haotian-liu/LLaVA/issues"
 
-[tool.comfy]
-PublisherId = "brantje"
-DisplayName = "ComfyUI MagicQuill (fixed)"
+[tool.setuptools.packages.find]
+exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]
+
+[tool.wheel]
+exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]

requirements.txt

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

scribble_color_edit.py

@@ -4,51 +4,23 @@ 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)
-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, '..')))
 sys.path.append(os.path.abspath(os.path.join(current_dir, '..', '..', 'comfy_extras')))
+print(sys.path)
 
-
-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_BrushNet.brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint, get_files_with_extension
 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():
@@ -68,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."""
@@ -124,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]
@@ -148,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,