FilesExpand file tree

master

/

nodes_marigold.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

45 lines (36 loc) · 1.61 KB

master

/

nodes_marigold.py

Copy path

Top

File metadata and controls

• Code
• Blame

45 lines (36 loc) · 1.61 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

from typing_extensions import override

import torch

import torch.nn.functional as F

from comfy_api.latest import ComfyExtension, io

class MarigoldV2PostProcess(io.ComfyNode):

@classmethod

def define_schema(cls) -> io.Schema:

return io.Schema(

node_id="MarigoldV2PostProcess",

display_name="Marigold V2 Post-Process",

category="image/geometry estimation",

description="Turns a decoded Marigold V2 prediction into an image: normalized depth with near as bright, unit surface normals, or sRGB albedo.",

inputs=[

io.Image.Input("image"),

io.Combo.Input("prediction", options=["depth", "normals", "albedo"]),

],

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

)

@classmethod

def execute(cls, image, prediction) -> io.NodeOutput:

if prediction == "normals":

return io.NodeOutput((F.normalize(image * 2.0 - 1.0, dim=-1) + 1.0) * 0.5)

if prediction == "albedo":

return io.NodeOutput(torch.where(image <= 0.0031308, image * 12.92, 1.055 * image.clamp(min=0.0031308) ** (1.0 / 2.4) - 0.055))

d = image.mean(dim=-1, keepdim=True)

lo = d.amin(dim=(1, 2, 3), keepdim=True)

hi = d.amax(dim=(1, 2, 3), keepdim=True)

return io.NodeOutput(((hi - d) / (hi - lo).clamp(min=1e-6)).repeat(1, 1, 1, 3))

class MarigoldExtension(ComfyExtension):

@override

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

return [

MarigoldV2PostProcess,

]

async def comfy_entrypoint() -> MarigoldExtension:

return MarigoldExtension()