Update downloader.py

downloader.py

@@ -1,494 +1,132 @@
-# ==============================================================================
-# 统一入口和依赖项
-# ==============================================================================
-import torch
-import numpy as np
-import random
-import os
 import yaml
-import argparse
+import os
 from pathlib import Path
-import imageio
-import tempfile
-from PIL import Image
 from huggingface_hub import hf_hub_download
-import shutil
-
-# 监听模式所需的依赖项
-import asyncio
-import websockets
-import subprocess
-import json
-import logging
-import sys
-import urllib.parse
-import requests
-
-from inference import (
-    create_ltx_video_pipeline,
-    create_latent_upsampler,
-    load_image_to_tensor_with_resize_and_crop,
-    seed_everething,
-    get_device,
-    calculate_padding,
-    load_media_file
-)
-from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline, LTXVideoPipeline
-from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
-
-# ==============================================================================
-# 日志配置
-# ==============================================================================
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# ==============================================================================
-# 监听模式的函数 (原 remote_client.py)
-# ==============================================================================
-
-# 全局变量，用于在监听模式下共享状态
-global_websocket = None
-global_machine_id = None
-global_card_id = None
-global_machine_secret = None
-global_server_url = None
-
-async def upload_file_to_server(file_path, card_id, machine_secret, machine_id):
-    """将文件上传到服务器的指定端点"""
-    try:
-        if not os.path.exists(file_path):
-            logger.error(f"[Uploader] File not found: {file_path}")
-            return False
-
-        upload_url = f"{global_server_url}/terminal/{card_id}/machine-upload?secret={urllib.parse.quote(machine_secret)}"
-        files = {'file': (os.path.basename(file_path), open(file_path, 'rb'), 'application/octet-stream')}
-        data = {'machine_id': machine_id}
-
-        logger.info(f"[Uploader] Uploading {os.path.basename(file_path)} to {upload_url}...")
-        response = requests.post(upload_url, files=files, data=data, timeout=120)
-        
-        if response.status_code == 200:
-            result = response.json()
-            if result and result.get("success"):
-                logger.info(f"[Uploader] Upload successful: {file_path}")
-                return True
-            else:
-                logger.error(f"[Uploader] Upload failed: {result.get('error', 'Unknown error')}")
-                return False
-        else:
-            logger.error(f"[Uploader] Upload failed with status code {response.status_code}: {response.text}")
-            return False
-            
-    except Exception as e:
-        logger.error(f"[Uploader] An exception occurred during upload: {e}")
-        return False
-
-async def watch_directory_for_uploads(dir_to_watch, card_id, secret, get_machine_id_func):
-    """
-    监视指定目录中的新文件，并自动上传。
-    """
-    processed_files = set()
-    logger.info(f"[Watcher] Starting to watch directory: {dir_to_watch}")
-    
-    # 初始扫描，将已存在的文件视为已处理
-    if os.path.isdir(dir_to_watch):
-        processed_files.update(os.listdir(dir_to_watch))
-        logger.info(f"[Watcher] Initial scan: {len(processed_files)} existing files ignored.")
-
-    while True:
-        await asyncio.sleep(5) # 每5秒检查一次
-        try:
-            if not os.path.isdir(dir_to_watch):
-                continue
-
-            current_files = set(os.listdir(dir_to_watch))
-            new_files = current_files - processed_files
-            
-            if new_files:
-                machine_id = get_machine_id_func()
-                if not machine_id:
-                    logger.warning("[Watcher] Machine ID not available, skipping upload cycle.")
-                    continue
-
-                logger.info(f"[Watcher] Detected {len(new_files)} new file(s): {', '.join(new_files)}")
-                for filename in new_files:
-                    file_path = os.path.join(dir_to_watch, filename)
-                    # 等待文件写入完成 (简单检查)
-                    await asyncio.sleep(2) 
-                    
-                    success = await upload_file_to_server(file_path, card_id, secret, machine_id)
-                    if success:
-                        logger.info(f"[Watcher] Successfully uploaded {filename}. Marking as processed.")
-                        processed_files.add(filename)
-                    else:
-                        logger.warning(f"[Watcher] Failed to upload {filename}. Will retry on next cycle.")
-
-            # 同步已处理列表，移除已删除的文件
-            processed_files.intersection_update(current_files)
-
-        except Exception as e:
-            logger.error(f"[Watcher] Error in file watching loop: {e}")
 
