initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.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

README.md

@@ -0,0 +1,99 @@
+---
+library_name: diffusers
+license: apache-2.0
+language:
+- ja
+pipeline_tag: text-to-image
+tags:
+- stable-diffusion
+---
+# 🐟 EvoSDXL-JP-v1
+
+🤗 [Models](https://huggingface.co/SakanaAI) | 📚 [Paper](https://arxiv.org/abs/2403.13187) | 📝 [Blog](https://sakana.ai/evosdxl-jp/) | 🐦 [Twitter](https://twitter.com/SakanaAILabs)
+
+
+**EvoSDXL-JP-v1** is an experimental education-purpose Japanese SDXL Lightning. 
+This model was created using the Evolutionary Model Merge method. 
+Please refer to our [report](https://arxiv.org/abs/2403.13187) and [blog](https://sakana.ai/evosdxl-jp/) for more details. 
+This model was produced by merging the following models. 
+We are grateful to the developers of the source models.
+
+- [SDXL](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
+- [Juggernaut-XL-v9](https://huggingface.co/RunDiffusion/Juggernaut-XL-v9)
+- [SDXL-DPO](https://huggingface.co/mhdang/dpo-sdxl-text2image-v1)
+- [JSDXL](https://huggingface.co/stabilityai/japanese-stable-diffusion-xl)
+- [SDXL-Lightning](https://huggingface.co/ByteDance/SDXL-Lightning)
+
+
+## Usage
+
+Use the code below to get started with the model.
+
+
+<details>
+<summary> Click to expand </summary>
+
+1. Git clone this model card
+   ```
+   git clone https://huggingface.co/SakanaAI/EvoSDXL-JP-v1
+   ```
+2. Install packages
+   ```
+   cd EvoSDXL-JP-v1
+   pip install -r requirements.txt
+   ```
+3. Run
+   ```python
+   from evosdxl_jp_v1 import load_evosdxl_jp
+    
+   prompt = "柴犬"
+   pipe = load_evosdxl_jp(device="cuda")
+   images = pipe(prompt, num_inference_steps=4, guidance_scale=0).images
+   images[0].save("image.png")
+   ```
+
+</details>
+
+
+ 
+## Model Details
+
+<!-- Provide a longer summary of what this model is. -->
+
+- **Developed by:** [Sakana AI](https://sakana.ai/)
+- **Model type:** Diffusion-based text-to-image generative model
+- **Language(s):** Japanese
+- **Repository:** [SakanaAI/evolutionary-model-merge](https://github.com/SakanaAI/evolutionary-model-merge)
+- **Paper:** https://arxiv.org/abs/2403.13187
+- **Blog:** https://sakana.ai/evosdxl-jp/
+
+
+## License
+The Python script included in this repository is licensed under the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0).
+Please note that the license for the model/pipeline generated by this script is inherited from the source models.
+
+## Uses
+This model is provided for research and development purposes only and should be considered as an experimental prototype. 
+It is not intended for commercial use or deployment in mission-critical environments. 
+Use of this model is at the user's own risk, and its performance and outcomes are not guaranteed. 
+Sakana AI shall not be liable for any direct, indirect, special, incidental, or consequential damages, or any loss arising from the use of this model, regardless of the results obtained. 
+Users must fully understand the risks associated with the use of this model and use it at their own discretion.
+
+
+## Acknowledgement
+
+We would like to thank the developers of the source models for their contributions and for making their work available. 
+
+
+## Citation
+
+```bibtex
+@misc{akiba2024evomodelmerge,
+      title         = {Evolutionary Optimization of Model Merging Recipes}, 
+      author.       = {Takuya Akiba and Makoto Shing and Yujin Tang and Qi Sun and David Ha},
+      year          = {2024},
+      eprint        = {2403.13187},
+      archivePrefix = {arXiv},
+      primaryClass  = {cs.NE}
+}
+```

evosdxl_jp_v1.py

@@ -0,0 +1,204 @@
+import os
+from typing import List, Dict, Union
+from tqdm import tqdm
+import torch
+import safetensors
+from huggingface_hub import hf_hub_download
+from transformers import AutoTokenizer, CLIPTextModelWithProjection
+from diffusers import (
+    StableDiffusionXLPipeline,
+    UNet2DConditionModel,
+    EulerDiscreteScheduler,
+)
+from diffusers.loaders import LoraLoaderMixin
+
+SDXL_REPO = "stabilityai/stable-diffusion-xl-base-1.0"
+JSDXL_REPO = "stabilityai/japanese-stable-diffusion-xl"
+L_REPO = "ByteDance/SDXL-Lightning"
+
+
+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 unsqeezing in the remaining dimenions.
+
+    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_evosdxl_jp(device="cuda") -> StableDiffusionXLPipeline:
+    sdxl_weights = split_conv_attn(load_from_pretrained(SDXL_REPO, device=device))
+    dpo_weights = split_conv_attn(
+        load_from_pretrained(
+            "mhdang/dpo-sdxl-text2image-v1",
+            "diffusion_pytorch_model.safetensors",
+            device=device,
+        )
+    )
+    jn_weights = split_conv_attn(
+        load_from_pretrained("RunDiffusion/Juggernaut-XL-v9", device=device)
+    )
+    jsdxl_weights = split_conv_attn(load_from_pretrained(JSDXL_REPO, device=device))
+    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,
+        ],
+    )
+    del sdxl_weights, dpo_weights, jn_weights, jsdxl_weights
+    torch.cuda.empty_cache()
+    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})
+    state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(
+        L_REPO, weight_name="sdxl_lightning_4step_lora.safetensors"
+    )
+    LoraLoaderMixin.load_lora_into_unet(state_dict, network_alphas, unet)
+    unet.fuse_lora(lora_scale=3.224682864579401)
+    new_weights = split_conv_attn(unet.state_dict())
+    l_weights = split_conv_attn(
+        load_from_pretrained(
+            L_REPO,
+            "sdxl_lightning_4step_unet.safetensors",
+            subfolder=None,
+            device=device,
+        )
+    )
+    jnl_weights = split_conv_attn(
+        load_from_pretrained(
+            "RunDiffusion/Juggernaut-XL-Lightning",
+            "diffusion_pytorch_model.bin",
+            device=device,
+        )
+    )
+    tensors = [l_weights, jnl_weights, new_weights]
+    new_conv = merge_models(
+        [sd["conv"] for sd in tensors],
+        [0.47222002022088533, 0.48419531030361584, 0.04358466947549889],
+    )
+    new_attn = merge_models(
+        [sd["attn"] for sd in tensors],
+        [0.023119324530758375, 0.04924981616469831, 0.9276308593045434],
+    )
+    new_weights = {**new_conv, **new_attn}
+    unet = UNet2DConditionModel.from_config(unet_config).to(device=device)
+    unet.load_state_dict({**new_conv, **new_attn})
+
+    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
+    )
+
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        SDXL_REPO,
+        unet=unet,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        torch_dtype=torch.float16,
+        variant="fp16",
+    )
+    # Ensure sampler uses "trailing" timesteps.
+    pipe.scheduler = EulerDiscreteScheduler.from_config(
+        pipe.scheduler.config, timestep_spacing="trailing"
+    )
+    pipe = pipe.to(device, dtype=torch.float16)
+    return pipe
+
+
+if __name__ == "__main__":
+    pipe: StableDiffusionXLPipeline = load_evosdxl_jp()
+    images = pipe("犬", num_inference_steps=4, guidance_scale=0).images
+    images[0].save("out.png")

requirements.txt

@@ -0,0 +1,4 @@
+diffusers==0.26.0
+sentencepiece
+transformers
+accelerate