🚀 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++许可协议。