video tab redesign and optimized ltxvideo

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

modules/framepack/framepack_ui.py

@@ -1,8 +1,9 @@
 import gradio as gr
 from modules import ui_sections, ui_common, ui_video_vlm
+from modules.video_models.video_utils import get_codecs
 from modules.framepack import framepack_load
 from modules.framepack.framepack_worker import get_latent_paddings
-from modules.framepack.framepack_wrappers import get_codecs, load_model, unload_model
+from modules.framepack.framepack_wrappers import load_model, unload_model
 from modules.framepack.framepack_wrappers import run_framepack # pylint: disable=wrong-import-order
 
 

modules/framepack/framepack_video.py

@@ -1,121 +0,0 @@
-import os
-import time
-import datetime
-import cv2
-import torch
-import einops
-from modules import shared, errors ,timer, rife
-
-
-def atomic_save_video(filename, tensor:torch.Tensor, fps:float=24, codec:str='libx264', pix_fmt:str='yuv420p', options:str='', metadata:dict={}, pbar=None):
-    try:
-        import av
-        av.logging.set_level(av.logging.ERROR) # pylint: disable=c-extension-no-member
-    except Exception as e:
-        shared.log.error(f'FramePack video: {e}')
-        return
-
-    frames, height, width, _channels = tensor.shape
-    rate = round(fps)
-    options_str = options
-    options = {}
-    for option in [option.strip() for option in options_str.split(',')]:
-        if '=' in option:
-            key, value = option.split('=', 1)
-        elif ':' in option:
-            key, value = option.split(':', 1)
-        else:
-            continue
-        options[key.strip()] = value.strip()
-    shared.log.info(f'FramePack video: file="{filename}" codec={codec} frames={frames} width={width} height={height} fps={rate} options={options}')
-    video_array = torch.as_tensor(tensor, dtype=torch.uint8).numpy(force=True)
-    task = pbar.add_task('encoding', total=frames) if pbar is not None else None
-    if task is not None:
-        pbar.update(task, description='video encoding')
-    with av.open(filename, mode="w") as container:
-        for k, v in metadata.items():
-            container.metadata[k] = v
-        stream: av.VideoStream = container.add_stream(codec, rate=rate, options=options)
-        stream.width = video_array.shape[2]
-        stream.height = video_array.shape[1]
-        stream.pix_fmt = pix_fmt
-        for img in video_array:
-            frame = av.VideoFrame.from_ndarray(img, format="rgb24")
-            for packet in stream.encode_lazy(frame):
-                container.mux(packet)
-            if task is not None:
-                pbar.update(task, advance=1)
-        for packet in stream.encode(): # flush
-            container.mux(packet)
-    shared.state.outputs(filename)
-
-
-def save_video(
-        pixels:torch.Tensor,
-        mp4_fps:int=24,
-        mp4_codec:str='libx264',
-        mp4_opt:str='',
-        mp4_ext:str='mp4',
-        mp4_sf:bool=False, # save safetensors
-        mp4_video:bool=True, # save video
-        mp4_frames:bool=False, # save frames
-        mp4_interpolate:int=0, # rife interpolation
-        stream=None, # async progress reporting stream
-        metadata:dict={}, # metadata for video
-        pbar=None, # progress bar for video
-    ):
-    if pixels is None:
-        return 0
-    t_save = time.time()
-    n, _c, t, h, w = pixels.shape
-    size = pixels.element_size() * pixels.numel()
-    shared.log.debug(f'FramePack video: video={mp4_video} export={mp4_frames} safetensors={mp4_sf} interpolate={mp4_interpolate}')
-    shared.log.debug(f'FramePack video: encode={t} raw={size} latent={pixels.shape} fps={mp4_fps} codec={mp4_codec} ext={mp4_ext} options="{mp4_opt}"')
-    try:
-        if stream is not None:
-            stream.output_queue.push(('progress', (None, 'Saving video...')))
-        if mp4_interpolate > 0:
-            x = pixels.squeeze(0).permute(1, 0, 2, 3)
-            interpolated = rife.interpolate_nchw(x, count=mp4_interpolate+1)
-            pixels = torch.stack(interpolated, dim=0)
-            pixels = pixels.permute(1, 2, 0, 3, 4)
-
-        n, _c, t, h, w = pixels.shape
-        x = torch.clamp(pixels.float(), -1., 1.) * 127.5 + 127.5
-        x = x.detach().cpu().to(torch.uint8)
-        x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=n)
-        x = x.contiguous()
-
-        timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
-        output_filename = os.path.join(shared.opts.outdir_video, f'{timestamp}-{mp4_codec}-f{t}')
-
-        if mp4_sf:
-            fn = f'{output_filename}.safetensors'
-            shared.log.info(f'FramePack export: file="{fn}" type=savetensors shape={x.shape}')
-            from safetensors.torch import save_file
-            shared.state.outputs(fn)
-            save_file({ 'frames': x }, fn, metadata={'format': 'video', 'frames': str(t), 'width': str(w), 'height': str(h), 'fps': str(mp4_fps), 'codec': mp4_codec, 'options': mp4_opt, 'ext': mp4_ext, 'interpolate': str(mp4_interpolate)})
-
-        if mp4_frames:
-            shared.log.info(f'FramePack frames: files="{output_filename}-00000.jpg" frames={t} width={w} height={h}')
-            for i in range(t):
-                image = cv2.cvtColor(x[i].numpy(), cv2.COLOR_RGB2BGR)
-                fn = f'{output_filename}-{i:05d}.jpg'
-                shared.state.outputs(fn)
-                cv2.imwrite(fn, image)
-
-        if mp4_video and (mp4_codec != 'none'):
-            fn = f'{output_filename}.{mp4_ext}'
-            atomic_save_video(fn, tensor=x, fps=mp4_fps, codec=mp4_codec, options=mp4_opt, metadata=metadata, pbar=pbar)
-            if stream is not None:
-                stream.output_queue.push(('progress', (None, f'Video {os.path.basename(fn)} | Codec {mp4_codec} | Size {w}x{h}x{t} | FPS {mp4_fps}')))
-                stream.output_queue.push(('file', fn))
-        else:
-            if stream is not None:
-                stream.output_queue.push(('progress', (None, '')))
-
-    except Exception as e:
-        shared.log.error(f'FramePack video: raw={size} {e}')
-        errors.display(e, 'FramePack video')
-    timer.process.add('save', time.time()-t_save)
-    return t

