nodes/ChenDarYen-ComfyUI-TimestepShiftModel
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

upload

Browse Source
This commit is contained in:
Chen Dar-Yen
2025-01-06 22:54:46 +00:00
parent d4ad3cc6b3
commit 14056d96d8
6 changed files with 551 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -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
ComfyUI_NitroSD_workflow.json
+448
View File
@@ -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
}
README.md
+23 -1
View File
@@ -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)
__init__.py
+2
View File
@@ -0,0 +1,2 @@
from .node import NODE_CLASS_MAPPINGS
__all__ = ['NODE_CLASS_MAPPINGS']
assets/workflow.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 324 KiB

node.py
+77
View File
@@ -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,
}