error CUDA out of memory. Tried to allocate 146.00 MiB.

MagicQuill/brushnet/powerpaint_utils.py

@@ -51,27 +51,7 @@ class TokenizerWrapper:
         Args:
             tokens (Union[str, List[str]]): The tokens to be added.
         """
-        # Check if tokens exist first to avoid assertion error in wrapped tokenizer
-        # and to ensure idempotency in shared environments.
-        if isinstance(tokens, str):
-            tokens_to_check = [tokens]
-        else:
-            tokens_to_check = tokens
-
-        # If all tokens are already in the vocabulary, skip adding them.
-        # This relies on the wrapped tokenizer's behavior or checking its vocab.
-        # Usually `add_tokens` returns 0 if all tokens exist.
-        # We just want to avoid the assertion error if they are already added.
-        
         num_added_tokens = self.wrapped.add_tokens(tokens, *args, **kwargs)
-        if num_added_tokens == 0:
-            # Check if they actually exist (idempotency case)
-            # If they exist, we don't assert error, just return.
-            # If they don't exist but add_tokens returned 0 (shouldn't happen for new tokens),
-            # then we might have an issue.
-            # For simplicity in fixing the leak/crash: if 0 added, assume they exist.
-            return
-
         assert num_added_tokens != 0, (
             f"The tokenizer already contains the token {tokens}. Please pass "
             "a different `placeholder_token` that is not already in the "
@@ -102,11 +82,6 @@ class TokenizerWrapper:
                 the added placeholder token.
             *args, **kwargs: The arguments for `self.wrapped.add_tokens`.
         """
-        
-        # Check if already in token_map (idempotency)
-        if placeholder_token in self.token_map:
-             return
-             
         output = []
         if num_vec_per_token == 1:
             self.try_adding_tokens(placeholder_token, *args, **kwargs)
@@ -301,29 +276,55 @@ class EmbeddingLayerWithFixes(nn.Module):
 
     def add_embeddings(self, embeddings: Optional[Union[dict, List[dict]]]):
         """Add external embeddings to this layer.
+
+        Use case:
+
+        >>> 1. Add token to tokenizer and get the token id.
+        >>> tokenizer = TokenizerWrapper('openai/clip-vit-base-patch32')
+        >>> # 'how much' in kiswahili
+        >>> tokenizer.add_placeholder_tokens('ngapi', num_vec_per_token=4)
+        >>>
+        >>> 2. Add external embeddings to the model.
+        >>> new_embedding = {
+        >>>     'name': 'ngapi',  # 'how much' in kiswahili
+        >>>     'embedding': torch.ones(1, 15) * 4,
+        >>>     'start': tokenizer.get_token_info('kwaheri')['start'],
+        >>>     'end': tokenizer.get_token_info('kwaheri')['end'],
+        >>>     'trainable': False  # if True, will registry as a parameter
+        >>> }
+        >>> embedding_layer = nn.Embedding(10, 15)
+        >>> embedding_layer_wrapper = EmbeddingLayerWithFixes(embedding_layer)
+        >>> embedding_layer_wrapper.add_embeddings(new_embedding)
+        >>>
+        >>> 3. Forward tokenizer and embedding layer!
+        >>> input_text = ['hello, ngapi!', 'hello my friend, ngapi?']
+        >>> input_ids = tokenizer(
+        >>>     input_text, padding='max_length', truncation=True,
+        >>>     return_tensors='pt')['input_ids']
+        >>> out_feat = embedding_layer_wrapper(input_ids)
+        >>>
+        >>> 4. Let's validate the result!
+        >>> assert (out_feat[0, 3: 7] == 2.3).all()
+        >>> assert (out_feat[2, 5: 9] == 2.3).all()
+
+        Args:
+            embeddings (Union[dict, list[dict]]): The external embeddings to
+                be added. Each dict must contain the following 4 fields: 'name'
+                (the name of this embedding), 'embedding' (the embedding
+                tensor), 'start' (the start token id of this embedding), 'end'
+                (the end token id of this embedding). For example:
+                `{name: NAME, start: START, end: END, embedding: torch.Tensor}`
         """
         if isinstance(embeddings, dict):
             embeddings = [embeddings]
 
