move the inference function to `helper.py`

app.py

@@ -1,24 +1,6 @@
-import datetime
 import os
-import sys
-import tempfile
-import time
-import zipfile
-from typing import List, Tuple
-
-import gradio as gr
-import spaces
-
-from gpu_info import stop_watcher, watch_gpu_memory
-
-PWD = os.path.dirname(__file__)
-CHECKPOINTS_PATH = "/data/checkpoints"
-LOG_DIR = os.path.join(PWD, "logs")
-os.makedirs(LOG_DIR, exist_ok=True)
 
 try:
-    import os
-
     from huggingface_hub import login
 
     # Try to login with token from environment variable
@@ -31,12 +13,8 @@ try:
 except Exception as e:
     print(f"Authentication failed: {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")
 
+import sys
 
 from test_environment import main as check_environment
 from test_environment import setup_environment
@@ -46,314 +24,21 @@ setup_environment()
 # setup env
 os.environ["CUDA_HOME"] = "/usr/local/cuda"
 os.environ["LD_LIBRARY_PATH"] = "$CUDA_HOME/lib:$CUDA_HOME/lib64:$LD_LIBRARY_PATH"
-os.environ["PATH"] = "$CUDA_HOME/bin:/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:$PATH"
+os.environ["PATH"] = "$CUDA_HOME/bin:/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:$PATH"
 
 if not check_environment():
     sys.exit(1)
 
