Merge branch 'main' into qwen-image-edit

diffsynth/configs/model_config.py

@@ -169,6 +169,7 @@ model_loader_configs = [
     (None, "8004730443f55db63092006dd9f7110e", ["qwen_image_text_encoder"], [QwenImageTextEncoder], "diffusers"),
     (None, "ed4ea5824d55ec3107b09815e318123a", ["qwen_image_vae"], [QwenImageVAE], "diffusers"),
     (None, "073bce9cf969e317e5662cd570c3e79c", ["qwen_image_blockwise_controlnet"], [QwenImageBlockWiseControlNet], "civitai"),
+    (None, "a9e54e480a628f0b956a688a81c33bab", ["qwen_image_blockwise_controlnet"], [QwenImageBlockWiseControlNet], "civitai"),
 ]
 huggingface_model_loader_configs = [
     # These configs are provided for detecting model type automatically.

diffsynth/models/qwen_image_controlnet.py

@@ -1,10 +1,7 @@
 import torch
 import torch.nn as nn
-from .qwen_image_dit import QwenEmbedRope, QwenImageTransformerBlock
-from ..vram_management import gradient_checkpoint_forward
-from einops import rearrange
-from .sd3_dit import TimestepEmbeddings, RMSNorm
-
+from .sd3_dit import RMSNorm
+from .utils import hash_state_dict_keys
 
 
 class BlockWiseControlBlock(torch.nn.Module):
@@ -35,10 +32,11 @@ class QwenImageBlockWiseControlNet(torch.nn.Module):
         self,
         num_layers: int = 60,
         in_dim: int = 64,
+        additional_in_dim: int = 0,
         dim: int = 3072,
     ):
         super().__init__()
-        self.img_in = nn.Linear(in_dim, dim)
+        self.img_in = nn.Linear(in_dim + additional_in_dim, dim)
         self.controlnet_blocks = nn.ModuleList(
             [
                 BlockWiseControlBlock(dim)
@@ -68,4 +66,9 @@ class QwenImageBlockWiseControlNetStateDictConverter():
         pass
 
     def from_civitai(self, state_dict):
-        return state_dict
+        hash_value = hash_state_dict_keys(state_dict)
+        extra_kwargs = {}
+        if hash_value == "a9e54e480a628f0b956a688a81c33bab":
+            # inpaint controlnet
+            extra_kwargs = {"additional_in_dim": 4}
+        return state_dict, extra_kwargs

diffsynth/pipelines/qwen_image.py

@@ -67,6 +67,7 @@ class QwenImagePipeline(BasePipeline):
             QwenImageUnit_ShapeChecker(),
             QwenImageUnit_NoiseInitializer(),
             QwenImageUnit_InputImageEmbedder(),
+            QwenImageUnit_Inpaint(),
             QwenImageUnit_PromptEmbedder(),
             QwenImageUnit_EntityControl(),
             QwenImageUnit_BlockwiseControlNet(),
@@ -259,6 +260,10 @@ class QwenImagePipeline(BasePipeline):
         # Image
         input_image: Image.Image = None,
         denoising_strength: float = 1.0,
+        # Inpaint
+        inpaint_mask: Image.Image = None,
+        inpaint_blur_size: int = None,
+        inpaint_blur_sigma: float = None,
         # Shape
         height: int = 1328,
         width: int = 1328,
@@ -297,6 +302,7 @@ class QwenImagePipeline(BasePipeline):
         inputs_shared = {
             "cfg_scale": cfg_scale,
             "input_image": input_image, "denoising_strength": denoising_strength,
+            "inpaint_mask": inpaint_mask, "inpaint_blur_size": inpaint_blur_size, "inpaint_blur_sigma": inpaint_blur_sigma,
             "height": height, "width": width,
             "seed": seed, "rand_device": rand_device,
             "enable_fp8_attention": enable_fp8_attention,
@@ -324,7 +330,7 @@ class QwenImagePipeline(BasePipeline):
                 noise_pred = noise_pred_posi
 
             # Scheduler
-            inputs_shared["latents"] = self.scheduler.step(noise_pred, self.scheduler.timesteps[progress_id], inputs_shared["latents"])
+            inputs_shared["latents"] = self.step(self.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs_shared)
         
         # Decode
         self.load_models_to_device(['vae'])
@@ -373,7 +379,26 @@ class QwenImageUnit_InputImageEmbedder(PipelineUnit):
             return {"latents": noise, "input_latents": input_latents}
         else:
             latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0])
-            return {"latents": latents, "input_latents": None}
+            return {"latents": latents, "input_latents": input_latents}
+
+
+
+class QwenImageUnit_Inpaint(PipelineUnit):
+    def __init__(self):
+        super().__init__(
+            input_params=("inpaint_mask", "height", "width", "inpaint_blur_size", "inpaint_blur_sigma"),
+        )
+
+    def process(self, pipe: QwenImagePipeline, inpaint_mask, height, width, inpaint_blur_size, inpaint_blur_sigma):
+        if inpaint_mask is None:
+            return {}
+        inpaint_mask = pipe.preprocess_image(inpaint_mask.convert("RGB").resize((width // 8, height // 8)), min_value=0, max_value=1)
+        inpaint_mask = inpaint_mask.mean(dim=1, keepdim=True)
+        if inpaint_blur_size is not None and inpaint_blur_sigma is not None:
+            from torchvision.transforms import GaussianBlur
+            blur = GaussianBlur(kernel_size=inpaint_blur_size * 2 + 1, sigma=inpaint_blur_sigma)
+            inpaint_mask = blur(inpaint_mask)
+        return {"inpaint_mask": inpaint_mask}
 
 
 class QwenImageUnit_PromptEmbedder(PipelineUnit):

diffsynth/utils/__init__.py

@@ -139,6 +139,20 @@ class BasePipeline(torch.nn.Module):
             else:
                 model.eval()
                 model.requires_grad_(False)
+                
+    
+    def blend_with_mask(self, base, addition, mask):
+        return base * (1 - mask) + addition * mask
+    
+    
+    def step(self, scheduler, latents, progress_id, noise_pred, input_latents=None, inpaint_mask=None, **kwargs):
+        timestep = scheduler.timesteps[progress_id]
+        if inpaint_mask is not None:
+            noise_pred_expected = scheduler.return_to_timestep(scheduler.timesteps[progress_id], latents, input_latents)
+            noise_pred = self.blend_with_mask(noise_pred_expected, noise_pred, inpaint_mask)
+        latents_next = scheduler.step(noise_pred, timestep, latents)
+        return latents_next
+
 
 
 @dataclass