Update evo_nishikie_v1.py

Adding `EvoNishikieConditioningImageProcessor` class

evo_nishikie_v1.py

@@ -1,7 +1,6 @@
 import gc
-from io import BytesIO
 import os
-from typing import Dict, List, Union
+from typing import Dict, List, Tuple, Union
 
 from PIL import Image, ImageFilter
 from controlnet_aux import LineartDetector
@@ -11,7 +10,6 @@ from diffusers import (
     UNet2DConditionModel,
 )
 from huggingface_hub import hf_hub_download
-import requests
 import safetensors
 import torch
 from tqdm import tqdm
@@ -30,8 +28,17 @@ UKIYOE_REPO = "SakanaAI/Evo-Ukiyoe-v1"
 # Evo-Nishikie
 NISHIKIE_REPO = "SakanaAI/Evo-Nishikie-v1"
 
-# Threshold for image binarization
-BINARY_THRESHOLD = 40
+
+class EvoNishikieConditioningImageProcessor:
+    def __init__(self, device="cpu"):
+        self.lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators").to(device)
+        self.image_filter = ImageFilter.MedianFilter(size=3)
+
+    def __call__(self, original_image: Image.Image) -> Image.Image:
+        lineart_image = self.lineart_detector(original_image, coarse=False, image_resolution=1024)
+        lineart_image_filtered = lineart_image.filter(self.image_filter)
+        conditioning_image = lineart_image_filtered.point(lambda p: 255 if p > 40 else 0).convert("L")
+        return conditioning_image
 
 
 def load_state_dict(checkpoint_file: Union[str, os.PathLike], device: str = "cpu"):
@@ -118,7 +125,9 @@ def split_conv_attn(weights):
     return {"conv": conv_tensors, "attn": attn_tensors}
 
 
-def load_evo_nishikie(device="cuda") -> StableDiffusionXLControlNetPipeline:
+def load_evo_nishikie(device="cuda", processor_device="cpu") -> Tuple[
+    StableDiffusionXLControlNetPipeline, EvoNishikieConditioningImageProcessor
+]:
     # Load base models
     sdxl_weights = split_conv_attn(load_from_pretrained(SDXL_REPO, device=device))
     dpo_weights = split_conv_attn(
@@ -190,4 +199,8 @@ def load_evo_nishikie(device="cuda") -> StableDiffusionXLControlNetPipeline:
     pipe.fuse_lora(lora_scale=1.0)
 
     pipe = pipe.to(device, dtype=torch.float16)
-    return pipe
+
+    # Load conditioning image processor
+    processor = EvoNishikieConditioningImageProcessor(device=processor_device)
+
+    return pipe, processor