-os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Workaround to suppress MP warning
-
-import copy
-import json
-import random
-from io import BytesIO
-
-import torch
-
-from cosmos_transfer1.checkpoints import (
-    BASE_7B_CHECKPOINT_AV_SAMPLE_PATH,
-    BASE_7B_CHECKPOINT_PATH,
-    EDGE2WORLD_CONTROLNET_DISTILLED_CHECKPOINT_PATH,
-)
-from cosmos_transfer1.diffusion.inference.inference_utils import (
-    validate_controlnet_specs,
-)
-from cosmos_transfer1.diffusion.inference.preprocessors import Preprocessors
-from cosmos_transfer1.diffusion.inference.world_generation_pipeline import (
-    DiffusionControl2WorldGenerationPipeline,
-    DistilledControl2WorldGenerationPipeline,
-)
-from cosmos_transfer1.utils import log, misc
-from cosmos_transfer1.utils.io import read_prompts_from_file, save_video
-from helper import parse_arguments
-
-torch.enable_grad(False)
-torch.serialization.add_safe_globals([BytesIO])
-
-
-def inference(cfg, control_inputs, chunking) -> Tuple[List[str], List[str]]:
-    video_paths = []
-    prompt_paths = []
-
-    control_inputs = validate_controlnet_specs(cfg, control_inputs)
-    misc.set_random_seed(cfg.seed)
-
-    device_rank = 0
-    process_group = None
-    if cfg.num_gpus > 1:
-        from megatron.core import parallel_state
-
-        from cosmos_transfer1.utils import distributed
-
-        distributed.init()
-        parallel_state.initialize_model_parallel(context_parallel_size=cfg.num_gpus)
-        process_group = parallel_state.get_context_parallel_group()
-
-        device_rank = distributed.get_rank(process_group)
-
-    preprocessors = Preprocessors()
-
-    if cfg.use_distilled:
-        assert not cfg.is_av_sample
-        checkpoint = EDGE2WORLD_CONTROLNET_DISTILLED_CHECKPOINT_PATH
-        pipeline = DistilledControl2WorldGenerationPipeline(
-            checkpoint_dir=cfg.checkpoint_dir,
-            checkpoint_name=checkpoint,
-            offload_network=cfg.offload_diffusion_transformer,
-            offload_text_encoder_model=cfg.offload_text_encoder_model,
-            offload_guardrail_models=cfg.offload_guardrail_models,
-            guidance=cfg.guidance,
-            num_steps=cfg.num_steps,
-            fps=cfg.fps,
-            seed=cfg.seed,
-            num_input_frames=cfg.num_input_frames,
-            control_inputs=control_inputs,
-            sigma_max=cfg.sigma_max,
-            blur_strength=cfg.blur_strength,
-            canny_threshold=cfg.canny_threshold,
-            upsample_prompt=cfg.upsample_prompt,
-            offload_prompt_upsampler=cfg.offload_prompt_upsampler,
-            process_group=process_group,
-        )
-    else:
-        checkpoint = BASE_7B_CHECKPOINT_AV_SAMPLE_PATH if cfg.is_av_sample else BASE_7B_CHECKPOINT_PATH
-
-        # Initialize transfer generation model pipeline
-        pipeline = DiffusionControl2WorldGenerationPipeline(
-            checkpoint_dir=cfg.checkpoint_dir,
-            checkpoint_name=checkpoint,
-            offload_network=cfg.offload_diffusion_transformer,
-            offload_text_encoder_model=cfg.offload_text_encoder_model,
-            offload_guardrail_models=cfg.offload_guardrail_models,
-            guidance=cfg.guidance,
-            num_steps=cfg.num_steps,
-            fps=cfg.fps,
-            seed=cfg.seed,
-            num_input_frames=cfg.num_input_frames,
-            control_inputs=control_inputs,
-            sigma_max=cfg.sigma_max,
-            blur_strength=cfg.blur_strength,
-            canny_threshold=cfg.canny_threshold,
-            upsample_prompt=cfg.upsample_prompt,
-            offload_prompt_upsampler=cfg.offload_prompt_upsampler,
-            process_group=process_group,
-            chunking=chunking,
-        )
-
-    if cfg.batch_input_path:
-        log.info(f"Reading batch inputs from path: {cfg.batch_input_path}")
-        prompts = read_prompts_from_file(cfg.batch_input_path)
-    else:
-        # Single prompt case
-        prompts = [{"prompt": cfg.prompt, "visual_input": cfg.input_video_path}]
-
-    batch_size = cfg.batch_size if hasattr(cfg, "batch_size") else 1
-    if any("upscale" in control_input for control_input in control_inputs) and batch_size > 1:
-        batch_size = 1
-        log.info("Setting batch_size=1 as upscale does not support batch generation")
-    os.makedirs(cfg.video_save_folder, exist_ok=True)
-    for batch_start in range(0, len(prompts), batch_size):
-        # Get current batch
-        batch_prompts = prompts[batch_start : batch_start + batch_size]
-        actual_batch_size = len(batch_prompts)
-        # Extract batch data
-        batch_prompt_texts = [p.get("prompt", None) for p in batch_prompts]
-        batch_video_paths = [p.get("visual_input", None) for p in batch_prompts]
-
-        batch_control_inputs = []
-        for i, input_dict in enumerate(batch_prompts):
-            current_prompt = input_dict.get("prompt", None)
-            current_video_path = input_dict.get("visual_input", None)
-
-            if cfg.batch_input_path:
-                video_save_subfolder = os.path.join(cfg.video_save_folder, f"video_{batch_start+i}")
-                os.makedirs(video_save_subfolder, exist_ok=True)
-            else:
-                video_save_subfolder = cfg.video_save_folder
-
-            current_control_inputs = copy.deepcopy(control_inputs)
-            if "control_overrides" in input_dict:
-                for hint_key, override in input_dict["control_overrides"].items():
-                    if hint_key in current_control_inputs:
-                        current_control_inputs[hint_key].update(override)
-                    else:
-                        log.warning(f"Ignoring unknown control key in override: {hint_key}")
 
