Partly SDXL working

install.py

@@ -11,11 +11,25 @@ hugging_face_repos = [
         'repo_id': 'LiuZichen/MagicQuill-models',
         'local_dir': models_dir
     },
+
+    # SDXL
     {
-        'repo_id': 'InstantX/FLUX.1-dev-Controlnet-Union',
-        'local_dir': os.path.join(models_dir, 'controlnet', 'FLUX1'),
+        'repo_id': 'xinsir/controlnet-scribble-sdxl-1.0',
+        'local_dir': os.path.join(models_dir, 'controlnet', 'SDXL'),
         'filename': 'diffusion_pytorch_model.safetensors',
-        'target_filename': 'FLUX.1-dev-Controlnet-Union.safetensors'
+        'target_filename': 'sdxl-controlnet-scribble-1.0.safetensors'
+    },
+    {
+        'repo_id': 'briaai/BRIA-2.3-ControlNet-ColorGrid',
+        'local_dir': os.path.join(models_dir, 'controlnet', 'SDXL'),
+        'filename': 'bria_23_controlnet_colorgrid.safetensors',
+        'target_filename': 'bria_23_controlnet_colorgrid.safetensors'
+    },
+    {
+        'repo_id': 'fullnonstop/random_mask_brushnet_ckpt_sdxl_v0',
+        'local_dir': os.path.join(models_dir, 'inpaint', 'SDXL'),
+        'filename': 'random_mask_brushnet_ckpt_sdxl_v0.safetensors',
+        'target_filename': 'random_mask_brushnet_ckpt_sdxl_v0.safetensors'
     }
 ]
 

magic_quill.py

@@ -284,14 +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}")
+        if "SDXL" in checkpoint_name:
+            base_model_version = "SDXL"
+        elif "SD1.5" in checkpoint_name:
+            base_model_version = "SD1.5"
+        else:
+            raise ValueError(f"Could not determine base model version from checkpoint name: {checkpoint_name} use base_model_version parameter")
+
         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 +388,7 @@ class MagicQuill(object):
                 "clip": ("CLIP",),
                 "vae": ("VAE",),
                 
-                "base_model_version": (['SD1.5', 'FLUX'], {"default": "SD1.5"}),
+                "base_model_version": (['SD1.5', 'SDXL'], {"default": "SD1.5"}),
                 "positive_prompt": ("STRING", {"default": ""}),
                 "negative_prompt": ("STRING", {"default": ""}),
                 "dtype": (['float16', 'bfloat16', 'float32', 'float64'], {"default": "float16"}),

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,23 @@ 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}")
-
+            edge_controlnet_name = "SDXL/sdxl-controlnet-scribble-1.0.safetensors"
+            color_controlnet_name = "SDXL/bria_23_controlnet_colorgrid.safetensors"
+            brushnet_name = os.path.join("SDXL", "random_mask_brushnet_ckpt_sdxl_v0.safetensors")
         else:
             raise ValueError("Invalid base_model_version, not supported yet!!!: {}".format(base_model_version))
+        print(f"Loading edge controlnet: {edge_controlnet_name} color controlnet: {color_controlnet_name} brushnet: {brushnet_name}")
+        
         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 +98,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 +122,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,