download result video & log file in zip

app.py

@@ -1,6 +1,9 @@
+import datetime
 import os
 import sys
+import tempfile
 import time
+import zipfile
 from typing import List, Tuple
 
 import gradio as gr
@@ -11,6 +14,8 @@ from gpu_info import watch_gpu_memory
 PWD = os.path.dirname(__file__)
 CHECKPOINTS_PATH = "/data/checkpoints"
 # CHECKPOINTS_PATH = os.path.join(PWD, "checkpoints")
+LOG_DIR = os.path.join(PWD, "logs")
+os.makedirs(LOG_DIR, exist_ok=True)
 
 try:
     import os
@@ -30,8 +35,8 @@ except Exception as e:
 # download checkpoints
 from download_checkpoints import main as download_checkpoints
 
-# os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
-# download_checkpoints(hf_token="", output_dir=CHECKPOINTS_PATH, model="7b_av")
+os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
+download_checkpoints(hf_token="", output_dir=CHECKPOINTS_PATH, model="7b_av")
 
 
 from test_environment import main as check_environment
@@ -271,6 +276,18 @@ def inference(cfg, control_inputs, chunking) -> Tuple[List[str], List[str]]:
     return video_paths, prompt_paths
 
 
+def create_zip_for_download(filename, files_to_zip):
+    temp_dir = tempfile.mkdtemp()
+    zip_path = os.path.join(temp_dir, f"{os.path.splitext(filename)[0]}.zip")
+
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for file_path in files_to_zip:
+            arcname = os.path.basename(file_path)
+            zipf.write(file_path, arcname)
+
+    return zip_path
+
+
 @spaces.GPU()
 def generate_video(
     rgb_video_path,
@@ -283,6 +300,10 @@ def generate_video(
     chunking=False,
     progress=gr.Progress(track_tqdm=True),
 ):
+    _dt = datetime.datetime.now(tz=datetime.timezone(datetime.timedelta(hours=8))).strftime("%Y-%m-%d_%H.%M.%S")
+    logfile_path = os.path.join(LOG_DIR, f"{_dt}.log")
+    log_handler = log.init_dev_loguru_file(logfile_path)
+
     if randomize_seed:
         actual_seed = random.randint(0, 1000000)
     else:
@@ -290,8 +311,8 @@ def generate_video(
 
     log.info(f"actual_seed: {actual_seed}")
 
-    if not os.path.isfile(rgb_video_path):
-        log.warning(f"File {rgb_video_path} does not exist")
+    if rgb_video_path is None or not os.path.isfile(rgb_video_path):
+        log.warning(f"File `{rgb_video_path}` does not exist")
         rgb_video_path = ""
 
     # add timer to calculate the generation time
@@ -315,7 +336,7 @@ def generate_video(
     )
 
     # watch gpu memory
-    watcher = watch_gpu_memory(10)
+    watcher = watch_gpu_memory(10, lambda x: log.debug(f"GPU memory usage: {x} (MiB)"))
 
     # start inference
     videos, prompts = inference(args, control_inputs, chunking)
@@ -328,7 +349,9 @@ def generate_video(
     watcher.cancel()
 
     video = videos[0]
-    return video, video, actual_seed
+
+    log.logger.remove(log_handler)
+    return video, create_zip_for_download(filename=logfile_path, files_to_zip=[video, logfile_path]), actual_seed
 
 
 # Define the Gradio Blocks interface
@@ -369,7 +392,7 @@ with gr.Blocks() as demo:
 
         with gr.Column():
             output_video = gr.Video(label="Generated Video", format="mp4")
-            output_file = gr.File(label="Download Video")
+            output_file = gr.File(label="Download Results")
 
     generate_button.click(
         fn=generate_video,

cosmos_transfer1/utils/log.py

@@ -76,10 +76,10 @@ def get_machine_format() -> str:
     return machine_format
 
 
-def init_loguru_file(path: str) -> None:
+def init_loguru_file(path: str) -> int:
     machine_format = get_machine_format()
     message_format = get_message_format()
-    logger.add(
+    return logger.add(
         path,
         encoding="utf8",
         level=LEVEL,
@@ -89,6 +89,19 @@ def init_loguru_file(path: str) -> None:
         enqueue=True,
     )
 
+def init_dev_loguru_file(path: str) -> int:
+    machine_format = get_machine_format()
+    message_format = get_message_format()
+    return logger.add(
+        path,
+        encoding="utf8",
+        level="DEBUG",
+        format="[<green>{time:MM-DD HH:mm:ss}</green>|" f"{machine_format}" f"{message_format}",
+        rotation="100 MB",
+        filter=lambda result: _rank0_only_filter(result) or not RANK0_ONLY,
+        enqueue=True,
+    )
+
 
 def get_message_format() -> str:
     message_format = "<level>{level}</level>|<cyan>{extra[relative_path]}:{line}:{function}</cyan>] {message}"