-            # if control inputs are not provided, run respective preprocessor (for seg and depth)
-            log.info("running preprocessor")
-            preprocessors(
-                current_video_path,
-                current_prompt,
-                current_control_inputs,
-                video_save_subfolder,
-                cfg.regional_prompts if hasattr(cfg, "regional_prompts") else None,
-            )
-            batch_control_inputs.append(current_control_inputs)
-
-        regional_prompts = []
-        region_definitions = []
-        if hasattr(cfg, "regional_prompts") and cfg.regional_prompts:
-            log.info(f"regional_prompts: {cfg.regional_prompts}")
-            for regional_prompt in cfg.regional_prompts:
-                regional_prompts.append(regional_prompt["prompt"])
-                if "region_definitions_path" in regional_prompt:
-                    log.info(f"region_definitions_path: {regional_prompt['region_definitions_path']}")
-                    region_definition_path = regional_prompt["region_definitions_path"]
-                    if isinstance(region_definition_path, str) and region_definition_path.endswith(".json"):
-                        with open(region_definition_path, "r") as f:
-                            region_definitions_json = json.load(f)
-                        region_definitions.extend(region_definitions_json)
-                    else:
-                        region_definitions.append(region_definition_path)
-
-        if hasattr(pipeline, "regional_prompts"):
-            pipeline.regional_prompts = regional_prompts
-        if hasattr(pipeline, "region_definitions"):
-            pipeline.region_definitions = region_definitions
-
-        # Generate videos in batch
-        batch_outputs = pipeline.generate(
-            prompt=batch_prompt_texts,
-            video_path=batch_video_paths,
-            negative_prompt=cfg.negative_prompt,
-            control_inputs=batch_control_inputs,
-            save_folder=video_save_subfolder,
-            batch_size=actual_batch_size,
-        )
-        if batch_outputs is None:
-            log.critical("Guardrail blocked generation for entire batch.")
-            continue
-
-        videos, final_prompts = batch_outputs
-        for i, (video, prompt) in enumerate(zip(videos, final_prompts)):
-            if cfg.batch_input_path:
-                video_save_subfolder = os.path.join(cfg.video_save_folder, f"video_{batch_start+i}")
-                video_save_path = os.path.join(video_save_subfolder, "output.mp4")
-                prompt_save_path = os.path.join(video_save_subfolder, "prompt.txt")
-            else:
-                video_save_path = os.path.join(cfg.video_save_folder, f"{cfg.video_save_name}.mp4")
-                prompt_save_path = os.path.join(cfg.video_save_folder, f"{cfg.video_save_name}.txt")
-            # Save video and prompt
-            if device_rank == 0:
-                os.makedirs(os.path.dirname(video_save_path), exist_ok=True)
-                save_video(
-                    video=video,
-                    fps=cfg.fps,
-                    H=video.shape[1],
-                    W=video.shape[2],
-                    video_save_quality=5,
-                    video_save_path=video_save_path,
-                )
-                video_paths.append(video_save_path)
-
-                # Save prompt to text file alongside video
-                with open(prompt_save_path, "wb") as f:
-                    f.write(prompt.encode("utf-8"))
-
-                prompt_paths.append(prompt_save_path)
-
-                log.info(f"Saved video to {video_save_path}")
-                log.info(f"Saved prompt to {prompt_save_path}")
-
-    # clean up properly
-    if cfg.num_gpus > 1:
-        parallel_state.destroy_model_parallel()
-        import torch.distributed as dist
-
-        dist.destroy_process_group()
-
-    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,
-    hdmap_video_input,
-    lidar_video_input,
-    prompt,
-    negative_prompt="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality.",  # noqa: E501
-    seed=42,
-    randomize_seed=False,
-    chunking=None,
-    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:
-        actual_seed = seed
-
-    log.info(f"actual_seed: {actual_seed}")
-
-    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
-    start_time = time.time()
-
-    # parse generation configs
-    args, control_inputs = parse_arguments(
-        controlnet_specs_in={
-            "hdmap": {"control_weight": 0.3, "input_control": hdmap_video_input},
-            "lidar": {"control_weight": 0.7, "input_control": lidar_video_input},
-        },
-        input_video_path=rgb_video_path,
-        checkpoint_dir=CHECKPOINTS_PATH,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        sigma_max=80,
-        offload_text_encoder_model=True,
-        is_av_sample=True,
-        num_gpus=1,
-        seed=seed,
-    )
-
-    # watch gpu memory
-    watcher = watch_gpu_memory(10, lambda x: log.debug(f"GPU memory usage: {x} (MiB)"))
-
-    # start inference
-    if chunking <= 0:
-        chunking = None
-    videos, prompts = inference(args, control_inputs, chunking)
-
-    # print the generation time
-    end_time = time.time()
-    log.info(f"Time taken: {end_time - start_time} s")
-
-    # stop the watcher
-    stop_watcher()
+from download_checkpoints import main as download_checkpoints
+from helper import CHECKPOINTS_PATH, generate_video
 
-    video = videos[0]
+# download checkpoints
+os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
+download_checkpoints(hf_token="", output_dir=CHECKPOINTS_PATH, model="7b_av")
 
-    log.logger.remove(log_handler)
-    return video, create_zip_for_download(filename=logfile_path, files_to_zip=[video, logfile_path]), actual_seed
 
+import gradio as gr
 
 # Define the Gradio Blocks interface
 with gr.Blocks() as demo:
@@ -412,3 +97,4 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.launch()
+    # demo.launch(server_name="0.0.0.0")

