update

iopaint/model/power_paint/power_paint_v2.py

@@ -1,6 +1,9 @@
+from itertools import chain
+
 import PIL.Image
 import cv2
 import torch
+from iopaint.model.original_sd_configs import get_config_files
 from loguru import logger
 from transformers import CLIPTextModel, CLIPTokenizer
 import numpy as np
@@ -14,9 +17,15 @@ from ..utils import (
     handle_from_pretrained_exceptions,
 )
 from .powerpaint_tokenizer import task_to_prompt
-from iopaint.schema import InpaintRequest
+from iopaint.schema import InpaintRequest, ModelType
 from .v2.BrushNet_CA import BrushNetModel
-from .v2.unet_2d_condition import UNet2DConditionModel
+from .v2.unet_2d_condition import UNet2DConditionModel_forward
+from .v2.unet_2d_blocks import (
+    CrossAttnDownBlock2D_forward,
+    DownBlock2D_forward,
+    CrossAttnUpBlock2D_forward,
+    UpBlock2D_forward,
+)
 
 
 class PowerPaintV2(DiffusionInpaintModel):
@@ -50,14 +59,7 @@ class PowerPaintV2(DiffusionInpaintModel):
             torch_dtype=torch_dtype,
             local_files_only=model_kwargs["local_files_only"],
         )
-        unet = handle_from_pretrained_exceptions(
-            UNet2DConditionModel.from_pretrained,
-            pretrained_model_name_or_path=self.model_id_or_path,
-            subfolder="unet",
-            variant="fp16",
-            torch_dtype=torch_dtype,
-            local_files_only=model_kwargs["local_files_only"],
-        )
+
         brushnet = BrushNetModel.from_pretrained(
             self.hf_model_id,
             subfolder="PowerPaint_Brushnet",
@@ -65,16 +67,32 @@ class PowerPaintV2(DiffusionInpaintModel):
             torch_dtype=torch_dtype,
             local_files_only=model_kwargs["local_files_only"],
         )
-        pipe = handle_from_pretrained_exceptions(
-            StableDiffusionPowerPaintBrushNetPipeline.from_pretrained,
-            pretrained_model_name_or_path=self.model_id_or_path,
-            torch_dtype=torch_dtype,
-            unet=unet,
-            brushnet=brushnet,
-            text_encoder_brushnet=text_encoder_brushnet,
-            variant="fp16",
-            **model_kwargs,
-        )
+
+        if self.model_info.is_single_file_diffusers:
+            if self.model_info.model_type == ModelType.DIFFUSERS_SD:
+                model_kwargs["num_in_channels"] = 4
+            else:
+                model_kwargs["num_in_channels"] = 9
+
+            pipe = StableDiffusionPowerPaintBrushNetPipeline.from_single_file(
+                self.model_id_or_path,
+                torch_dtype=torch_dtype,
+                load_safety_checker=False,
+                original_config_file=get_config_files()["v1"],
+                brushnet=brushnet,
+                text_encoder_brushnet=text_encoder_brushnet,
+                **model_kwargs,
+            )
+        else:
+            pipe = handle_from_pretrained_exceptions(
+                StableDiffusionPowerPaintBrushNetPipeline.from_pretrained,
+                pretrained_model_name_or_path=self.model_id_or_path,
+                torch_dtype=torch_dtype,
+                brushnet=brushnet,
+                text_encoder_brushnet=text_encoder_brushnet,
+                variant="fp16",
+                **model_kwargs,
+            )
         pipe.tokenizer = PowerPaintTokenizer(
             CLIPTokenizer.from_pretrained(self.hf_model_id, subfolder="tokenizer")
         )
@@ -95,6 +113,34 @@ class PowerPaintV2(DiffusionInpaintModel):
 
         self.callback = kwargs.pop("callback", None)
 
+        # Monkey patch the forward method of the UNet to use the brushnet_unet_forward method
+        self.model.unet.forward = UNet2DConditionModel_forward.__get__(
+            self.model.unet, self.model.unet.__class__
+        )
+
+        # Monkey patch unet down_blocks to use CrossAttnDownBlock2D_forward
+        for down_block in chain(
+            self.model.unet.down_blocks, self.model.brushnet.down_blocks
+        ):
+            if down_block.__class__.__name__ == "CrossAttnDownBlock2D":
+                down_block.forward = CrossAttnDownBlock2D_forward.__get__(
+                    down_block, down_block.__class__
+                )
+            else:
+                down_block.forward = DownBlock2D_forward.__get__(
+                    down_block, down_block.__class__
+                )
+
+        for up_block in chain(self.model.unet.up_blocks, self.model.brushnet.up_blocks):
+            if up_block.__class__.__name__ == "CrossAttnUpBlock2D":
+                up_block.forward = CrossAttnUpBlock2D_forward.__get__(
+                    up_block, up_block.__class__
+                )
+            else:
+                up_block.forward = UpBlock2D_forward.__get__(
+                    up_block, up_block.__class__
+                )
+
     def forward(self, image, mask, config: InpaintRequest):
         """Input image and output image have same size
         image: [H, W, C] RGB
@@ -129,11 +175,10 @@ class PowerPaintV2(DiffusionInpaintModel):
             brushnet_conditioning_scale=1.0,
             guidance_scale=config.sd_guidance_scale,
             output_type="np",
-            callback=self.callback,
+            callback_on_step_end=self.callback,
             height=img_h,
             width=img_w,
             generator=torch.manual_seed(config.sd_seed),
-            callback_steps=1,
         ).images[0]
 
         output = (output * 255).round().astype("uint8")