-        # Idempotency check: filter out embeddings that are already present by name
-        existing_names = {emb["name"] for emb in self.external_embeddings}
-        new_embeddings = []
-        for emb in embeddings:
-            if emb["name"] not in existing_names:
-                new_embeddings.append(emb)
-            # Optional: Warn or check if existing embedding matches the new one?
-            # For now, assume if name exists, it's the same token being re-added.
-
-        if not new_embeddings:
-            return
-
-        self.external_embeddings += new_embeddings
+        self.external_embeddings += embeddings
         self.check_duplicate_names(self.external_embeddings)
         self.check_ids_overlap(self.external_embeddings)
 
         # set for trainable
         added_trainable_emb_info = []
-        for embedding in new_embeddings:
+        for embedding in embeddings:
             trainable = embedding.get("trainable", False)
             if trainable:
                 name = embedding["name"]
@@ -331,7 +332,7 @@ class EmbeddingLayerWithFixes(nn.Module):
                 self.trainable_embeddings[name] = embedding["embedding"]
                 added_trainable_emb_info.append(name)
 
-        added_emb_info = [emb["name"] for emb in new_embeddings]
+        added_emb_info = [emb["name"] for emb in embeddings]
         added_emb_info = ", ".join(added_emb_info)
         print(f"Successfully add external embeddings: {added_emb_info}.", "current")
 