-
-async def start_listener_mode(card_id, machine_secret, watch_dir):
+def download_ltx_models():
     """
-    启动监听模式的主函数。
+    独立下载LTX-Video模型的脚本
+    保持与主程序相同的路径和配置
     """
-    global global_websocket, global_machine_id, global_card_id, global_machine_secret, global_server_url
-
-    global_card_id = card_id
-    global_machine_secret = machine_secret
-
-    server_hostname = "remote-terminal-worker.nianxi4563.workers.dev" # 或者您的服务器域名
-    global_server_url = f"https://{server_hostname}"
-    encoded_secret = urllib.parse.quote(machine_secret)
-    uri = f"wss://{server_hostname}/terminal/{card_id}?secret={encoded_secret}"
-
-    # 启动文件监视器
-    def get_machine_id(): return global_machine_id
-    watcher_task = asyncio.create_task(watch_directory_for_uploads(watch_dir, card_id, machine_secret, get_machine_id))
-
-    while True: # 自动重连循环
-        try:
-            logger.info(f"[Listener] Attempting to connect to {uri}")
-            async with websockets.connect(uri, ping_interval=20, ping_timeout=60) as websocket:
-                global_websocket = websocket
-                logger.info("[Listener] Connected to WebSocket server.")
-
-                # 循环以获取 machine_id
-                while global_machine_id is None:
-                    try:
-                        response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
-                        data = json.loads(response)
-                        if data.get("type") == "connected" and "machine_id" in data:
-                            global_machine_id = data["machine_id"]
-                            logger.info(f"[Listener] Assigned machine ID: {global_machine_id}")
-                            break
-                    except asyncio.TimeoutError:
-                        logger.debug("[Listener] Waiting for machine ID...")
-                    except Exception as e:
-                        logger.error(f"[Listener] Error receiving machine ID: {e}")
-                        await asyncio.sleep(5) # 等待后重试
-                        break # break inner loop to reconnect
-
-                if not global_machine_id:
-                    continue # continue outer loop to reconnect
-
-                # 主消息处理循环
-                while True:
-                    message = await websocket.recv()
-                    data = json.loads(message)
-                    logger.debug(f"[Listener] Received message: {data}")
-
-                    if data.get("type") == "command":
-                        command = data["command"]
-                        logger.info(f"[Listener] Received command: {command}")
-                        
-                        # 使用 subprocess 在新进程中执行命令
-                        # 这使得监听器可以继续工作，而推理在后台运行
-                        try:
-                            # 将命令包装在 `python app.py ...` 中
-                            full_command = f"python app.py {command}"
-                            logger.info(f"Executing subprocess: {full_command}")
-                            subprocess.run(full_command, shell=True, check=True)
-                            logger.info("Subprocess finished successfully.")
-                            # 结果文件将由 watcher 自动上传
-                        except subprocess.CalledProcessError as e:
-                            logger.error(f"Command execution failed with return code {e.returncode}")
-                            error_output = e.stderr if e.stderr else e.stdout
-                            if global_websocket:
-                                await global_websocket.send(json.dumps({
-                                    "type": "error", "data": f"Command failed: {error_output}", "machine_id": global_machine_id
-                                }))
-                        except Exception as e:
-                            logger.error(f"Failed to run command: {e}")
-
-        except websockets.exceptions.ConnectionClosed as e:
-            logger.warning(f"[Listener] WebSocket closed: code={e.code}, reason={e.reason}. Reconnecting in 10 seconds...")
-        except Exception as e:
-            logger.error(f"[Listener] Connection failed: {e}. Reconnecting in 10 seconds...")
-        
-        global_websocket = None
-        global_machine_id = None
-        await asyncio.sleep(10)
-
-
-# ==============================================================================
-# 推理模式的函数 (原 app.py)
-# ==============================================================================
-config_file_path = "configs/ltxv-13b-0.9.7-distilled.yaml"
-with open(config_file_path, "r") as file:
-    PIPELINE_CONFIG_YAML = yaml.safe_load(file)
-
-LTX_REPO = "Lightricks/LTX-Video"
-MAX_IMAGE_SIZE = PIPELINE_CONFIG_YAML.get("max_resolution", 1280)
-MAX_NUM_FRAMES = 257
-FPS = 30.0 
-
-# 全局变量以缓存加载的模型
-pipeline_instance = None
-latent_upsampler_instance = None
-models_dir = "downloaded_models_gradio_cpu_init"
-Path(models_dir).mkdir(parents=True, exist_ok=True)
-output_dir = "output" # 所有模式共用的输出目录
-Path(output_dir).mkdir(parents=True, exist_ok=True)
-
-def initialize_models():
-    """加载并初始化所有AI模型（如果尚未加载）。"""
-    global pipeline_instance, latent_upsampler_instance
-
-    if pipeline_instance is not None:
-        logger.info("Models already initialized.")
-        return
-
-    logger.info("Initializing models for the first time...")
-    logger.info("Downloading models (if not present)...")
-    distilled_model_actual_path = hf_hub_download(
-        repo_id=LTX_REPO, filename=PIPELINE_CONFIG_YAML["checkpoint_path"], local_dir=models_dir, local_dir_use_symlinks=False
-    )
-    PIPELINE_CONFIG_YAML["checkpoint_path"] = distilled_model_actual_path
-    logger.info(f"Distilled model path: {distilled_model_actual_path}")
-
-    SPATIAL_UPSCALER_FILENAME = PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"]
-    spatial_upscaler_actual_path = hf_hub_download(
-        repo_id=LTX_REPO, filename=SPATIAL_UPSCALER_FILENAME, local_dir=models_dir, local_dir_use_symlinks=False
-    )
-    PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"] = spatial_upscaler_actual_path
-    logger.info(f"Spatial upscaler model path: {spatial_upscaler_actual_path}")
-
-    logger.info("Creating LTX Video pipeline on CPU...")
-    pipeline_instance = create_ltx_video_pipeline(
-        ckpt_path=PIPELINE_CONFIG_YAML["checkpoint_path"],
-        precision=PIPELINE_CONFIG_YAML["precision"],
-        text_encoder_model_name_or_path=PIPELINE_CONFIG_YAML["text_encoder_model_name_or_path"],
-        sampler=PIPELINE_CONFIG_YAML["sampler"],
-        device="cpu",
-        enhance_prompt=False,
-        prompt_enhancer_image_caption_model_name_or_path=PIPELINE_CONFIG_YAML["prompt_enhancer_image_caption_model_name_or_path"],
-        prompt_enhancer_llm_model_name_or_path=PIPELINE_CONFIG_YAML["prompt_enhancer_llm_model_name_or_path"],
-    )
-    logger.info("LTX Video pipeline created on CPU.")
-
-    if PIPELINE_CONFIG_YAML.get("spatial_upscaler_model_path"):
-        logger.info("Creating latent upsampler on CPU...")
-        latent_upsampler_instance = create_latent_upsampler(
-            PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"], device="cpu"
-        )
-        logger.info("Latent upsampler created on CPU.")
-
-    target_inference_device = "cuda" if torch.cuda.is_available() else "cpu"
-    logger.info(f"Moving models to target inference device: {target_inference_device}")
-    pipeline_instance.to(target_inference_device)
-    if latent_upsampler_instance: 
-        latent_upsampler_instance.to(target_inference_device)
-    logger.info("Model initialization complete.")
-
-
-def generate(prompt, negative_prompt="worst quality, inconsistent motion, blurry, jittery, distorted", 
-             input_image_filepath=None, input_video_filepath=None,
-             height_ui=512, width_ui=704, mode="text-to-video",
-             duration_ui=2.0, ui_frames_to_use=9,
-             seed_ui=42, randomize_seed=True, ui_guidance_scale=None, improve_texture_flag=True):
-
-    # 确保模型已加载
-    initialize_models()
     
