Update evo_nishikie_v1.py

evo_nishikie_v1.py

@@ -3,8 +3,8 @@ from io import BytesIO
 import os
 from typing import Dict, List, Union
 
-from PIL import Image
-from controlnet_aux import CannyDetector
+from PIL import Image, ImageFilter
+from controlnet_aux import LineartDetector
 from diffusers import (
     ControlNetModel,
     StableDiffusionXLControlNetPipeline,
@@ -17,7 +17,8 @@ import torch
 from tqdm import tqdm
 from transformers import AutoTokenizer, CLIPTextModelWithProjection
 
-# Base models (fine-tuned from SDXL-1.0)
+
+# Base models
 SDXL_REPO = "stabilityai/stable-diffusion-xl-base-1.0"
 DPO_REPO = "mhdang/dpo-sdxl-text2image-v1"
 JN_REPO = "RunDiffusion/Juggernaut-XL-v9"
@@ -29,6 +30,9 @@ UKIYOE_REPO = "SakanaAI/Evo-Ukiyoe-v1"
 # Evo-Nishikie
 NISHIKIE_REPO = "SakanaAI/Evo-Nishikie-v1"
 
+# Threshold for image binarization
+BINARY_THRESHOLD = 40
+
 
 def load_state_dict(checkpoint_file: Union[str, os.PathLike], device: str = "cpu"):
     file_extension = os.path.basename(checkpoint_file).split(".")[-1]
@@ -124,6 +128,7 @@ def load_evo_nishikie(device="cuda") -> StableDiffusionXLControlNetPipeline:
     )
     jn_weights = split_conv_attn(load_from_pretrained(JN_REPO, device=device))
     jsdxl_weights = split_conv_attn(load_from_pretrained(JSDXL_REPO, device=device))
+
     # Merge base models
     tensors = [sdxl_weights, dpo_weights, jn_weights, jsdxl_weights]
     new_conv = merge_models(
@@ -144,11 +149,14 @@ def load_evo_nishikie(device="cuda") -> StableDiffusionXLControlNetPipeline:
             0.2198623756106564,
         ],
     )
+
+    # Delete no longer needed variables to free
     del sdxl_weights, dpo_weights, jn_weights, jsdxl_weights
     gc.collect()
     if "cuda" in device:
         torch.cuda.empty_cache()
 
+    # Instantiate UNet
     unet_config = UNet2DConditionModel.load_config(SDXL_REPO, subfolder="unet")
     unet = UNet2DConditionModel.from_config(unet_config).to(device=device)
     unet.load_state_dict({**new_conv, **new_attn})
@@ -181,26 +189,5 @@ def load_evo_nishikie(device="cuda") -> StableDiffusionXLControlNetPipeline:
     # Load Evo-Ukiyoe weights
     pipe.load_lora_weights(UKIYOE_REPO)
     pipe.fuse_lora(lora_scale=1.0)
-    return pipe
-
 
-if __name__ == "__main__":
-    url = "https://sakana.ai/assets/nedo-grant/nedo_grant.jpeg"
-    original_image = Image.open(
-        BytesIO(requests.get(url).content)
-    ).resize((1024, 1024), Image.Resampling.LANCZOS)
-    canny_detector = CannyDetector()
-    canny_image = canny_detector(original_image, image_resolution=1024)
-    pipe: StableDiffusionXLControlNetPipeline = load_evo_nishikie()
-    images = pipe(
-        prompt="銀杏が色づく。草木が生えた地面と青空の富士山。最高品質の輻の浮世絵。",
-        negative_prompt="暗い。",
-        image=canny_image,
-        guidance_scale=8.0,
-        controlnet_conditioning_scale=0.6,
-        num_inference_steps=50,
-        generator=torch.Generator().manual_seed(0),
-        num_images_per_prompt=1,
-        output_type="pil",
-    ).images
-    images[0].save("out.png")
+    return pipe