NukeDepthCrafter/nukedc.py at b79dcefdee3c9c872723a90ef2688cfd33c03022 · Theo-SAMINADIN-td/NukeDepthCrafter · GitHub

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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import nuke
from DepthCrafterPlugin.utils import *
from diffusers.training_utils import set_seed


def getInputInfos():
    f = nuke.thisNode().dependencies()

    try :

        for i in f:
            getInputInfos.read = i


        try :
            getInputInfos.FrameNumber = getInputInfos.read.knob('last').getValue()
        except :
            getInputInfos.FrameNumber = nuke.root().knob('last_frame').getValue()


        try :
            getInputInfos.path = getInputInfos.read.knob('file').getValue()
        except :
            getInputInfos.path = ''
    except :
        getInputInfos.FrameNumber = nuke.root().knob('last_frame').getValue()
        getInputInfos.path = ''


def UpdateBtn():
    getInputInfos()

    FilePath = getInputInfos.path

    nuke.thisNode().knob('FilePath').setValue(FilePath)


def GenerateDepthAction():

    if (nuke.thisNode().knob('FileType').value() == "mp4") :
        VideoExportBool = 1
    else :
        VideoExportBool = 0
    depthcrafter_demo = DepthCrafterDemo(
        unet_path=os.path.expandvars(r"C:\Users\$USERNAME\.nuke\DepthCrafterPlugin"),
        pre_train_path="stabilityai/stable-video-diffusion-img2vid-xt",
        cpu_offload=nuke.thisNode().knob('CPUOFF_OPT').value(),
    )
    # process the videos, the video paths are separated by comma
    video_paths = [nuke.thisNode().knob('FilePath').getValue()]
    #args.video_path.split(",")
    for video in video_paths:
        depthcrafter_demo.infer(
            video,
            nuke.thisNode().knob('InferSteps').value(),                    # args.num_inference_steps,
            nuke.thisNode().knob('CFG').value(),                    # args.guidance_scale,
            nuke.thisNode().knob('OutputPath').getValue(),       #args.save_folder,
            window_size= 110,       #args.window_size,
            process_length=nuke.thisNode().knob('FrameNumber').value(),    #args.process_length,
            overlap= 25,              #args.overlap,
            max_res=nuke.thisNode().knob('MaxRes').value(),              #args.max_res,
            target_fps=nuke.thisNode().knob('FPS').value(),           #args.target_fps,
            seed= 42  ,                 #args.seed,
            track_time= False   ,    #args.track_time,
            save_npz= False,          #args.save_npz,
            video_export=VideoExportBool
        )

        # clear the cache for the next video
        gc.collect()
        torch.cuda.empty_cache()
        nuke.message('Render ' + video + ' Done')


def CreateDCNode():
    getInputInfos()
    nuke.createNode('NoOp')
    s = nuke.selectedNode()
    s.knob('name').setValue('DepthCrafter')
    s.addKnob(nuke.File_Knob('FilePath', 'File Path'))
    s.addKnob(nuke.PyScript_Knob('UpdatePath', 'Update Path', 'UpdateBtn()' ))

    s.addKnob(nuke.Text_Knob(''))

    s.addKnob(nuke.Enumeration_Knob('CPUOFF_OPT', 'CPU Offload Options', ['model', 'sequential', 'none']))
    s.addKnob(nuke.Int_Knob("FPS", 'Output Frame Rate'))
    s.addKnob(nuke.Int_Knob("InferSteps", 'Inference Steps')) #NEED TO CREATE FUNCTION TO ROUND UP
    s.addKnob(nuke.Double_Knob("CFG", 'Guidance scale'))
    s.addKnob(nuke.Int_Knob("FrameNumber", 'Number of frame'))
    s.addKnob(nuke.Int_Knob("MaxRes", 'Maximum Resolution'))

    s.addKnob(nuke.Text_Knob(' ', ''))

    s.addKnob(nuke.Enumeration_Knob('FileType', 'File type', ['exr', 'mp4']))

    s.addKnob(nuke.File_Knob('OutputPath', 'Output Path'))
    s.addKnob(nuke.PyScript_Knob('GenerateDepth', 'Generate Depth', 'GenerateDepthAction()'))


### SETTING RANGES, DEFAULT VALUES, TOOLTIP & FORMATING ###
    s['FPS'].setValue(int(nuke.root().knob('fps').getValue())) #ADD ROOT FPS BY DEFAULT
    s['InferSteps'].setValue(25)
    s['CFG'].setValue(1.2)
    s['FrameNumber'].setValue(int(getInputInfos.FrameNumber))
    s['MaxRes'].setValue(1024)

    s['InferSteps'].setRange(1, 40)
    s['CFG'].setRange(1, 20)


    s['CPUOFF_OPT'].setFlag(nuke.STARTLINE)
    s['FPS'].setFlag(nuke.STARTLINE)
    s['InferSteps'].setFlag(nuke.STARTLINE)
    s['CFG'].setFlag(nuke.STARTLINE)
    s['FrameNumber'].setFlag(nuke.STARTLINE)
    s['MaxRes'].setFlag(nuke.STARTLINE)
    s['UpdatePath'].setFlag(nuke.STARTLINE)
    s['GenerateDepth'].setFlag(nuke.STARTLINE)

    s['CPUOFF_OPT'].setTooltip("To save memory, we can offload the model to CPU. Model is the default one, Sequential will be slower but save more memory")
    s['FPS'].setTooltip("Target FPS for the output video")
    s['InferSteps'].setTooltip("Number of inference steps")
    s['CFG'].setTooltip("Guidance scale/CFG")
    s['FrameNumber'].setTooltip("Number of frame to generate")
    s['MaxRes'].setTooltip("Output resolution")
    s['OutputPath'].setTooltip("path/to/your/folder/ Extension and filename are automaticly set to the selected file type and input file name ")
    s['GenerateDepth'].setTooltip("Generate Depth")

    print(nuke.thisNode().allKnobs())