mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-31 03:17:23 +08:00
174208df6b
Some checks are pending
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run
* Move dataset/text nodes to text category * Rename category utils into utilities * Rename category api node into partner * Move categories conditioning, latent, sampling, model_patches, training, etc. under model category * Dispatch partner nodes in to 3d, audio, image, text, video categories * Move PreviewAny node to utilities category
82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
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/video",
|
|
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 = model_management.get_torch_device()
|
|
memory_required = model_management.module_size(upscale_model)
|
|
|
|
model_dtype = next(upscale_model.parameters()).dtype
|
|
latents = samples["samples"]
|
|
input_dtype = latents.dtype
|
|
|
|
memory_required += math.prod(latents.shape) * 3000.0 # TODO: more accurate
|
|
model_management.free_memory(memory_required, device)
|
|
|
|
try:
|
|
upscale_model.to(device) # TODO: use the comfy model management system.
|
|
|
|
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 = upscale_model(latents)
|
|
finally:
|
|
upscale_model.cpu()
|
|
|
|
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()
|