-    target_inference_device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    if randomize_seed:
-        seed_ui = random.randint(0, 2**32 - 1)
-    seed_everething(int(seed_ui))
+    # 读取配置文件
+    config_file_path = "configs/ltxv-13b-0.9.7-distilled.yaml"
     
-    if ui_guidance_scale is None:
-        ui_guidance_scale = PIPELINE_CONFIG_YAML.get("first_pass", {}).get("guidance_scale", 1.0)
+    if not os.path.exists(config_file_path):
+        print(f"错误: 配置文件 {config_file_path} 不存在")
+        print("请确保配置文件在正确的位置")
+        return False
     
-    target_frames_ideal = duration_ui * FPS
-    target_frames_rounded = round(target_frames_ideal)
-    if target_frames_rounded < 1: 
-        target_frames_rounded = 1
+    with open(config_file_path, "r") as file:
+        PIPELINE_CONFIG_YAML = yaml.safe_load(file)
     
-    n_val = round((float(target_frames_rounded) - 1.0) / 8.0)
-    actual_num_frames = int(n_val * 8 + 1)
-
-    actual_num_frames = max(9, actual_num_frames)
-    actual_num_frames = min(MAX_NUM_FRAMES, actual_num_frames)
+    # 设置常量
+    LTX_REPO = "Lightricks/LTX-Video"
+    models_dir = "downloaded_models_gradio_cpu_init"
     
