Remove repository cloning code from install.py and update README.md to include notes on required repositories.

.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

@@ -1,41 +1,10 @@
 import os
 import subprocess
-from huggingface_hub import snapshot_download, hf_hub_download
+from huggingface_hub import snapshot_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)
 
-
-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)
+# 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)

magic_quill.py

@@ -284,14 +284,6 @@ 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
@@ -388,7 +380,7 @@ class MagicQuill(object):
                 "clip": ("CLIP",),
                 "vae": ("VAE",),
                 
-                "base_model_version": (['SD1.5', 'FLUX'], {"default": "SD1.5"}),
+                "base_model_version": (['SD1.5'], {"default": "SD1.5"}),
                 "positive_prompt": ("STRING", {"default": ""}),
                 "negative_prompt": ("STRING", {"default": ""}),
                 "dtype": (['float16', 'bfloat16', 'float32', 'float64'], {"default": "float16"}),

pyproject.toml

@@ -1,35 +1,77 @@
+[tool.poetry]
+name = "ComfyUI-MagicQuill"
+version = "1.0.0"
+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.0" # 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)"

scribble_color_edit.py

@@ -16,11 +16,7 @@ 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 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 import ControlNetLoader, ControlNetApplyAdvanced, CLIPTextEncode, KSampler, VAEDecode
 from nodes_mask import GrowMask
 
 class ScribbleColorEditModel():
@@ -40,36 +36,17 @@ 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]
-        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
+        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]
 
     def safe_vae_decode(self, vae, latent_samples):
         """Safe VAE decoding that handles inference tensors correctly."""
@@ -115,8 +92,6 @@ 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]
@@ -141,39 +116,20 @@ class ScribbleColorEditModel():
                 lineart_output[bool_add_mask_resized] = 1.0
             else:
                 lineart_output[bool_remove_mask_resized & ~bool_add_mask_resized] = 0.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
-                )
+            positive, negative = self.controlnet_apply.apply_controlnet(positive, negative, self.edge_controlnet, lineart_output, edge_strength, 0.0, 1.0)
 
+        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,