Nitrofusion

🚀 NitroFusion

NitroFusion是一個通過動態對抗訓練實現高保真單步擴散的項目。該項目提出了新的方法，在文本到圖像生成領域展現出高效且優質的圖像生成能力，為相關研究和應用提供了新的思路和工具。

基礎信息

🚀 快速開始

你可以通過以下鏈接快速瞭解和體驗NitroFusion：

• arXiv論文
• 項目頁面
• Hugging Face在線演示

✨ 主要特性

模型概覽

• nitrosd - realism_unet.safetensors：能生成具有精細細節的逼真圖像。
• nitrosd - vibrant_unet.safetensors：生成的圖像具有鮮豔、飽和的色彩特徵。
• 兩個模型均支持1到4步推理。

最新消息

• 2025年1月6日：發佈了ComfyUI檢查點nitrosd - realism_comfyui.safetensors和nitrosd - vibrant_comfyui.safetensors，以及一個工作流。
• 2024年12月4日：論文在arXiv上發佈，項目頁面公開。
• 2024年11月30日：單步文本到圖像演示在🤗 Hugging Face Space上公開。
• 2024年11月29日：發佈了兩個檢查點：NitroSD - Realism和NitroSD - Vibrant。

💻 使用示例

基礎用法

首先，我們需要實現帶有時間步偏移的調度器以進行多步推理：

from diffusers import LCMScheduler
class TimestepShiftLCMScheduler(LCMScheduler):
    def __init__(self, *args, shifted_timestep=250, **kwargs):
        super().__init__(*args, **kwargs)
        self.register_to_config(shifted_timestep=shifted_timestep)
    def set_timesteps(self, *args, **kwargs):
        super().set_timesteps(*args, **kwargs)
        self.origin_timesteps = self.timesteps.clone()
        self.shifted_timesteps = (self.timesteps * self.config.shifted_timestep / self.config.num_train_timesteps).long()
        self.timesteps = self.shifted_timesteps
    def step(self, model_output, timestep, sample, generator=None, return_dict=True):
        if self.step_index is None:
            self._init_step_index(timestep)
        self.timesteps = self.origin_timesteps
        output = super().step(model_output, timestep, sample, generator, return_dict)
        self.timesteps = self.shifted_timesteps
        return output

然後，我們可以使用擴散器管道：

import torch
from diffusers import DiffusionPipeline, UNet2DConditionModel
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
# Load model.
base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
repo = "ChenDY/NitroFusion"
# NitroSD-Realism
ckpt = "nitrosd-realism_unet.safetensors"
unet = UNet2DConditionModel.from_config(base_model_id, subfolder="unet").to("cuda", torch.float16)
unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
scheduler = TimestepShiftLCMScheduler.from_pretrained(base_model_id, subfolder="scheduler", shifted_timestep=250)
scheduler.config.original_inference_steps = 4
# # NitroSD-Vibrant
# ckpt = "nitrosd-vibrant_unet.safetensors"
# unet = UNet2DConditionModel.from_config(base_model_id, subfolder="unet").to("cuda", torch.float16)
# unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
# scheduler = TimestepShiftLCMScheduler.from_pretrained(base_model_id, subfolder="scheduler", shifted_timestep=500)
# scheduler.config.original_inference_steps = 4
pipe = DiffusionPipeline.from_pretrained(
    base_model_id,
    unet=unet,
    scheduler=scheduler,
    torch_dtype=torch.float16,
    variant="fp16",
).to("cuda")
prompt = "a photo of a cat"
image = pipe(
    prompt=prompt,
    num_inference_steps=1,  # NotroSD-Realism and -Vibrant both support 1 - 4 inference steps.
    guidance_scale=0,
).images[0]

ComfyUI使用方法

1. 下載nitrosd - realism_comfyui.safetensors和nitrosd - vibrant_comfyui.safetensors，並將它們放在ComfyUI/models/checkpoints目錄下。
2. 將ComfyUI - TimestepShiftModel倉庫克隆到ComfyUI/custom_nodes目錄下。
3. 嘗試使用工作流！

📄 許可證

• NitroSD - Realism遵循其基礎模型DMD2的許可協議，採用cc - by - nc - 4.0。
• NitroSD - Vibrant採用openrail++許可協議。