Merge aa52bc2d34 into 873de5f37a

comfy_extras/nodes_lt_upsampler.py

@@ -70,6 +70,82 @@ class LTXVLatentUpsampler:
         return (return_dict,)
 
 
+def ltxLatentUpscalerBySizeWithModel(model, samples, upscale_method, width, height, crop):
+    if width == 0 and height == 0:
+        return io.NodeOutput(samples)
+    else:
+        if width == 0:
+            height = max(64, height)
+            width = max(64, round(samples.shape[-1] * height / samples.shape[-2]))
+        elif height == 0:
+            width = max(64, width)
+            height = max(64, round(samples.shape[-2] * width / samples.shape[-1]))
+        else:
+            width = max(64, width)
+            height = max(64, height)
+        s = comfy.utils.common_upscale(samples, width // 64, height // 64, upscale_method, crop)
+        s = model(s)
+        
+        return s
+
+
+class LTXVLatentUpsamplerBySize: 
+    methods = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]  
+    options = ["disabled", "center"] 
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"samples": ("LATENT",),
+                     "upscale_method": (s.methods, {"default": "bilinear"}),
+                     "upscale_model": ("LATENT_UPSCALE_MODEL",),
+                     "vae": ("VAE",),
+                     "width": ("INT", {"default": 1280, "min": 0, "max": 16384, "step": 8}),
+                     "height": ("INT", {"default": 720, "min": 0, "max": 16384, "step": 8}),
+                     "crop": (s.options,),
+                    },
+                }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "upsample_latent"
+    CATEGORY = "latent/video"
+    DESCRIPTION = "Upscale latents to the desired size"
+
+    def upsample_latent(cls, samples, upscale_method, upscale_model, vae, width, height, crop) -> tuple:
+        
+        #-------------------------------------------------------------------
+        device = comfy.model_management.get_torch_device()
+        memory_required = comfy.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
+        comfy.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 = ltxLatentUpscalerBySizeWithModel(upscale_model, latents, upscale_method, width, height, crop)
+        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=comfy.model_management.intermediate_device())
+        return_dict = samples.copy()
+        return_dict["samples"] = upsampled_latents
+        return_dict.pop("noise_mask", None)
+
+        return (return_dict,)
+        
+
 NODE_CLASS_MAPPINGS = {
     "LTXVLatentUpsampler": LTXVLatentUpsampler,
+    "LTXVLatentUpsamplerBySize": LTXVLatentUpsamplerBySize,
 }