wip

lama_cleaner/model/paint_by_example.py

@@ -4,8 +4,9 @@ import cv2
 import torch
 from loguru import logger
 
+from lama_cleaner.helper import decode_base64_to_image
 from lama_cleaner.model.base import DiffusionInpaintModel
-from lama_cleaner.schema import Config
+from lama_cleaner.schema import InpaintRequest
 
 
 class PaintByExample(DiffusionInpaintModel):
@@ -38,16 +39,21 @@ class PaintByExample(DiffusionInpaintModel):
         else:
             self.model = self.model.to(device)
 
-    def forward(self, image, mask, config: Config):
+    def forward(self, image, mask, config: InpaintRequest):
         """Input image and output image have same size
         image: [H, W, C] RGB
         mask: [H, W, 1] 255 means area to repaint
         return: BGR IMAGE
         """
+        if config.paint_by_example_example_image is None:
+            raise ValueError("paint_by_example_example_image is required")
+        example_image, _, _ = decode_base64_to_image(
+            config.paint_by_example_example_image
+        )
         output = self.model(
             image=PIL.Image.fromarray(image),
             mask_image=PIL.Image.fromarray(mask[:, :, -1], mode="L"),
-            example_image=config.paint_by_example_example_image,
+            example_image=PIL.Image.fromarray(example_image),
             num_inference_steps=config.sd_steps,
             guidance_scale=config.sd_guidance_scale,
             negative_prompt="out of frame, lowres, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, disfigured, gross proportions, malformed limbs, watermark, signature",