-    actual_height = int(height_ui)
-    actual_width = int(width_ui)
-
-    height_padded = ((actual_height - 1) // 32 + 1) * 32
-    width_padded = ((actual_width - 1) // 32 + 1) * 32
-    num_frames_padded = ((actual_num_frames - 2) // 8 + 1) * 8 + 1
+    # 创建模型目录
+    Path(models_dir).mkdir(parents=True, exist_ok=True)
+    print(f"模型下载目录: {Path(models_dir).resolve()}")
     
-    padding_values = calculate_padding(actual_height, actual_width, height_padded, width_padded)
-
-    call_kwargs = {
-        "prompt": prompt, "negative_prompt": negative_prompt, "height": height_padded, "width": width_padded,
-        "num_frames": num_frames_padded, "frame_rate": int(FPS), 
-        "generator": torch.Generator(device=target_inference_device).manual_seed(int(seed_ui)),
-        "output_type": "pt", "conditioning_items": None, "media_items": None,
-        "decode_timestep": PIPELINE_CONFIG_YAML["decode_timestep"], "decode_noise_scale": PIPELINE_CONFIG_YAML["decode_noise_scale"],
-        "stochastic_sampling": PIPELINE_CONFIG_YAML["stochastic_sampling"], "image_cond_noise_scale": 0.15,
-        "is_video": True, "vae_per_channel_normalize": True, "mixed_precision": (PIPELINE_CONFIG_YAML["precision"] == "mixed_precision"),
-        "offload_to_cpu": False, "enhance_prompt": False,
-    }
-
-    stg_mode_str = PIPELINE_CONFIG_YAML.get("stg_mode", "attention_values")
-    if stg_mode_str.lower() in ["stg_av", "attention_values"]:
-        call_kwargs["skip_layer_strategy"] = SkipLayerStrategy.AttentionValues
-    elif stg_mode_str.lower() in ["stg_as", "attention_skip"]:
-        call_kwargs["skip_layer_strategy"] = SkipLayerStrategy.AttentionSkip
-    elif stg_mode_str.lower() in ["stg_r", "residual"]:
-        call_kwargs["skip_layer_strategy"] = SkipLayerStrategy.Residual
-    elif stg_mode_str.lower() in ["stg_t", "transformer_block"]:
-        call_kwargs["skip_layer_strategy"] = SkipLayerStrategy.TransformerBlock
-    else:
-        raise ValueError(f"Invalid stg_mode: {stg_mode_str}")
-
-    if mode == "image-to-video" and input_image_filepath:
-        try:
-            media_tensor = load_image_to_tensor_with_resize_and_crop(input_image_filepath, actual_height, actual_width)
-            media_tensor = torch.nn.functional.pad(media_tensor, padding_values)
-            call_kwargs["conditioning_items"] = [ConditioningItem(media_tensor.to(target_inference_device), 0, 1.0)]
-        except Exception as e:
-            logger.error(f"Error loading image {input_image_filepath}: {e}")
-            raise RuntimeError(f"Could not load image: {e}")
-    elif mode == "video-to-video" and input_video_filepath:
-        try:
-            call_kwargs["media_items"] = load_media_file(
-                media_path=input_video_filepath, height=actual_height, width=actual_width,
-                max_frames=int(ui_frames_to_use), padding=padding_values
-            ).to(target_inference_device)
-        except Exception as e:
-            logger.error(f"Error loading video {input_video_filepath}: {e}")
-            raise RuntimeError(f"Could not load video: {e}")
-
-    active_latent_upsampler = latent_upsampler_instance if improve_texture_flag and latent_upsampler_instance else None
-    result_images_tensor = None
-
-    if improve_texture_flag:
-        if not active_latent_upsampler:
-            raise RuntimeError("Spatial upscaler model not loaded or improve_texture not selected.")
+    try:
+        # 下载主模型
+        print("\n开始下载主模型...")
+        print(f"模型文件: {PIPELINE_CONFIG_YAML['checkpoint_path']}")
         
-        multi_scale_pipeline_obj = LTXMultiScalePipeline(pipeline_instance, active_latent_upsampler)
-        first_pass_args = {**PIPELINE_CONFIG_YAML.get("first_pass", {}), "guidance_scale": float(ui_guidance_scale)}
-        second_pass_args = {**PIPELINE_CONFIG_YAML.get("second_pass", {}), "guidance_scale": float(ui_guidance_scale)}
+        distilled_model_actual_path = hf_hub_download(
+            repo_id=LTX_REPO,
+            filename=PIPELINE_CONFIG_YAML["checkpoint_path"],
+            local_dir=models_dir,
+            local_dir_use_symlinks=False
+        )
+        print(f"✅ 主模型下载完成: {distilled_model_actual_path}")
         
-        multi_scale_call_kwargs = {
-            **call_kwargs, "downscale_factor": PIPELINE_CONFIG_YAML["downscale_factor"],
-            "first_pass": first_pass_args, "second_pass": second_pass_args
-        }
+        # 下载空间上采样器模型
+        print("\n开始下载空间上采样器模型...")
+        SPATIAL_UPSCALER_FILENAME = PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"]
+        print(f"模型文件: {SPATIAL_UPSCALER_FILENAME}")
         
-        logger.info(f"Calling multi-scale pipeline on {target_inference_device}")
-        result_images_tensor = multi_scale_pipeline_obj(**multi_scale_call_kwargs).images
-    else:
-        single_pass_call_kwargs = {**call_kwargs, **PIPELINE_CONFIG_YAML.get("first_pass", {}), "guidance_scale": float(ui_guidance_scale)}
-        logger.info(f"Calling base pipeline on {target_inference_device}")
-        result_images_tensor = pipeline_instance(**single_pass_call_kwargs).images
-
-    if result_images_tensor is None:
-        raise RuntimeError("Generation failed, result tensor is None.")
+        spatial_upscaler_actual_path = hf_hub_download(
+            repo_id=LTX_REPO,
+            filename=SPATIAL_UPSCALER_FILENAME,
+            local_dir=models_dir,
+            local_dir_use_symlinks=False
+        )
+        print(f"✅ 空间上采样器模型下载完成: {spatial_upscaler_actual_path}")
+        
+        # 显示下载摘要
+        print("\n" + "="*60)
+        print("模型下载完成摘要:")
+        print("="*60)
+        print(f"下载目录: {models_dir}")
+        print(f"主模型: {os.path.basename(distilled_model_actual_path)}")
+        print(f"上采样器: {os.path.basename(spatial_upscaler_actual_path)}")
+        
+        # 检查文件大小
+        main_size = os.path.getsize(distilled_model_actual_path) / (1024**3)  # GB
+        upscaler_size = os.path.getsize(spatial_upscaler_actual_path) / (1024**3)  # GB
+        total_size = main_size + upscaler_size
+        
+        print(f"\n文件大小:")
+        print(f"主模型: {main_size:.2f} GB")
+        print(f"上采样器: {upscaler_size:.2f} GB")
+        print(f"总计: {total_size:.2f} GB")
+        
+        return True
+        
+    except Exception as e:
+        print(f"\n❌ 下载过程中出现错误: {e}")
+        print("可能的解决方案:")
+        print("1. 检查网络连接")
+        print("2. 确认Hugging Face访问权限")
+        print("3. 检查磁盘空间是否足够")
+        return False
 
-    pad_left, pad_right, pad_top, pad_bottom = padding_values
-    slice_h_end = -pad_bottom if pad_bottom > 0 else None
-    slice_w_end = -pad_right if pad_right > 0 else None
+def check_models_exist():
+    """
+    检查模型是否已经存在
+    """
+    config_file_path = "configs/ltxv-13b-0.9.7-distilled.yaml"
     
-    result_images_tensor = result_images_tensor[:, :, :actual_num_frames, pad_top:slice_h_end, pad_left:slice_w_end]
-
-    video_np = (result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).clip(0, 255).astype(np.uint8)
+    if not os.path.exists(config_file_path):
+        return False
     
-    # 使用随机数确保文件名几乎不重复
-    timestamp = random.randint(10000, 99999)
-    output_video_path = os.path.join(output_dir, f"output_{timestamp}_{seed_ui}.mp4")
+    with open(config_file_path, "r") as file:
+        config = yaml.safe_load(file)
     
-    try:
-        with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], macro_block_size=1) as writer:
-            for frame in video_np:
-                writer.append_data(frame)
-    except Exception:
-        with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], format='FFMPEG', codec='libx264') as writer:
-            for frame in video_np:
-                writer.append_data(frame)
+    models_dir = "downloaded_models_gradio_cpu_init"
+    main_model = os.path.join(models_dir, config["checkpoint_path"])
+    upscaler_model = os.path.join(models_dir, config["spatial_upscaler_model_path"])
     
