Update app.py

app.py

@@ -5,7 +5,6 @@ import base64
 import uuid
 
 from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
-from diffusers.utils import export_to_video
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 from PIL import Image
@@ -30,6 +29,32 @@ dtype = torch.float16
 pipe = AnimateDiffPipeline.from_pretrained(bases[base_loaded], torch_dtype=dtype).to(device)
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
 
+import tempfile
+import numpy as np
+import cv2
+
+def export_to_video_file(video_frames, output_video_path=None, fps=10):
+    if output_video_path is None:
+        output_video_path = tempfile.NamedTemporaryFile(suffix=".webm").name
+
+    if isinstance(video_frames[0], np.ndarray):
+        video_frames = [(frame * 255).astype(np.uint8) for frame in video_frames]
+    elif isinstance(video_frames[0], Image.Image):
+        video_frames = [np.array(frame) for frame in video_frames]
+
+    # Use VP8 codec
+    fourcc = cv2.VideoWriter_fourcc(*'VP80')
+    h, w, c = video_frames[0].shape
+    video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (w, h), True)
+
+    for frame in video_frames:
+        # Ensure the video frame is in the correct color format
+        img = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+        video_writer.write(img)
+    video_writer.release()
+
+    return output_video_path
+
 def generate_image(secret_token, prompt, base, width, height, motion, step):
     if secret_token != SECRET_TOKEN:
         raise gr.Error(
@@ -69,18 +94,18 @@ def generate_image(secret_token, prompt, base, width, height, motion, step):
     )
 
     name = str(uuid.uuid4()).replace("-", "")
-    path = f"/tmp/{name}.mp4"
+    path = f"/tmp/{name}.webm"
     
-    # I think we are looking time here too, converting to mp4 is too slow, we should return
+    # I think we are looking time here too, converting to webm is too slow, we should return
     # the frames unencoded to the frontend renderer
-    export_to_video(output.frames[0], path, fps=10)
+    export_to_video_file(output.frames[0], path, fps=10)
 
     # Read the content of the video file and encode it to base64
     with open(path, "rb") as video_file:
         video_base64 = base64.b64encode(video_file.read()).decode('utf-8')
 
     # Prepend the appropriate data URI header with MIME type
-    video_data_uri = 'data:video/mp4;base64,' + video_base64
+    video_data_uri = 'data:video/webm;base64,' + video_base64
     
     # clean-up (otherwise there is always a risk of "ghosting", eg. someone seeing the previous generated video",
     # of one of the steps go wrong)
@@ -94,8 +119,8 @@ with gr.Blocks() as demo:
     gr.HTML("""
         <div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
         <div style="text-align: center; color: black;">
-        <p style="color: black;">This space is a REST API to programmatically generate MP4 videos for AiTube, the next generation video platform.</p>
-        <p style="color: black;">Interested in using it? Look no further than the <a href="https://huggingface.co/spaces/ByteDance/AnimateDiff-Lightning" target="_blank">original space</a>!</p>
+        <p style="color: black;">This space is a headless component of the cloud rendering engine used by AiTube.</p>
+        <p style="color: black;">It is not available for public use, but you can use the <a href="https://huggingface.co/spaces/ByteDance/AnimateDiff-Lightning" target="_blank">original space</a>.</p>
         </div>
         </div>""")