@@ -459,39 +460,17 @@ def add_tokens(
         assert len(initialize_tokens) == len(
             placeholder_tokens
         ), "placeholder_token should be the same length as initialize_token"
-    
-    # Safe to call multiple times (idempotent)
     for ii in range(len(placeholder_tokens)):
         tokenizer.add_placeholder_token(placeholder_tokens[ii], num_vec_per_token=num_vectors_per_token)
 
     # text_encoder.set_embedding_layer()
     embedding_layer = text_encoder.text_model.embeddings.token_embedding
-    if not isinstance(embedding_layer, EmbeddingLayerWithFixes):
-        text_encoder.text_model.embeddings.token_embedding = EmbeddingLayerWithFixes(embedding_layer)
-        embedding_layer = text_encoder.text_model.embeddings.token_embedding
+    text_encoder.text_model.embeddings.token_embedding = EmbeddingLayerWithFixes(embedding_layer)
+    embedding_layer = text_encoder.text_model.embeddings.token_embedding
 
     assert embedding_layer is not None, (
         "Do not support get embedding layer for current text encoder. " "Please check your configuration."
     )
-    
-    # Only calculate initialization for tokens that are NOT already in the layer
-    existing_names = {emb["name"] for emb in embedding_layer.external_embeddings}
-    tokens_to_add = []
-    init_tokens_to_add = []
-    
-    for ii, token in enumerate(placeholder_tokens):
-        # This check assumes the placeholder token name matches the embedding name
-        # TokenizerWrapper adds suffix _0, _1 etc if num_vec > 1. 
-        # The logic below handles generic case, but here we assume 1-to-1 or we check the main token.
-        # Actually EmbeddingLayer uses specific names. TokenizerWrapper.add_placeholder_token generates them.
-        # If num_vec_per_token > 1, TokenizerWrapper generates token_0, token_1...
-        # Let's check if the embedding layer already has them.
-        
-        # The original code below generated embeddings for ALL input tokens.
-        # add_embeddings will filter them out.
-        # But we need to be careful not to re-initialize them differently if they exist.
-        pass
-
     initialize_embedding = []
     if initialize_tokens is not None:
         for ii in range(len(placeholder_tokens)):
@@ -510,12 +489,8 @@ def add_tokens(
 
     token_info_all = []
     for ii in range(len(placeholder_tokens)):
-        # get_token_info relies on the token being in tokenizer. 
-        # add_placeholder_token ensures it's there (idempotent now).
         token_info = tokenizer.get_token_info(placeholder_tokens[ii])
         token_info["embedding"] = initialize_embedding[ii]
         token_info["trainable"] = True
         token_info_all.append(token_info)
-    
-    # Idempotency is handled inside add_embeddings now
     embedding_layer.add_embeddings(token_info_all)

MagicQuill/brushnet_nodes.py

@@ -149,13 +149,7 @@ class PowerPaintCLIPLoader:
 
         print('PowerPaint base CLIP file: ', base_CLIP_file)
 
-        # Reuse TokenizerWrapper if already wrapped
-        if isinstance(pp_clip.tokenizer.clip_l.tokenizer, TokenizerWrapper):
-            pp_tokenizer = pp_clip.tokenizer.clip_l.tokenizer
-        else:
-            pp_tokenizer = TokenizerWrapper(pp_clip.tokenizer.clip_l.tokenizer)
-            pp_clip.tokenizer.clip_l.tokenizer = pp_tokenizer
-
+        pp_tokenizer = TokenizerWrapper(pp_clip.tokenizer.clip_l.tokenizer)
         pp_text_encoder = pp_clip.patcher.model.clip_l.transformer
 
         add_tokens(
@@ -170,8 +164,7 @@ class PowerPaintCLIPLoader:
 
         print('PowerPaint CLIP file: ', pp_CLIP_file)
 
-        # Already assigned above if new, or reused if existing
-        # pp_clip.tokenizer.clip_l.tokenizer = pp_tokenizer 
+        pp_clip.tokenizer.clip_l.tokenizer = pp_tokenizer
         pp_clip.patcher.model.clip_l.transformer = pp_text_encoder
 
         return (pp_clip,)
@@ -287,7 +280,7 @@ class PowerPaint:
         # unload vae and CLIPs
         del vae
         del clip
-        for loaded_model in list(comfy.model_management.current_loaded_models):
+        for loaded_model in comfy.model_management.current_loaded_models:
             if type(loaded_model.model.model) in ModelsToUnload:
                 comfy.model_management.current_loaded_models.remove(loaded_model)
                 loaded_model.model_unload()
@@ -375,7 +368,7 @@ class BrushNet:
 
         # unload vae
         del vae
-        for loaded_model in list(comfy.model_management.current_loaded_models):
+        for loaded_model in comfy.model_management.current_loaded_models:
             if type(loaded_model.model.model) in ModelsToUnload:
                 comfy.model_management.current_loaded_models.remove(loaded_model)
                 loaded_model.model_unload()
@@ -948,20 +941,6 @@ def brushnet_inference(x, timesteps, transformer_options, debug):
                     debug=debug,
                 )
 
-def forward_patched_by_brushnet(self, x, *args, **kwargs):
-    h = self.original_forward(x, *args, **kwargs)
-    if hasattr(self, 'add_sample_after') and type(self):
-        to_add = self.add_sample_after
-        if torch.is_tensor(to_add):
-            # interpolate due to RAUNet
-            if h.shape[2] != to_add.shape[2] or h.shape[3] != to_add.shape[3]:
-                to_add = torch.nn.functional.interpolate(to_add, size=(h.shape[2], h.shape[3]), mode='bicubic')                  
-            h += to_add.to(h.dtype).to(h.device)
-        else:
-            h += self.add_sample_after
-        self.add_sample_after = 0
-    return h
-
 
 # This is main patch function
 def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, 
@@ -973,7 +952,7 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents,
     is_SDXL = isinstance(model.model.model_config, comfy.supported_models.SDXL)
 
     if is_SDXL:
-        input_blocks = [[0, comfy.ops.manual_cast.Conv2d],
+        input_blocks = [[0, comfy.ops.disable_weight_init.Conv2d],
                         [1, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
                         [2, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
                         [3, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample],
@@ -995,7 +974,7 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents,
                         [7, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
                         [8, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock]]
     else:
-        input_blocks = [[0, comfy.ops.manual_cast.Conv2d],
+        input_blocks = [[0, comfy.ops.disable_weight_init.Conv2d],
                         [1, comfy.ldm.modules.attention.SpatialTransformer],
                         [2, comfy.ldm.modules.attention.SpatialTransformer],
                         [3, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample],
@@ -1080,33 +1059,36 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents,
     bo['latent_id'] = 0
 
     # patch layers `forward` so we can apply brushnet
-    # Modified to prevent re-patching leak and closure creation
+    def forward_patched_by_brushnet(self, x, *args, **kwargs):
+        h = self.original_forward(x, *args, **kwargs)
+        if hasattr(self, 'add_sample_after') and type(self):
+            to_add = self.add_sample_after
+            if torch.is_tensor(to_add):
+                # interpolate due to RAUNet
+                if h.shape[2] != to_add.shape[2] or h.shape[3] != to_add.shape[3]:
+                    to_add = torch.nn.functional.interpolate(to_add, size=(h.shape[2], h.shape[3]), mode='bicubic')                  
+                h += to_add.to(h.dtype).to(h.device)
+            else:
+                h += self.add_sample_after
+            self.add_sample_after = 0
+        return h
+
     for i, block in enumerate(model.model.diffusion_model.input_blocks):
         for j, layer in enumerate(block):
             if not hasattr(layer, 'original_forward'):
                 layer.original_forward = layer.forward
-            
-            # Only patch if not already patched by us
-            if getattr(layer.forward, '__func__', None) != forward_patched_by_brushnet:
-                 layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
-                 
+            layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
             layer.add_sample_after = 0
 
     for j, layer in enumerate(model.model.diffusion_model.middle_block):
         if not hasattr(layer, 'original_forward'):
             layer.original_forward = layer.forward
-        
-        if getattr(layer.forward, '__func__', None) != forward_patched_by_brushnet:
-             layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
-             
+        layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
         layer.add_sample_after = 0
 
     for i, block in enumerate(model.model.diffusion_model.output_blocks):
         for j, layer in enumerate(block):
             if not hasattr(layer, 'original_forward'):
                 layer.original_forward = layer.forward
-            
-            if getattr(layer.forward, '__func__', None) != forward_patched_by_brushnet:
-                 layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
-                 
-            layer.add_sample_after = 0
+            layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
+            layer.add_sample_after = 0

MagicQuill/comfy/cli_args.py

@@ -59,7 +59,7 @@ fp_group.add_argument("--force-fp16", action="store_true", help="Force fp16.")
 fpunet_group = parser.add_mutually_exclusive_group()
 fpunet_group.add_argument("--bf16-unet", action="store_true", help="Run the UNET in bf16. This should only be used for testing stuff.")
 fpunet_group.add_argument("--fp16-unet", action="store_true", help="Store unet weights in fp16.")
-fpunet_group.add_argument("--fp8_e4m3fn-unet", type=bool, default=True, help="Store unet weights in fp8_e4m3fn.")
+fpunet_group.add_argument("--fp8_e4m3fn-unet", action="store_true", help="Store unet weights in fp8_e4m3fn.")
 fpunet_group.add_argument("--fp8_e5m2-unet", action="store_true", help="Store unet weights in fp8_e5m2.")
 
 fpvae_group = parser.add_mutually_exclusive_group()

MagicQuill/scribble_color_edit.py

@@ -41,8 +41,7 @@ class ScribbleColorEditModel():
         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]
-
-    @torch.inference_mode()
+    
     def process(self, ckpt_name, image, colored_image, positive_prompt, negative_prompt, mask, add_mask, remove_mask, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler, base_model_version='SD1.5', dtype='float16', palette_resolution=2048):
         if ckpt_name != self.ckpt_name:
             self.ckpt_name = ckpt_name

app.py

@@ -5,7 +5,6 @@ subprocess.run(
         "pip install ./gradio_magicquill-0.0.1-py3-none-any.whl"
     )
 )
-import spaces
 import gradio as gr
 from gradio_magicquill import MagicQuill
 import random
@@ -14,7 +13,8 @@ import numpy as np
 from PIL import Image, ImageOps
 import base64
 import io
-from fastapi import Request
+from fastapi import FastAPI, Request
+import uvicorn
 from MagicQuill import folder_paths
 from MagicQuill.scribble_color_edit import ScribbleColorEditModel
 from gradio_client import Client, handle_file
@@ -23,18 +23,11 @@ import tempfile
 import cv2
 import os
 import requests
-import gc
 
 snapshot_download(repo_id="LiuZichen/MagicQuill-models", repo_type="model", local_dir="models")
 # HF_TOKEN = os.environ.get("HF_TOKEN")
 # The client has been made public. Welcome to duplicate our repo.
-_HELPER_SPACE = "LiuZichen/MagicQuillHelper"
-_client = None
-def client():
-    global _client
-    if _client is None:
-        _client = Client(_HELPER_SPACE)
-    return _client
+client = Client("LiuZichen/DrawNGuess")
 scribbleColorEditModel = ScribbleColorEditModel()
 
 def tensor_to_numpy(tensor):
@@ -136,11 +129,10 @@ def guess_prompt_handler(original_image, add_color_image, add_edge_image):
     add_color_image_file.close()
     add_edge_mask_file.close()
     
-    res = client().predict(
+    res = client.predict(
         handle_file(original_image_file.name),
         handle_file(add_color_image_file.name),
-        handle_file(add_edge_mask_file.name),
-        api_name="/guess_prompt"
+        handle_file(add_edge_mask_file.name)
     )
     
     if original_image_file and os.path.exists(original_image_file.name):
@@ -185,20 +177,9 @@ def generate(ckpt_name, total_mask, original_image, add_color_image, add_edge_im
     )
 
     final_image_base64 = tensor_to_base64(final_image)
-    
-    del latent_samples, final_image, lineart_output, color_output
-    gc.collect()
-    torch.cuda.empty_cache()
-    
     return final_image_base64
 
-
-from MagicQuill.comfy import model_management as _mm
-
-
-@spaces.GPU(duration=20)
 def generate_image_handler(x, ckpt_name, negative_prompt, fine_edge, grow_size, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
-    _mm.unload_all_models()
     if seed == -1:
         seed = random.randint(0, 2**32 - 1)
     ms_data = x['from_frontend']
@@ -225,7 +206,6 @@ with gr.Blocks(css=css, head=head) as demo:
         """
         # Welcome to MagicQuill! The paper has been accepted to CVPR 2025.
         Click the [link](https://magicquill.art) to view our demo and tutorial. Give us a [GitHub star](https://github.com/magic-quill/magicquill) if you are interested. 
-        MagicQuillV2 is available!!! Check our [demo](https://magicquill.art/v2/).
         """)
     with gr.Row(elem_classes="row"):
         ms = MagicQuill()
@@ -234,12 +214,11 @@ with gr.Blocks(css=css, head=head) as demo:
             btn = gr.Button("Run", variant="primary")
         with gr.Column():
             with gr.Accordion("parameters", open=False):
-                ckpt_value = os.path.join('SD1.5', 'realisticVisionV60B1_v51VAE.safetensors')
                 ckpt_name = gr.Dropdown(
-                    label="Base Model (fixed for demo)",
-                    choices=[ckpt_value],
-                    value=ckpt_value,
-                    interactive=False
+                    label="Base Model Name",
+                    choices=folder_paths.get_filename_list("checkpoints"),
+                    value='SD1.5/realisticVisionV60B1_v51VAE.safetensors',
+                    interactive=True
                 )
                 negative_prompt = gr.Textbox(
                     label="Negative Prompt",
@@ -333,19 +312,23 @@ with gr.Blocks(css=css, head=head) as demo:
         """)
     demo.queue(max_size=20, status_update_rate=0.1)
 
+app = FastAPI()
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, prevent_thread_lock=True)
+@app.post("/magic_quill/guess_prompt")
+async def guess_prompt(request: Request):
+    data = await request.json()
+    res = guess_prompt_handler(data['original_image'], data['add_color_image'], data['add_edge_image'])
+    return res
 
-    @demo.app.post("/magic_quill/guess_prompt")
-    async def guess_prompt(request: Request):
-        data = await request.json()
-        return guess_prompt_handler(data['original_image'], data['add_color_image'], data['add_edge_image'])
+@app.post("/magic_quill/process_background_img")
+async def process_background_img(request: Request):
+    img = await request.json()
+    resized_img_tensor = load_and_resize_image(img)
+    resized_img_base64 = "data:image/png;base64," + tensor_to_base64(resized_img_tensor)
+    return resized_img_base64
 
-    @demo.app.post("/magic_quill/process_background_img")
-    async def process_background_img(request: Request):
-        img = await request.json()
-        resized_img_tensor = load_and_resize_image(img)
-        return "data:image/png;base64," + tensor_to_base64(resized_img_tensor)
+app = gr.mount_gradio_app(app, demo, "/")
 
-    demo.block_thread()
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)
+    # demo.launch()

requirements.txt

@@ -14,7 +14,7 @@ anyio==4.4.0
 async-timeout==4.0.3
 attrs==23.2.0
 beautifulsoup4==4.12.3
-bitsandbytes
+bitsandbytes==0.43.3
 certifi==2024.7.4
 cffi==1.16.0
 chardet==5.2.0
@@ -33,7 +33,7 @@ einops-exts==0.0.4
 embreex==2.17.7.post5
 eval-type-backport==0.2.0
 exceptiongroup==1.2.2
-fastapi<0.112
+fastapi
 ffmpy==0.4.0
 filelock==3.15.4
 flatbuffers==24.3.25
@@ -132,11 +132,11 @@ sounddevice==0.4.7
 soupsieve==2.5
 spandrel==0.3.4
 stanza==1.1.1
-starlette<0.38
+starlette
 svg-path==6.3
 svglib==1.5.1
 svgwrite==1.4.3
-sympy==1.13.3
+sympy==1.13.1
 tabulate==0.9.0
 termcolor==2.4.0
 threadpoolctl==3.5.0
@@ -151,8 +151,7 @@ tqdm==4.66.5
 trampoline==0.1.2
 transformers==4.37.2
 trimesh==4.4.3
-torch==2.8.0
-torchvision==0.23.0
+triton==2.1.0
 torchsde==0.2.6
 typer==0.12.5
 typing-extensions==4.12.2