-    logger.info(f"Video saved successfully to: {output_video_path}")
-    return output_video_path, seed_ui
-
-def run_inference(args):
-    """处理命令行参数并运行AI推理。"""
-    logger.info(f"Starting single-run inference...")
-    logger.info(f"Prompt: {args.prompt}")
-    logger.info(f"Mode: {args.mode}")
-    logger.info(f"Duration: {args.duration}s")
-    logger.info(f"Resolution: {args.width}x{args.height}")
-    logger.info(f"Output directory: {os.path.abspath(output_dir)}")
+    main_exists = os.path.exists(main_model)
+    upscaler_exists = os.path.exists(upscaler_model)
     
-    try:
-        output_path, used_seed = generate(
-            prompt=args.prompt, negative_prompt=args.negative_prompt,
-            input_image_filepath=args.input_image, input_video_filepath=args.input_video,
-            height_ui=args.height, width_ui=args.width, mode=args.mode,
-            duration_ui=args.duration, ui_frames_to_use=args.frames_to_use,
-            seed_ui=args.seed, randomize_seed=args.randomize_seed,
-            ui_guidance_scale=args.guidance_scale, improve_texture_flag=not args.no_improve_texture
-        )
-        logger.info(f"\n✅ Video generation completed!")
-        logger.info(f"📁 Output saved to: {output_path}")
-        logger.info(f"🎲 Used seed: {used_seed}")
-        
-    except Exception as e:
-        logger.error(f"❌ Error during generation: {e}", exc_info=True)
-        sys.exit(1)
-
+    print("模型存在性检查:")
+    print(f"主模型: {'✅ 存在' if main_exists else '❌ 不存在'}")
+    print(f"上采样器: {'✅ 存在' if upscaler_exists else '❌ 不存在'}")
+    
+    return main_exists and upscaler_exists
 
