diff --git a/.gitignore b/.gitignore index 15201ac..d86e4ff 100644 --- a/.gitignore +++ b/.gitignore @@ -165,7 +165,7 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +.idea/ # PyPI configuration file .pypirc diff --git a/ComfyUI_NitroSD_workflow.json b/ComfyUI_NitroSD_workflow.json new file mode 100644 index 0000000..7fa4e0e --- /dev/null +++ b/ComfyUI_NitroSD_workflow.json @@ -0,0 +1,448 @@ +{ + "last_node_id": 9, + "last_link_id": 10, + "nodes": [ + { + "id": 5, + "type": "CLIPTextEncode", + "pos": [ + -40, + 30 + ], + "size": [ + 400, + 76 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 4 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 5 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "" + ] + }, + { + "id": 1, + "type": "CheckpointLoaderSimple", + "pos": [ + -800, + 30 + ], + "size": [ + 315, + 98 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 1 + ], + "slot_index": 0 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 3, + 4 + ], + "slot_index": 1 + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 9 + ], + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "nitrosd-realism_comfyui.safetensors" + ] + }, + { + "id": 4, + "type": "CLIPTextEncode", + "pos": [ + -40, + -230 + ], + "size": [ + 400, + 200 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 3 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 6 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "Breathtaking 8k photograph of Patagonian peaks at sunset. Dynamic composition with foreground wildflowers." + ] + }, + { + "id": 2, + "type": "Timestep Shift Model", + "pos": [ + -420, + 10 + ], + "size": [ + 320, + 60 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 2 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Timestep Shift Model" + }, + "widgets_values": [ + 250 + ] + }, + { + "id": 3, + "type": "MarkdownNote", + "pos": [ + -420, + 130 + ], + "size": [ + 320, + 100 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Value of shifted_timestep", + "properties": {}, + "widgets_values": [ + "`nitrosd-realism_comfyui.safetensors` use `shifted_timestep` 250.\n\n`nitrosd-vibrant_comfyui.safetensors` use `shifted_timestep` 500." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 6, + "type": "EmptyLatentImage", + "pos": [ + 40, + 170 + ], + "size": [ + 320, + 110 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 7 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage" + }, + "widgets_values": [ + 1024, + 1024, + 1 + ] + }, + { + "id": 8, + "type": "VAEDecode", + "pos": [ + 820, + -170 + ], + "size": [ + 210, + 46 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 8 + }, + { + "name": "vae", + "type": "VAE", + "link": 9 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 10 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + }, + "widgets_values": [] + }, + { + "id": 9, + "type": "SaveImage", + "pos": [ + 1090, + -250 + ], + "size": [ + 380, + 540 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 10 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 7, + "type": "KSampler", + "pos": [ + 440, + -170 + ], + "size": [ + 315, + 262 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 2 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 6 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 5 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 7 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 8 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 2025, + "randomize", + 1, + 1, + "lcm", + "normal", + 1 + ] + } + ], + "links": [ + [ + 1, + 1, + 0, + 2, + 0, + "MODEL" + ], + [ + 2, + 2, + 0, + 7, + 0, + "MODEL" + ], + [ + 3, + 1, + 1, + 4, + 0, + "CLIP" + ], + [ + 4, + 1, + 1, + 5, + 0, + "CLIP" + ], + [ + 5, + 5, + 0, + 7, + 2, + "CONDITIONING" + ], + [ + 6, + 4, + 0, + 7, + 1, + "CONDITIONING" + ], + [ + 7, + 6, + 0, + 7, + 3, + "LATENT" + ], + [ + 8, + 7, + 0, + 8, + 0, + "LATENT" + ], + [ + 9, + 1, + 2, + 8, + 1, + "VAE" + ], + [ + 10, + 8, + 0, + 9, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.6018254139530037, + "offset": [ + 948.422224215607, + 431.8031367211276 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/README.md b/README.md index f983c3c..e9b116a 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,24 @@ # ComfyUI-TimestepShiftModel -ComfyUI implemtation for timestep shift used in NitroSD +This is a ComfyUI implementation of the timestep shift technique used in [NitroFusion: High-Fidelity Single-Step Diffusion through Dynamic Adversarial Training](https://arxiv.org/abs/2412.02030). + +For more details, visit the [official NitroFusion GitHub repository](https://github.com/ChenDarYen/NitroFusion). + +## Usage +Clone this repository into the `ComfyUI/custom_nodes` directory: +```bash +git clone https://github.com/ChenDarYen/ComfyUI-TimestepShiftModel.git +``` + +Download the NitroSD models from [Hugging Face](https://huggingface.co/ChenDY/NitroFusion) or use the following commands: +```bash +wget https://huggingface.co/ChenDY/NitroFusion/resolve/main/nitrosd-realism_comfyui.safetensors +wget https://huggingface.co/ChenDY/NitroFusion/resolve/main/nitrosd-vibrant_comfyui.safetensors +``` + +The `Timestep Shift Model` node takes a model and a `shifted_timestep` value as input, producing a timestep-shifted model. + +Apart from this step, the workflow is the same as a standard text-to-image generation workflow. + +Have fun with the example workflow [ComfyUI_NitroSD_workflow.json](./ComfyUI_NitroSD_workflow.json)! + +![](./assets/workflow.png) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..0b56e61 --- /dev/null +++ b/__init__.py @@ -0,0 +1,2 @@ +from .node import NODE_CLASS_MAPPINGS +__all__ = ['NODE_CLASS_MAPPINGS'] \ No newline at end of file diff --git a/assets/workflow.png b/assets/workflow.png new file mode 100644 index 0000000..5fa8694 Binary files /dev/null and b/assets/workflow.png differ diff --git a/node.py b/node.py new file mode 100644 index 0000000..60bb617 --- /dev/null +++ b/node.py @@ -0,0 +1,77 @@ +from types import MethodType +from functools import partial + +import torch +from comfy.model_base import BaseModel + + +def apply_model_with_shifted_timestep( + self: BaseModel, + x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, + shifted_timestep: int = None, + **kwargs, +): + sigma = t + xc = self.model_sampling.calculate_input(sigma, x) + if c_concat is not None: + xc = torch.cat([xc] + [c_concat], dim=1) + + context = c_crossattn + dtype = self.get_dtype() + + if self.manual_cast_dtype is not None: + dtype = self.manual_cast_dtype + + xc = xc.to(dtype) + if shifted_timestep is None: + t = self.model_sampling.timestep(t).float() + else: + num_train_timesteps = len(self.model_sampling.log_sigmas) + t = (self.model_sampling.timestep(t) * (shifted_timestep / num_train_timesteps)).long() + + context = context.to(dtype) + extra_conds = {} + for o in kwargs: + extra = kwargs[o] + if hasattr(extra, "dtype"): + if extra.dtype != torch.int and extra.dtype != torch.long: + extra = extra.to(dtype) + extra_conds[o] = extra + + model_output = self.diffusion_model(xc, t, context=context, control=control, + transformer_options=transformer_options, **extra_conds).float() + + if shifted_timestep is None: + return self.model_sampling.calculate_denoised(sigma, model_output, x) + + denoised_sigma = self.model_sampling.sigma(t) + denoised_sigma = denoised_sigma.view(denoised_sigma.shape[:1] + (1,) * (x.ndim - 1)) + x = xc * ((denoised_sigma ** 2 + self.model_sampling.sigma_data ** 2) ** 0.5) + return self.model_sampling.calculate_denoised(denoised_sigma, model_output, x) + + + +class TimestepShiftModel: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "model": ("MODEL",), + "shifted_timestep": ("INT", {"default": 250, "min": 1, "max": 1000}), + } + } + RETURN_TYPES = ("MODEL",) + CATEGORY = "test" + FUNCTION = "shift_model_timestep" + + def shift_model_timestep(self, model, shifted_timestep): + model.model._apply_model = MethodType( + partial(apply_model_with_shifted_timestep, shifted_timestep=shifted_timestep), + model.model, + ) + return (model, ) + + +NODE_CLASS_MAPPINGS = { + "Timestep Shift Model": TimestepShiftModel, +} \ No newline at end of file