helper.py

@@ -1,10 +1,45 @@
 import argparse
+import copy
+import datetime
+import json
+import os
+import random
 import sys
-from typing import Any, Dict, Literal, Optional
+import tempfile
+import time
+import zipfile
+from io import BytesIO
+from typing import Any, Dict, List, Literal, Optional, Tuple
 
-sys.path.append("./cosmos-transfer1")
+import torch
 
-from cosmos_transfer1.diffusion.inference.inference_utils import valid_hint_keys
+from cosmos_transfer1.checkpoints import (
+    BASE_7B_CHECKPOINT_AV_SAMPLE_PATH,
+    BASE_7B_CHECKPOINT_PATH,
+    EDGE2WORLD_CONTROLNET_DISTILLED_CHECKPOINT_PATH,
+)
+from cosmos_transfer1.diffusion.inference.inference_utils import (
+    valid_hint_keys,
+    validate_controlnet_specs,
+)
+from cosmos_transfer1.diffusion.inference.preprocessors import Preprocessors
+from cosmos_transfer1.diffusion.inference.world_generation_pipeline import (
+    DiffusionControl2WorldGenerationPipeline,
+    DistilledControl2WorldGenerationPipeline,
+)
+from cosmos_transfer1.utils import log, misc
+from cosmos_transfer1.utils.io import read_prompts_from_file, save_video
+from gpu_info import stop_watcher, watch_gpu_memory
+
+PWD = os.path.dirname(__file__)
+CHECKPOINTS_PATH = "/data/checkpoints"
+LOG_DIR = os.path.join(PWD, "logs")
+os.makedirs(LOG_DIR, exist_ok=True)
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Workaround to suppress MP warning
+
+torch.enable_grad(False)
+torch.serialization.add_safe_globals([BytesIO])
 
 
 def load_controlnet_specs(controlnet_specs_in: dict) -> Dict[str, Any]:
