From e8cc10b9f5717cd1d65d253c25890ac89c05a471 Mon Sep 17 00:00:00 2001
From: hadxp <hadxpk@gmail.com>
Date: Tue, 21 Jul 2026 12:45:53 +0200
Subject: [PATCH] update for torch

---
 Dockerfile                |  2 +-
 diffusers_helper/utils.py | 34 +++++++++++++++++++++++++---------
 requirements.txt          |  2 +-
 3 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 3103733b..0cd5618e 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -16,7 +16,7 @@ COPY requirements.txt .
 # Install dependencies
 RUN pip3 install --no-cache-dir -r requirements.txt
 RUN export CUDA_SHORT_VERSION=$(echo "${CUDA_VERSION}" | sed 's/\.//g' | cut -c 1-3) && \
-    pip3 install --no-cache-dir torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/cu${CUDA_SHORT_VERSION}"
+    pip3 install --no-cache-dir torch torchvision torchaudio torchcodec --index-url "https://download.pytorch.org/whl/cu${CUDA_SHORT_VERSION}"
 
 # Copy the source code to /app
 COPY . .
diff --git a/diffusers_helper/utils.py b/diffusers_helper/utils.py
index ce95d665..dd7e684c 100644
--- a/diffusers_helper/utils.py
+++ b/diffusers_helper/utils.py
@@ -8,6 +8,7 @@
 import numpy as np
 import datetime
 import torchvision
+from torchcodec.encoders import VideoEncoder
 
 from PIL import Image
 
@@ -292,16 +293,31 @@ def save_bcthw_as_mp4(x, output_filename, fps=10, crf=0):
     )
     x = torch.clamp(x.float(), -1.0, 1.0) * 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=per_row)
-    torchvision.io.write_video(
-        output_filename,
-        x,
-        fps=fps,
-        video_codec="libx264",
-        options={"crf": str(int(crf))},
-    )
-    return x
 
+    try:
+        # torchvision wants (T, H, W, C)
+        x_thwc = einops.rearrange(x, "(m n) c t h w -> t (m h) (n w) c", n=per_row)
+        torchvision.io.write_video(
+            output_filename,
+            x_thwc,
+            fps=fps,
+            video_codec="libx264",
+            options={"crf": str(int(crf))},
+        )
+        x = x_thwc
+    except AttributeError:
+        # torchcodec wants (T, C, H, W)
+        x_tchw = einops.rearrange(x, "(m n) c t h w -> t c (m h) (n w)", n=per_row)
+        encoder = VideoEncoder(frames=x_tchw, frame_rate=fps)
+        encoder.to_file(
+            dest=output_filename,
+            codec="libx264",
+            pixel_format="yuv420p",
+            crf=crf,
+        )
+        x = x_tchw
+
+    return x
 
 def save_bcthw_as_png(x, output_filename):
     os.makedirs(
diff --git a/requirements.txt b/requirements.txt
index b0a13e87..696c3aed 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -34,7 +34,7 @@ gfpgan>=1.3.5
 psutil
 realesrgan
 colorlog
-
+torchcodec
 
 # mmaudio
 moviepy