FilesExpand file tree

master

/

nodes_lt_upsampler.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

75 lines (60 loc) · 2.41 KB

master

/

nodes_lt_upsampler.py

Copy path

Top

File metadata and controls

• Code
• Blame

75 lines (60 loc) · 2.41 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

74

75

from comfy import model_management

from comfy_api.latest import ComfyExtension, IO

from typing_extensions import override

import math

class LTXVLatentUpsampler(IO.ComfyNode):

"""

Upsamples a video latent by a factor of 2.

"""

@classmethod

def define_schema(cls):

return IO.Schema(

node_id="LTXVLatentUpsampler",

category="model/latent/ltxv",

is_experimental=True,

inputs=[

IO.Latent.Input("samples"),

IO.LatentUpscaleModel.Input("upscale_model"),

IO.Vae.Input("vae"),

],

outputs=[

IO.Latent.Output(),

],

)

@classmethod

def execute(cls, samples, upscale_model, vae) -> IO.NodeOutput:

"""

Upsample the input latent using the provided model.

Args:

samples (dict): Input latent samples

upscale_model (LatentUpsampler): Loaded upscale model

vae: VAE model for normalization

Returns:

tuple: Tuple containing the upsampled latent

"""

device = upscale_model.load_device

model = upscale_model.model

model_dtype = upscale_model.model_dtype()

latents = samples["samples"]

input_dtype = latents.dtype

memory_required = math.prod(latents.shape) * 3000.0 # TODO: more accurate

model_management.load_models_gpu([upscale_model], memory_required=memory_required)

latents = latents.to(dtype=model_dtype, device=device)

"""Upsample latents without tiling."""

latents = vae.first_stage_model.per_channel_statistics.un_normalize(latents)

upsampled_latents = model(latents)

upsampled_latents = vae.first_stage_model.per_channel_statistics.normalize(

upsampled_latents

)

upsampled_latents = upsampled_latents.to(dtype=input_dtype, device=model_management.intermediate_device())

return_dict = samples.copy()

return_dict["samples"] = upsampled_latents

return_dict.pop("noise_mask", None)

return IO.NodeOutput(return_dict)

upsample_latent = execute # TODO: remove

class LTXVLatentUpsamplerExtension(ComfyExtension):

@override

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

return [LTXVLatentUpsampler]

async def comfy_entrypoint() -> LTXVLatentUpsamplerExtension:

return LTXVLatentUpsamplerExtension()