initial commit

.gitattributes

@@ -0,0 +1,39 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+sample1.jpg filter=lfs diff=lfs merge=lfs -text
+sample2.jpg filter=lfs diff=lfs merge=lfs -text
+sample3.jpg filter=lfs diff=lfs merge=lfs -text
+sample4.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,12 @@
+---
+title: Evo-Nishikie
+emoji: 🐠
+colorFrom: red
+colorTo: indigo
+sdk: gradio
+sdk_version: 4.26.0
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,162 @@
+import random
+
+from PIL import Image, ImageFilter
+from controlnet_aux import LineartDetector
+from diffusers import EulerDiscreteScheduler
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+
+# torch._inductor.config.conv_1x1_as_mm = True
+# torch._inductor.config.coordinate_descent_tuning = True
+# torch._inductor.config.epilogue_fusion = False
+# torch._inductor.config.coordinate_descent_check_all_directions = True
+from evo_nishikie_v1 import load_evo_nishikie
+
+
+DESCRIPTION = """# 🐟 Evo-Nishikie
+🤗 [モデル一覧](https://huggingface.co/SakanaAI) | 📝 [ブログ](https://sakana.ai/evo-ukiyoe/) | 🐦 [Twitter](https://twitter.com/SakanaAILabs)
+
+[Evo-Nishikie](https://huggingface.co/SakanaAI/Evo-Nishikie-v1)は[Sakana AI](https://sakana.ai/)が教育目的で開発した浮世絵に特化した画像生成モデルです。
+入力した単色摺の浮世絵（墨摺絵等）を日本語プロンプトに沿って多色摺の浮世絵（錦絵）風に変換した画像を生成することができます。より詳しくは、上記のブログをご参照ください。
+"""
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
+
+MAX_SEED = np.iinfo(np.int32).max
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+NUM_IMAGES_PER_PROMPT = 1
+SAFETY_CHECKER = True
+if SAFETY_CHECKER:
+    from safety_checker import StableDiffusionSafetyChecker
+    from transformers import CLIPFeatureExtractor
+
+    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker"
+    ).to(device)
+    feature_extractor = CLIPFeatureExtractor.from_pretrained(
+        "openai/clip-vit-base-patch32"
+    )
+
+    def check_nsfw_images(
+        images: list[Image.Image],
+    ) -> tuple[list[Image.Image], list[bool]]:
+        safety_checker_input = feature_extractor(images, return_tensors="pt").to(device)
+        has_nsfw_concepts = safety_checker(
+            images=[images], clip_input=safety_checker_input.pixel_values.to(device)
+        )
+
+        return images, has_nsfw_concepts
+
+
+pipe = load_evo_nishikie(device)
+pipe.scheduler = EulerDiscreteScheduler.from_config(
+    pipe.scheduler.config, use_karras_sigmas=True,
+)
+# pipe.unet.to(memory_format=torch.channels_last)
+# pipe.controlnet.to(memory_format=torch.channels_last)
+# pipe.vae.to(memory_format=torch.channels_last)
+# # Compile the UNet, ControlNet and VAE.
+# pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
+# pipe.controlnet = torch.compile(pipe.controlnet, mode="max-autotune", fullgraph=True)
+# pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)
+
+lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
+image_filter = ImageFilter.MedianFilter(size=3)
+BINARY_THRESHOLD = 40
+
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
+@spaces.GPU
+@torch.inference_mode()
+def generate(
+    input_image: Image.Image,
+    prompt: str,
+    seed: int = 0,
+    randomize_seed: bool = False,
+    progress=gr.Progress(track_tqdm=True),
+):
+    pipe.to(device)
+
+    lineart_image = lineart_detector(input_image, coarse=False, image_resolution=1024)
+    lineart_image_filtered = lineart_image.filter(image_filter)
+    conditioning_image = lineart_image_filtered.point(lambda p: 255 if p > BINARY_THRESHOLD else 0).convert("L")
+
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.Generator().manual_seed(seed)
+
+    images = pipe(
+        prompt=prompt + "最高品質の輻の浮世絵。超詳細。",
+        negative_prompt="暗い",
+        image=conditioning_image,
+        guidance_scale=7.0,
+        controlnet_conditioning_scale=0.8,
+        num_inference_steps=50,
+        generator=generator,
+        num_images_per_prompt=NUM_IMAGES_PER_PROMPT,
+        output_type="pil",
+    ).images
+
+    if SAFETY_CHECKER:
+        images, has_nsfw_concepts = check_nsfw_images(images)
+        if any(has_nsfw_concepts):
+            gr.Warning("NSFW content detected.")
+            return Image.new("RGB", (512, 512), "WHITE"), seed
+    return images[0], seed
+
+
+examples = [
+    ["./sample1.jpg", "女性がやかんと鍋を持ち、小屋の前に立っています。背景には室内で会話する人々がいます。"],
+    ["./sample2.jpg", "着物を着た女性が、赤ん坊を抱え、もう一人の子どもが手押し車を引いています。背景には木があります。"],
+    ["./sample3.jpg", "女性が花柄の着物を着ており、他の人物たちが座りながら会話しています。背景には家の内部があります。"],
+    ["./sample4.jpg", "花柄や模様入りの着物を着た男女が室内で集まり、煎茶の準備をしています。背景に木材の装飾があります。"],
+]
+
+css = """
+.gradio-container{max-width: 1380px !important}
+h1{text-align:center}
+"""
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(DESCRIPTION)
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(image_mode="RGB", type="pil", show_label=False)
+            prompt = gr.Textbox(placeholder="日本語でプロンプトを入力してください。", show_label=False)
+            submit = gr.Button()
+            with gr.Accordion("詳細設定", open=False):
+                seed = gr.Slider(label="シード値", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                randomize_seed = gr.Checkbox(label="ランダムにシード値を決定", value=True)
+        with gr.Column():
+            result = gr.Image(label="Evo-Nishikieからの生成結果", type="pil", show_label=False)
+    gr.Examples(examples=examples, inputs=[input_image, prompt], outputs=[result, seed], fn=generate)
+    gr.on(
+        triggers=[
+            submit.click,
+        ],
+        fn=generate,
+        inputs=[
+            input_image,
+            prompt,
+            seed,
+            randomize_seed,
+        ],
+        outputs=[result, seed],
+        api_name="run",
+    )
+    gr.Markdown("""⚠️ 本モデルは実験段階のプロトタイプであり、教育および研究開発の目的でのみ提供されています。商用利用や、障害が重大な影響を及ぼす可能性のある環境（ミッションクリティカルな環境）での使用には適していません。
+                本モデルの使用は、利用者の自己責任で行われ、その性能や結果については何ら保証されません。
+                Sakana AIは、本モデルの使用によって生じた直接的または間接的な損失に対して、結果に関わらず、一切の責任を負いません。
+                利用者は、本モデルの使用に伴うリスクを十分に理解し、自身の判断で使用することが必要です。
+                アップロードされた画像は画像生成のみに使用され、サーバー上に保存されることはありません。
+
+                出典：サンプル画像はすべて[日本古典籍データセット（国文学研究資料館蔵）『絵本玉かつら』](http://codh.rois.ac.jp/pmjt/book/200013861/)から引用しました。""")
+
+demo.queue().launch()

evo_nishikie_v1.py

@@ -0,0 +1,187 @@
+import gc
+import os
+from typing import Dict, List, Union
+
+from diffusers import (
+    ControlNetModel,
+    StableDiffusionXLControlNetPipeline,
+    UNet2DConditionModel,
+)
+from huggingface_hub import hf_hub_download
+import safetensors
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer, CLIPTextModelWithProjection
+
+
+# Base models
+SDXL_REPO = "stabilityai/stable-diffusion-xl-base-1.0"
+DPO_REPO = "mhdang/dpo-sdxl-text2image-v1"
+JN_REPO = "RunDiffusion/Juggernaut-XL-v9"
+JSDXL_REPO = "stabilityai/japanese-stable-diffusion-xl"
+
+# Evo-Ukiyoe
+UKIYOE_REPO = "SakanaAI/Evo-Ukiyoe-v1"
+
+# Evo-Nishikie
+NISHIKIE_REPO = "SakanaAI/Evo-Nishikie-v1"
+
+
+def load_state_dict(checkpoint_file: Union[str, os.PathLike], device: str = "cpu"):
+    file_extension = os.path.basename(checkpoint_file).split(".")[-1]
+    if file_extension == "safetensors":
+        return safetensors.torch.load_file(checkpoint_file, device=device)
+    else:
+        return torch.load(checkpoint_file, map_location=device)
+
+
+def load_from_pretrained(
+    repo_id,
+    filename="diffusion_pytorch_model.fp16.safetensors",
+    subfolder="unet",
+    device="cuda",
+) -> Dict[str, torch.Tensor]:
+    return load_state_dict(
+        hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            subfolder=subfolder,
+        ),
+        device=device,
+    )
+
+
+def reshape_weight_task_tensors(task_tensors, weights):
+    """
+    Reshapes `weights` to match the shape of `task_tensors` by unsqueezing in the remaining dimensions.
+
+    Args:
+        task_tensors (`torch.Tensor`): The tensors that will be used to reshape `weights`.
+        weights (`torch.Tensor`): The tensor to be reshaped.
+
+    Returns:
+        `torch.Tensor`: The reshaped tensor.
+    """
+    new_shape = weights.shape + (1,) * (task_tensors.dim() - weights.dim())
+    weights = weights.view(new_shape)
+    return weights
+
+
+def linear(task_tensors: List[torch.Tensor], weights: torch.Tensor) -> torch.Tensor:
+    """
+    Merge the task tensors using `linear`.
+
+    Args:
+        task_tensors(`List[torch.Tensor]`):The task tensors to merge.
+        weights (`torch.Tensor`):The weights of the task tensors.
+
+    Returns:
+        `torch.Tensor`: The merged tensor.
+    """
+    task_tensors = torch.stack(task_tensors, dim=0)
+    # weighted task tensors
+    weights = reshape_weight_task_tensors(task_tensors, weights)
+    weighted_task_tensors = task_tensors * weights
+    mixed_task_tensors = weighted_task_tensors.sum(dim=0)
+    return mixed_task_tensors
+
+
+def merge_models(task_tensors, weights):
+    keys = list(task_tensors[0].keys())
+    weights = torch.tensor(weights, device=task_tensors[0][keys[0]].device)
+    state_dict = {}
+    for key in tqdm(keys, desc="Merging"):
+        w_list = []
+        for i, sd in enumerate(task_tensors):
+            w = sd.pop(key)
+            w_list.append(w)
+        new_w = linear(task_tensors=w_list, weights=weights)
+        state_dict[key] = new_w
+    return state_dict
+
+
+def split_conv_attn(weights):
+    attn_tensors = {}
+    conv_tensors = {}
+    for key in list(weights.keys()):
+        if any(k in key for k in ["to_k", "to_q", "to_v", "to_out.0"]):
+            attn_tensors[key] = weights.pop(key)
+        else:
+            conv_tensors[key] = weights.pop(key)
+    return {"conv": conv_tensors, "attn": attn_tensors}
+
+
+def load_evo_nishikie(device="cuda") -> StableDiffusionXLControlNetPipeline:
+    # Load base models
+    sdxl_weights = split_conv_attn(load_from_pretrained(SDXL_REPO, device=device))
+    dpo_weights = split_conv_attn(
+        load_from_pretrained(
+            DPO_REPO, "diffusion_pytorch_model.safetensors", device=device
+        )
+    )
+    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(
+        [sd["conv"] for sd in tensors],
+        [
+            0.15928833971605916,
+            0.1032449268871776,
+            0.6503217149752791,
+            0.08714501842148402,
+        ],
+    )
+    new_attn = merge_models(
+        [sd["attn"] for sd in tensors],
+        [
+            0.1877279276437178,
+            0.20014114603909822,
+            0.3922685507065275,
+            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})
+
+    # Load other modules
+    text_encoder = CLIPTextModelWithProjection.from_pretrained(
+        JSDXL_REPO, subfolder="text_encoder", torch_dtype=torch.float16, variant="fp16",
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        JSDXL_REPO, subfolder="tokenizer", use_fast=False,
+    )
+
+    # Load Evo-Nishikie weights
+    controlnet = ControlNetModel.from_pretrained(
+        NISHIKIE_REPO, torch_dtype=torch.float16, device=device,
+    )
+
+    # Load pipeline
+    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+        SDXL_REPO,
+        unet=unet,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        controlnet=controlnet,
+        torch_dtype=torch.float16,
+        variant="fp16",
+    )
+
+    # Load Evo-Ukiyoe weights
+    pipe.load_lora_weights(UKIYOE_REPO)
+    pipe.fuse_lora(lora_scale=1.0)
+
+    pipe = pipe.to(device, dtype=torch.float16)
+
+    return pipe

