ltx2 vae fix

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/ltx/ltx_process.py

@@ -1,10 +1,9 @@
-"""
-- modernui
-- teacache and others
-"""
 import os
 import time
 import threading
+import torch
+from PIL import Image
+
 from modules import shared, errors, timer, memstats, progress, processing, sd_models, sd_samplers, extra_networks
 from modules.video_models.video_save import save_video
 from modules.video_models.video_utils import check_av
@@ -128,6 +127,7 @@ def run_ltx(task_id,
         t0 = time.time()
         shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
         t1 = time.time()
+        output_type = 'np' if 'LTX2' in shared.sd_model.__class__.__name__ else 'latent'
         base_args = {
             "prompt": prompt,
             "negative_prompt": negative,
@@ -137,11 +137,8 @@ def run_ltx(task_id,
             "num_inference_steps": steps,
             "generator": get_generator(seed),
             "callback_on_step_end": diffusers_callback,
+            "output_type": output_type,
         }
-        if 'LTX2' in shared.sd_model.__class__.__name__:
-            base_args["output_type"] = "np"
-        else:
-            base_args["output_type"] = "latent"
         if 'Condition' in shared.sd_model.__class__.__name__:
             base_args["image_cond_noise_scale"] = image_cond_noise_scale
         shared.log.debug(f'Video: cls={shared.sd_model.__class__.__name__} op=base {base_args}')
@@ -175,7 +172,7 @@ def run_ltx(task_id,
                 "width": get_bucket(upsample_ratio * width),
                 "height": get_bucket(upsample_ratio * height),
                 "generator": get_generator(seed),
-                "output_type": "latent",
+                "output_type": output_type,
             }
             if latents.ndim == 4:
                 latents = latents.unsqueeze(0) # add batch dimension
@@ -212,7 +209,7 @@ def run_ltx(task_id,
                 "image_cond_noise_scale": image_cond_noise_scale,
                 "generator": get_generator(seed),
                 "callback_on_step_end": diffusers_callback,
-                "output_type": "latent",
+                "output_type": output_type,
             }
             if latents.ndim == 4:
                 latents = latents.unsqueeze(0) # add batch dimension
@@ -240,7 +237,12 @@ def run_ltx(task_id,
 
         yield None, 'LTX: VAE decode in progress...'
         try:
-            frames = vae_decode(latents, decode_timestep, seed)
+            if torch.is_tensor(latents):
+                frames = vae_decode(latents, decode_timestep, seed)
+            else:
+                frames = latents
+        except TypeError as e:
+            frames = latents # likely because the latents are already decoded
         except AssertionError as e:
             yield from abort(e, ok=True, p=p)
             return
@@ -267,7 +269,14 @@ def run_ltx(task_id,
         )
 
         t_end = time.time()
-        _n, _c, _t, h, w = frames.shape
+        if isinstance(frames, list) and isinstance(frames[0], Image.Image):
+            w, h = frames[0].size
+        elif frames.ndim == 5:
+            _n, _c, _t, h, w = frames.shape
+        elif frames.ndim == 4:
+            _n, h, w, _c = frames.shape
+        else:
+            h, w = frames.shape[-2], frames.shape[-1]
         resolution = f'{w}x{h}' if num_frames > 0 else None
         summary = timer.process.summary(min_time=0.25, total=False).replace('=', ' ')
         memory = shared.mem_mon.summary()

modules/video_models/video_save.py

@@ -62,6 +62,18 @@ def images_to_tensor(images):
     return tensor
 
 
+def numpy_to_tensor(images):
+    if images is None or len(images) == 0:
+        return None
+    array = [torch.from_numpy(images[i]) for i in range(images.shape[0])]
+    tensor = torch.stack(array, dim=0) # n h w c
+    tensor = tensor.unsqueeze(0) # 1, n, h, w, c
+    tensor = tensor.permute(0, 4, 1, 2, 3).contiguous() # 1, c, n, h, w
+    # tensor = (tensor.float() / 127.5) - 1.0 # from [0,255] to [-1,1]
+    # shared.log.debug(f'Video output: images={len(images)} tensor={tensor.shape}')
+    return tensor
+
+
 def atomic_save_video(filename, tensor:torch.Tensor, fps:float=24, codec:str='libx264', pix_fmt:str='yuv420p', options:str='', metadata:dict={}, pbar=None):
     av = check_av()
     if av is None or av is False:
@@ -141,6 +153,10 @@ def save_video(
 
     if pixels is None:
         return 0, output_video
+    if isinstance(pixels, np.ndarray):
+        pixels = numpy_to_tensor(pixels)
+    if isinstance(pixels, list) and isinstance(pixels[0], Image.Image):
+        pixels = images_to_tensor(pixels)
     if not torch.is_tensor(pixels):
         shared.log.error(f'Video: type={type(pixels)} not a tensor')
         return 0, output_video