modules/framepack/framepack_worker.py

@@ -4,7 +4,7 @@ import rich.progress as rp
 from modules import shared, errors ,devices, sd_models, timer, memstats
 from modules.framepack import framepack_vae # pylint: disable=wrong-import-order
 from modules.framepack import framepack_hijack # pylint: disable=wrong-import-order
-from modules.framepack import framepack_video # pylint: disable=wrong-import-order
+from modules.video_models.video_save import save_video # pylint: disable=wrong-import-order
 
 
 stream = None # AsyncStream
@@ -302,12 +302,12 @@ def worker(
                 if is_last_section:
                     break
 
-            total_generated_frames = framepack_video.save_video(history_pixels, mp4_fps, mp4_codec, mp4_opt, mp4_ext, mp4_sf, mp4_video, mp4_frames, mp4_interpolate, pbar=pbar, stream=stream, metadata=metadata)
+            total_generated_frames, _video_filename = save_video(history_pixels, mp4_fps, mp4_codec, mp4_opt, mp4_ext, mp4_sf, mp4_video, mp4_frames, mp4_interpolate, pbar=pbar, stream=stream, metadata=metadata)
 
     except AssertionError:
         shared.log.info('FramePack: interrupted')
         if shared.opts.keep_incomplete:
-            framepack_video.save_video(history_pixels, mp4_fps, mp4_codec, mp4_opt, mp4_ext, mp4_sf, mp4_video, mp4_frames, mp4_interpolate=0, stream=stream, metadata=metadata)
+            save_video(history_pixels, mp4_fps, mp4_codec, mp4_opt, mp4_ext, mp4_sf, mp4_video, mp4_frames, mp4_interpolate=0, stream=stream, metadata=metadata)
     except Exception as e:
         shared.log.error(f'FramePack: {e}')
         errors.display(e, 'FramePack')

modules/framepack/framepack_wrappers.py

@@ -6,6 +6,7 @@ import numpy as np
 import torch
 import gradio as gr
 from modules import shared, processing, timer, paths, extra_networks, progress, ui_video_vlm
+from modules.video_models.video_utils import check_av
 from modules.framepack import framepack_install # pylint: disable=wrong-import-order
 from modules.framepack import framepack_load # pylint: disable=wrong-import-order
 from modules.framepack import framepack_worker # pylint: disable=wrong-import-order
@@ -20,38 +21,6 @@ queue_lock = threading.Lock()
 loaded_variant = None
 
 
-def check_av():
-    try:
-        import av
-    except Exception as e:
-        shared.log.error(f'av package: {e}')
-        return False
-    return av
-
-
-def get_codecs():
-    av = check_av()
-    if av is None:
-        return []
-    codecs = []
-    for codec in av.codecs_available:
-        try:
-            c = av.Codec(codec, mode='w')
-            if c.type == 'video' and c.is_encoder and len(c.video_formats) > 0:
-                if not any(c.name == ca.name for ca in codecs):
-                    codecs.append(c)
-        except Exception:
-            pass
-    hw_codecs = [c for c in codecs if (c.capabilities & 0x40000 > 0) or (c.capabilities & 0x80000 > 0)]
-    sw_codecs = [c for c in codecs if c not in hw_codecs]
-    shared.log.debug(f'Video codecs: hardware={len(hw_codecs)} software={len(sw_codecs)}')
-    # for c in hw_codecs:
-    #     shared.log.trace(f'codec={c.name} cname="{c.canonical_name}" decs="{c.long_name}" intra={c.intra_only} lossy={c.lossy} lossless={c.lossless} capabilities={c.capabilities} hw=True')
-    # for c in sw_codecs:
-    #     shared.log.trace(f'codec={c.name} cname="{c.canonical_name}" decs="{c.long_name}" intra={c.intra_only} lossy={c.lossy} lossless={c.lossless} capabilities={c.capabilities} hw=False')
-    return ['none'] + [c.name for c in hw_codecs + sw_codecs]
-
-
 def prepare_image(image, resolution):
     from modules.framepack.pipeline.utils import resize_and_center_crop
     buckets = [