requirements.txt

@@ -0,0 +1,10 @@
+torch
+torchvision
+
+accelerate==0.32.0
+controlnet-aux==0.0.9
+diffusers==0.29.2
+gradio==4.38.1
+sentencepiece==0.2.0
+transformers==4.42.3
+peft==0.11.1

safety_checker.py

@@ -0,0 +1,137 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel
+
+
+def cosine_distance(image_embeds, text_embeds):
+    normalized_image_embeds = nn.functional.normalize(image_embeds)
+    normalized_text_embeds = nn.functional.normalize(text_embeds)
+    return torch.mm(normalized_image_embeds, normalized_text_embeds.t())
+
+
+class StableDiffusionSafetyChecker(PreTrainedModel):
+    config_class = CLIPConfig
+
+    _no_split_modules = ["CLIPEncoderLayer"]
+
+    def __init__(self, config: CLIPConfig):
+        super().__init__(config)
+
+        self.vision_model = CLIPVisionModel(config.vision_config)
+        self.visual_projection = nn.Linear(
+            config.vision_config.hidden_size, config.projection_dim, bias=False
+        )
+
+        self.concept_embeds = nn.Parameter(
+            torch.ones(17, config.projection_dim), requires_grad=False
+        )
+        self.special_care_embeds = nn.Parameter(
+            torch.ones(3, config.projection_dim), requires_grad=False
+        )
+
+        self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False)
+        self.special_care_embeds_weights = nn.Parameter(
+            torch.ones(3), requires_grad=False
+        )
+
+    @torch.no_grad()
+    def forward(self, clip_input, images):
+        pooled_output = self.vision_model(clip_input)[1]  # pooled_output
+        image_embeds = self.visual_projection(pooled_output)
+
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        special_cos_dist = (
+            cosine_distance(image_embeds, self.special_care_embeds)
+            .cpu()
+            .float()
+            .numpy()
+        )
+        cos_dist = (
+            cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy()
+        )
+
+        result = []
+        batch_size = image_embeds.shape[0]
+        for i in range(batch_size):
+            result_img = {
+                "special_scores": {},
+                "special_care": [],
+                "concept_scores": {},
+                "bad_concepts": [],
+            }
+
+            # increase this value to create a stronger `nfsw` filter
+            # at the cost of increasing the possibility of filtering benign images
+            adjustment = 0.0
+
+            for concept_idx in range(len(special_cos_dist[0])):
+                concept_cos = special_cos_dist[i][concept_idx]
+                concept_threshold = self.special_care_embeds_weights[concept_idx].item()
+                result_img["special_scores"][concept_idx] = round(
+                    concept_cos - concept_threshold + adjustment, 3
+                )
+                if result_img["special_scores"][concept_idx] > 0:
+                    result_img["special_care"].append(
+                        {concept_idx, result_img["special_scores"][concept_idx]}
+                    )
+                    adjustment = 0.01
+
+            for concept_idx in range(len(cos_dist[0])):
+                concept_cos = cos_dist[i][concept_idx]
+                concept_threshold = self.concept_embeds_weights[concept_idx].item()
+                result_img["concept_scores"][concept_idx] = round(
+                    concept_cos - concept_threshold + adjustment, 3
+                )
+                if result_img["concept_scores"][concept_idx] > 0:
+                    result_img["bad_concepts"].append(concept_idx)
+
+            result.append(result_img)
+
+        has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result]
+
+        return has_nsfw_concepts
+
+    @torch.no_grad()
+    def forward_onnx(self, clip_input: torch.FloatTensor, images: torch.FloatTensor):
+        pooled_output = self.vision_model(clip_input)[1]  # pooled_output
+        image_embeds = self.visual_projection(pooled_output)
+
+        special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds)
+        cos_dist = cosine_distance(image_embeds, self.concept_embeds)
+
+        # increase this value to create a stronger `nsfw` filter
+        # at the cost of increasing the possibility of filtering benign images
+        adjustment = 0.0
+
+        special_scores = (
+            special_cos_dist - self.special_care_embeds_weights + adjustment
+        )
+        # special_scores = special_scores.round(decimals=3)
+        special_care = torch.any(special_scores > 0, dim=1)
+        special_adjustment = special_care * 0.01
+        special_adjustment = special_adjustment.unsqueeze(1).expand(
+            -1, cos_dist.shape[1]
+        )
+
+        concept_scores = (cos_dist - self.concept_embeds_weights) + special_adjustment
+        # concept_scores = concept_scores.round(decimals=3)
+        has_nsfw_concepts = torch.any(concept_scores > 0, dim=1)
+
+        images[has_nsfw_concepts] = 0.0  # black image
+
+        return images, has_nsfw_concepts