FilesExpand file tree

master

/

nodes_differential_diffusion.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

73 lines (59 loc) · 2.47 KB

master

/

nodes_differential_diffusion.py

Copy path

Top

File metadata and controls

• Code
• Blame

73 lines (59 loc) · 2.47 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

# code adapted from https://github.com/exx8/differential-diffusion

from typing_extensions import override

import torch

from comfy_api.latest import ComfyExtension, io

class DifferentialDiffusion(io.ComfyNode):

@classmethod

def define_schema(cls):

return io.Schema(

node_id="DifferentialDiffusion",

search_aliases=["inpaint gradient", "variable denoise strength"],

display_name="Differential Diffusion",

category="experimental",

inputs=[

io.Model.Input("model"),

io.Float.Input(

"strength",

default=1.0,

min=0.0,

max=1.0,

step=0.01,

optional=True,

),

],

outputs=[io.Model.Output()],

is_experimental=True,

)

@classmethod

def execute(cls, model, strength=1.0) -> io.NodeOutput:

model = model.clone()

model.set_model_denoise_mask_function(lambda *args, **kwargs: cls.forward(*args, **kwargs, strength=strength))

return io.NodeOutput(model)

@classmethod

def forward(cls, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict, strength: float):

model = extra_options["model"]

step_sigmas = extra_options["sigmas"]

sigma_to = model.inner_model.model_sampling.sigma_min

if step_sigmas[-1] > sigma_to:

sigma_to = step_sigmas[-1]

sigma_from = step_sigmas[0]

ts_from = model.inner_model.model_sampling.timestep(sigma_from)

ts_to = model.inner_model.model_sampling.timestep(sigma_to)

current_ts = model.inner_model.model_sampling.timestep(sigma[0])

threshold = (current_ts - ts_to) / (ts_from - ts_to)

# Generate the binary mask based on the threshold

binary_mask = (denoise_mask >= threshold).to(denoise_mask.dtype)

# Blend binary mask with the original denoise_mask using strength

if strength and strength < 1:

blended_mask = strength * binary_mask + (1 - strength) * denoise_mask

return blended_mask

else:

return binary_mask

class DifferentialDiffusionExtension(ComfyExtension):

@override

async def get_node_list(self) -> list[type[io.ComfyNode]]:

return [

DifferentialDiffusion,

]

async def comfy_entrypoint() -> DifferentialDiffusionExtension:

return DifferentialDiffusionExtension()