@@ -62,7 +97,8 @@ def parse_arguments(
     :param int num_input_frames: Number of conditional frames for long video generation
     :param float sigma_max: sigma_max for partial denoising
     :param str blur_strength: blur strength
-    :param str canny_threshold: blur strength of canny threshold applied to input. Lower means less blur or more detected edges, which means higher fidelity to input
+    :param str canny_threshold: blur strength of canny threshold applied to input. Lower means less blur or more detected edges,
+                                which means higher fidelity to input
     :param bool is_av_sample: Whether the model is an driving post-training model
     :param str checkpoint_dir: Base directory containing model checkpoints
     :param str tokenizer_dir: Tokenizer weights directory relative to checkpoint_dir
@@ -121,3 +157,281 @@ def parse_arguments(
             setattr(cmd_args, key, json_args[key])
 
     return cmd_args, control_inputs
+
+
+def inference(cfg, control_inputs, chunking) -> Tuple[List[str], List[str]]:
+    video_paths = []
+    prompt_paths = []
+
+    control_inputs = validate_controlnet_specs(cfg, control_inputs)
+    misc.set_random_seed(cfg.seed)
+
+    device_rank = 0
+    process_group = None
+    if cfg.num_gpus > 1:
+        from megatron.core import (
+            parallel_state,  # pyright: ignore[reportMissingImports]
+        )
+
+        from cosmos_transfer1.utils import distributed
+
+        distributed.init()
+        parallel_state.initialize_model_parallel(context_parallel_size=cfg.num_gpus)
+        process_group = parallel_state.get_context_parallel_group()
+
+        device_rank = distributed.get_rank(process_group)
+
+    preprocessors = Preprocessors()
+
+    if cfg.use_distilled:
+        assert not cfg.is_av_sample
+        checkpoint = EDGE2WORLD_CONTROLNET_DISTILLED_CHECKPOINT_PATH
+        pipeline = DistilledControl2WorldGenerationPipeline(
+            checkpoint_dir=cfg.checkpoint_dir,
+            checkpoint_name=checkpoint,
+            offload_network=cfg.offload_diffusion_transformer,
+            offload_text_encoder_model=cfg.offload_text_encoder_model,
+            offload_guardrail_models=cfg.offload_guardrail_models,
+            guidance=cfg.guidance,
+            num_steps=cfg.num_steps,
+            fps=cfg.fps,
+            seed=cfg.seed,
+            num_input_frames=cfg.num_input_frames,
+            control_inputs=control_inputs,
+            sigma_max=cfg.sigma_max,
+            blur_strength=cfg.blur_strength,
+            canny_threshold=cfg.canny_threshold,
+            upsample_prompt=cfg.upsample_prompt,
+            offload_prompt_upsampler=cfg.offload_prompt_upsampler,
+            process_group=process_group,
+        )
+    else:
+        checkpoint = BASE_7B_CHECKPOINT_AV_SAMPLE_PATH if cfg.is_av_sample else BASE_7B_CHECKPOINT_PATH
+
+        # Initialize transfer generation model pipeline
+        pipeline = DiffusionControl2WorldGenerationPipeline(
+            checkpoint_dir=cfg.checkpoint_dir,
+            checkpoint_name=checkpoint,
+            offload_network=cfg.offload_diffusion_transformer,
+            offload_text_encoder_model=cfg.offload_text_encoder_model,
+            offload_guardrail_models=cfg.offload_guardrail_models,
+            guidance=cfg.guidance,
+            num_steps=cfg.num_steps,
+            fps=cfg.fps,
+            seed=cfg.seed,
+            num_input_frames=cfg.num_input_frames,
+            control_inputs=control_inputs,
+            sigma_max=cfg.sigma_max,
+            blur_strength=cfg.blur_strength,
+            canny_threshold=cfg.canny_threshold,
+            upsample_prompt=cfg.upsample_prompt,
+            offload_prompt_upsampler=cfg.offload_prompt_upsampler,
+            process_group=process_group,
+            chunking=chunking,
+        )
+
+    if cfg.batch_input_path:
+        log.info(f"Reading batch inputs from path: {cfg.batch_input_path}")
+        prompts = read_prompts_from_file(cfg.batch_input_path)
+    else:
+        # Single prompt case
+        prompts = [{"prompt": cfg.prompt, "visual_input": cfg.input_video_path}]
+
+    batch_size = cfg.batch_size if hasattr(cfg, "batch_size") else 1
+    if any("upscale" in control_input for control_input in control_inputs) and batch_size > 1:
+        batch_size = 1
+        log.info("Setting batch_size=1 as upscale does not support batch generation")
+    os.makedirs(cfg.video_save_folder, exist_ok=True)
+    for batch_start in range(0, len(prompts), batch_size):
+        # Get current batch
+        batch_prompts = prompts[batch_start : batch_start + batch_size]
+        actual_batch_size = len(batch_prompts)
+        # Extract batch data
+        batch_prompt_texts = [p.get("prompt", None) for p in batch_prompts]
+        batch_video_paths = [p.get("visual_input", None) for p in batch_prompts]
+
+        batch_control_inputs = []
+        for i, input_dict in enumerate(batch_prompts):
+            current_prompt = input_dict.get("prompt", None)
+            current_video_path = input_dict.get("visual_input", None)
+
+            if cfg.batch_input_path:
+                video_save_subfolder = os.path.join(cfg.video_save_folder, f"video_{batch_start+i}")
+                os.makedirs(video_save_subfolder, exist_ok=True)
+            else:
+                video_save_subfolder = cfg.video_save_folder
+
+            current_control_inputs = copy.deepcopy(control_inputs)
+            if "control_overrides" in input_dict:
+                for hint_key, override in input_dict["control_overrides"].items():
+                    if hint_key in current_control_inputs:
+                        current_control_inputs[hint_key].update(override)
+                    else:
+                        log.warning(f"Ignoring unknown control key in override: {hint_key}")
+
+            # if control inputs are not provided, run respective preprocessor (for seg and depth)
+            log.info("running preprocessor")
+            preprocessors(
+                current_video_path,
+                current_prompt,
+                current_control_inputs,
+                video_save_subfolder,
+                cfg.regional_prompts if hasattr(cfg, "regional_prompts") else None,
+            )
+            batch_control_inputs.append(current_control_inputs)
+
+        regional_prompts = []
+        region_definitions = []
+        if hasattr(cfg, "regional_prompts") and cfg.regional_prompts:
+            log.info(f"regional_prompts: {cfg.regional_prompts}")
+            for regional_prompt in cfg.regional_prompts:
+                regional_prompts.append(regional_prompt["prompt"])
+                if "region_definitions_path" in regional_prompt:
+                    log.info(f"region_definitions_path: {regional_prompt['region_definitions_path']}")
+                    region_definition_path = regional_prompt["region_definitions_path"]
+                    if isinstance(region_definition_path, str) and region_definition_path.endswith(".json"):
+                        with open(region_definition_path, "r") as f:
+                            region_definitions_json = json.load(f)
+                        region_definitions.extend(region_definitions_json)
+                    else:
+                        region_definitions.append(region_definition_path)
+
+        if hasattr(pipeline, "regional_prompts"):
+            pipeline.regional_prompts = regional_prompts
+        if hasattr(pipeline, "region_definitions"):
+            pipeline.region_definitions = region_definitions
+
+        # Generate videos in batch
+        batch_outputs = pipeline.generate(
+            prompt=batch_prompt_texts,
+            video_path=batch_video_paths,
+            negative_prompt=cfg.negative_prompt,
+            control_inputs=batch_control_inputs,
+            save_folder=video_save_subfolder,
+            batch_size=actual_batch_size,
+        )
+        if batch_outputs is None:
+            log.critical("Guardrail blocked generation for entire batch.")
+            continue
+
+        videos, final_prompts = batch_outputs
+        for i, (video, prompt) in enumerate(zip(videos, final_prompts)):
+            if cfg.batch_input_path:
+                video_save_subfolder = os.path.join(cfg.video_save_folder, f"video_{batch_start+i}")
+                video_save_path = os.path.join(video_save_subfolder, "output.mp4")
+                prompt_save_path = os.path.join(video_save_subfolder, "prompt.txt")
+            else:
+                video_save_path = os.path.join(cfg.video_save_folder, f"{cfg.video_save_name}.mp4")
+                prompt_save_path = os.path.join(cfg.video_save_folder, f"{cfg.video_save_name}.txt")
+            # Save video and prompt
+            if device_rank == 0:
+                os.makedirs(os.path.dirname(video_save_path), exist_ok=True)
+                save_video(
+                    video=video,
+                    fps=cfg.fps,
+                    H=video.shape[1],
+                    W=video.shape[2],
+                    video_save_quality=5,
+                    video_save_path=video_save_path,
+                )
+                video_paths.append(video_save_path)
+
+                # Save prompt to text file alongside video
+                with open(prompt_save_path, "wb") as f:
+                    f.write(prompt.encode("utf-8"))
+
+                prompt_paths.append(prompt_save_path)
+
+                log.info(f"Saved video to {video_save_path}")
+                log.info(f"Saved prompt to {prompt_save_path}")
+
+    # clean up properly
+    if cfg.num_gpus > 1:
+        parallel_state.destroy_model_parallel()
+        import torch.distributed as dist
+
+        dist.destroy_process_group()
+
+    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
+
+
+import gradio as gr
+
+
+def generate_video(
+    rgb_video_path,
+    hdmap_video_input,
+    lidar_video_input,
+    prompt,
+    negative_prompt="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality.",  # noqa: E501
+    seed=42,
+    randomize_seed=False,
+    chunking=None,
+    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:
+        actual_seed = seed
+
+    log.info(f"actual_seed: {actual_seed}")
+
+    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
+    start_time = time.time()
+
+    # parse generation configs
+    args, control_inputs = parse_arguments(
+        controlnet_specs_in={
+            "hdmap": {"control_weight": 0.3, "input_control": hdmap_video_input},
+            "lidar": {"control_weight": 0.7, "input_control": lidar_video_input},
+        },
+        input_video_path=rgb_video_path,
+        checkpoint_dir=CHECKPOINTS_PATH,
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        sigma_max=80,
+        offload_text_encoder_model=True,
+        is_av_sample=True,
+        num_gpus=1,
+        seed=seed,
+    )
+
+    # watch gpu memory
+    watcher = watch_gpu_memory(10, lambda x: log.debug(f"GPU memory usage: {x} (MiB)"))
+
+    # start inference
+    if chunking <= 0:
+        chunking = None
+    videos, prompts = inference(args, control_inputs, chunking)
+
+    # print the generation time
+    end_time = time.time()
+    log.info(f"Time taken: {end_time - start_time} s")
+
+    # stop the watcher
+    stop_watcher()
+
+    video = videos[0]
+
+    log.logger.remove(log_handler)
+    return video, create_zip_for_download(filename=logfile_path, files_to_zip=[video, logfile_path]), actual_seed