-# ==============================================================================
-# 主入口和参数解析
-# ==============================================================================
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="LTX Video Generation and Server Client")
+def main():
+    print("LTX-Video 模型下载器")
+    print("="*40)
     
-    # --- 模式选择 ---
-    group = parser.add_argument_group('运行模式')
-    group.add_argument("--listen", action="store_true", help="以监听模式运行，连接到服务器等待指令。")
+    # 检查模型是否已存在
+    if check_models_exist():
+        print("\n所有模型已存在，无需重新下载。")
+        choice = input("是否要重新下载？(y/N): ").lower().strip()
+        if choice != 'y':
+            print("取消下载。")
+            return
     
-    # --- 监听模式参数 ---
-    listener_group = parser.add_argument_group('监听模式参数 (需配合 --listen)')
-    listener_group.add_argument("--card-id", help="用于向服务器认证的Card ID。")
-    listener_group.add_argument("--secret", help="用于向服务器认证的Machine Secret。")
-    listener_group.add_argument("--watch-dir", default=output_dir, help=f"监听新文件并自动上传的目录 (默认: {output_dir})")
-
-    # --- 推理模式参数 ---
-    inference_group = parser.add_argument_group('推理模式参数 (默认模式)')
-    inference_group.add_argument("--prompt", help="用于视频生成的文本提示。")
-    inference_group.add_argument("--negative-prompt", default="worst quality, inconsistent motion, blurry, jittery, distorted", help="负面提示。")
-    inference_group.add_argument("--mode", choices=["text-to-video", "image-to-video", "video-to-video"], default="text-to-video", help="生成模式。")
-    inference_group.add_argument("--input-image", help="输入图像路径 (用于 image-to-video 模式)。")
-    inference_group.add_argument("--input-video", help="输入视频路径 (用于 video-to-video 模式)。")
-    inference_group.add_argument("--duration", type=float, default=2.0, help="视频时长 (秒, 0.3-8.5)。")
-    inference_group.add_argument("--height", type=int, default=512, help="视频高度 (将被调整为32的倍数)。")
-    inference_group.add_argument("--width", type=int, default=704, help="视频宽度 (将被调整为32的倍数)。")
-    inference_group.add_argument("--seed", type=int, default=42, help="随机种子。")
-    inference_group.add_argument("--randomize-seed", action="store_true", help="使用一个随机的种子。")
-    inference_group.add_argument("--guidance-scale", type=float, help="引导比例。")
-    inference_group.add_argument("--no-improve-texture", action="store_true", help="禁用纹理增强 (更快，但质量可能较低)。")
-    inference_group.add_argument("--frames-to-use", type=int, default=9, help="从输入视频中使用多少帧 (用于 video-to-video)。")
-
-    args = parser.parse_args()
-
-    # 根据模式分发任务
-    if args.listen:
-        if not args.card_id or not args.secret:
-            parser.error("--card-id 和 --secret 是 --listen 模式的必需参数。")
-        logger.info(f"启动监听模式... Card ID: {args.card_id}, Watch Dir: {args.watch_dir}")
-        try:
-            asyncio.run(start_listener_mode(args.card_id, args.secret, args.watch_dir))
-        except KeyboardInterrupt:
-            logger.info("监听模式已停止。")
+    print("\n开始下载模型...")
+    success = download_ltx_models()
+    
+    if success:
+        print("\n🎉 所有模型下载成功！")
+        print("现在可以运行主程序了。")
     else:
-        if not args.prompt:
-            parser.error("--prompt 是推理模式的必需参数。")
-        
-        # 确保尺寸是32的倍数
-        args.height = ((args.height - 1) // 32 + 1) * 32
-        args.width = ((args.width - 1) // 32 + 1) * 32
-        
-        run_inference(args)
+        print("\n💥 模型下载失败，请检查错误信息并重试。")
+
+if __name__ == "__main__":
+    main()