Upload pipeline.py with huggingface_hub

pipeline.py

@@ -0,0 +1,150 @@
+import os
+from typing import Optional, Tuple
+
+import cv2
+import numpy as np
+import torch
+from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline
+from PIL import Image
+from rembg import remove
+
+
+DEFAULT_PROMPT = (
+    "16x16 pixel art sprite, NES style, cute stuffed animal character, "
+    "strictly pixelated, chunky visible pixels, limited flat color palette, "
+    "sharp pixel boundaries, no anti-aliasing, no gradients, no shading, "
+    "indie RPG game sprite style, warm saturated color palette, "
+    "chibi proportions, thick dark outlines, flat 2-tone coloring, "
+    "white background, full body, front-facing, "
+    "bold black outlines, clean pixel edges"
+)
+
+DEFAULT_NEGATIVE_PROMPT = (
+    "realistic, 3d render, blurry, smooth, photograph, gradient, shadow, "
+    "anti-aliasing, soft edges, painterly, watercolor, sketch, detailed texture"
+)
+
+
+class PixelArtLoRAPipeline:
+    """End-to-end inference pipeline for the uploaded SDXL LoRA."""
+
+    def __init__(
+        self,
+        lora_path: str = ".",
+        base_model: str = "stabilityai/stable-diffusion-xl-base-1.0",
+        controlnet_model: str = "diffusers/controlnet-canny-sdxl-1.0",
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        self.lora_path = lora_path
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.dtype = dtype or (torch.float16 if self.device == "cuda" else torch.float32)
+
+        controlnet = ControlNetModel.from_pretrained(
+            controlnet_model,
+            torch_dtype=self.dtype,
+            use_safetensors=True,
+        )
+
+        self.pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            base_model,
+            controlnet=controlnet,
+            torch_dtype=self.dtype,
+            use_safetensors=True,
+        )
+        self.pipe.load_lora_weights(lora_path)
+        self.pipe.to(self.device)
+        self.pipe.enable_attention_slicing()
+
+    @staticmethod
+    def remove_background(image: Image.Image) -> Image.Image:
+        removed = remove(image.convert("RGBA"))
+        white_bg = Image.new("RGBA", removed.size, (255, 255, 255, 255))
+        white_bg.paste(removed, mask=removed.split()[3])
+        return white_bg.convert("RGB")
+
+    @staticmethod
+    def rembg_succeeded(image: Image.Image, threshold: float = 0.40) -> bool:
+        arr = np.array(image.convert("RGB"))
+        bg_mask = (arr[:, :, 0] >= 245) & (arr[:, :, 1] >= 245) & (arr[:, :, 2] >= 245)
+        return float(bg_mask.sum()) / bg_mask.size >= threshold
+
+    @staticmethod
+    def extract_canny(image: Image.Image, low: int = 80, high: int = 180) -> Image.Image:
+        gray = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2GRAY)
+        edges = cv2.Canny(gray, low, high)
+        return Image.fromarray(np.stack([edges] * 3, axis=-1))
+
+    @staticmethod
+    def quantize_colors(image: Image.Image, colors: int = 32) -> Image.Image:
+        quantized = image.convert("RGB").quantize(
+            colors=colors,
+            method=Image.Quantize.MEDIANCUT,
+            dither=Image.Dither.NONE,
+        )
+        return quantized.convert("RGB")
+
+    def prepare_image(
+        self,
+        image: Image.Image,
+        size: Tuple[int, int] = (512, 512),
+        bg_min_ratio: float = 0.40,
+        canny_low: int = 80,
+        canny_high: int = 180,
+    ) -> Tuple[Image.Image, Image.Image, bool]:
+        original = image.convert("RGB").resize(size)
+        bg_removed = self.remove_background(original)
+
+        if self.rembg_succeeded(bg_removed, bg_min_ratio):
+            source = bg_removed
+            rembg_ok = True
+        else:
+            source = original
+            rembg_ok = False
+
+        canny = self.extract_canny(source, canny_low, canny_high)
+        return source, canny, rembg_ok
+
+    def __call__(
+        self,
+        image: Image.Image,
+        prompt: str = DEFAULT_PROMPT,
+        negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        controlnet_conditioning_scale: float = 0.8,
+        strength: float = 0.75,
+        quantize: bool = True,
+        n_colors: int = 32,
+        seed: Optional[int] = None,
+    ) -> dict:
+        source, canny, rembg_ok = self.prepare_image(image)
+
+        generator = None
+        if seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        result = self.pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=source,
+            control_image=canny,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            strength=strength,
+            generator=generator,
+        ).images[0]
+
+        final = self.quantize_colors(result, n_colors) if quantize else result
+        return {
+            "image": final,
+            "raw_image": result,
+            "source_image": source,
+            "canny_image": canny,
+            "rembg_ok": rembg_ok,
+        }
+
+
+def load_pipeline(lora_path: Optional[str] = None) -> PixelArtLoRAPipeline:
+    return PixelArtLoRAPipeline(lora_path=lora_path or os.getenv("LORA_PATH", "."))