Update downloader.py

downloader.py

@@ -1,132 +1,494 @@
-import yaml
+# ==============================================================================
+# 统一入口和依赖项
+# ==============================================================================
+import torch
+import numpy as np
+import random
 import os
+import yaml
+import argparse
 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}")
 
-def download_ltx_models():
+
+async def start_listener_mode(card_id, machine_secret, watch_dir):
     """
-    独立下载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()
     
-    # 读取配置文件
-    config_file_path = "configs/ltxv-13b-0.9.7-distilled.yaml"
+    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))
     
-    if not os.path.exists(config_file_path):
-        print(f"错误: 配置文件 {config_file_path} 不存在")
-        print("请确保配置文件在正确的位置")
-        return False
+    if ui_guidance_scale is None:
+        ui_guidance_scale = PIPELINE_CONFIG_YAML.get("first_pass", {}).get("guidance_scale", 1.0)
     
-    with open(config_file_path, "r") as file:
-        PIPELINE_CONFIG_YAML = yaml.safe_load(file)
+    target_frames_ideal = duration_ui * FPS
+    target_frames_rounded = round(target_frames_ideal)
+    if target_frames_rounded < 1: 
+        target_frames_rounded = 1
     
-    # 设置常量
-    LTX_REPO = "Lightricks/LTX-Video"
-    models_dir = "downloaded_models_gradio_cpu_init"
+    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)
     
-    # 创建模型目录
-    Path(models_dir).mkdir(parents=True, exist_ok=True)
-    print(f"模型下载目录: {Path(models_dir).resolve()}")
+    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
     
-    try:
-        # 下载主模型
-        print("\n开始下载主模型...")
-        print(f"模型文件: {PIPELINE_CONFIG_YAML['checkpoint_path']}")
-        
-        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}")
-        
-        # 下载空间上采样器模型
-        print("\n开始下载空间上采样器模型...")
-        SPATIAL_UPSCALER_FILENAME = PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"]
-        print(f"模型文件: {SPATIAL_UPSCALER_FILENAME}")
-        
-        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
+    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.")
         
-        print(f"\n文件大小:")
-        print(f"主模型: {main_size:.2f} GB")
-        print(f"上采样器: {upscaler_size:.2f} GB")
-        print(f"总计: {total_size:.2f} GB")
+        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)}
         
-        return True
+        multi_scale_call_kwargs = {
+            **call_kwargs, "downscale_factor": PIPELINE_CONFIG_YAML["downscale_factor"],
+            "first_pass": first_pass_args, "second_pass": second_pass_args
+        }
         
-    except Exception as e:
-        print(f"\n❌ 下载过程中出现错误: {e}")
-        print("可能的解决方案:")
-        print("1. 检查网络连接")
-        print("2. 确认Hugging Face访问权限")
-        print("3. 检查磁盘空间是否足够")
-        return False
+        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
 
-def check_models_exist():
-    """
-    检查模型是否已经存在
-    """
-    config_file_path = "configs/ltxv-13b-0.9.7-distilled.yaml"
-    
-    if not os.path.exists(config_file_path):
-        return False
-    
-    with open(config_file_path, "r") as file:
-        config = yaml.safe_load(file)
+    if result_images_tensor is None:
+        raise RuntimeError("Generation failed, result tensor is None.")
+
+    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
     
-    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"])
+    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)
     
-    main_exists = os.path.exists(main_model)
-    upscaler_exists = os.path.exists(upscaler_model)
+    # 使用随机数确保文件名几乎不重复
+    timestamp = random.randint(10000, 99999)
+    output_video_path = os.path.join(output_dir, f"output_{timestamp}_{seed_ui}.mp4")
     
-    print("模型存在性检查:")
-    print(f"主模型: {'✅ 存在' if main_exists else '❌ 不存在'}")
-    print(f"上采样器: {'✅ 存在' if upscaler_exists else '❌ 不存在'}")
+    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)
     
-    return main_exists and upscaler_exists
+    logger.info(f"Video saved successfully to: {output_video_path}")
+    return output_video_path, seed_ui
 
-def main():
-    print("LTX-Video 模型下载器")
-    print("="*40)
+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)}")
     
-    # 检查模型是否已存在
-    if check_models_exist():
-        print("\n所有模型已存在，无需重新下载。")
-        choice = input("是否要重新下载？(y/N): ").lower().strip()
-        if choice != 'y':
-            print("取消下载。")
-            return
+    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)
+
+
+# ==============================================================================
+# 主入口和参数解析
+# ==============================================================================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="LTX Video Generation and Server Client")
     
-    print("\n开始下载模型...")
-    success = download_ltx_models()
+    # --- 模式选择 ---
+    group = parser.add_argument_group('运行模式')
+    group.add_argument("--listen", action="store_true", help="以监听模式运行，连接到服务器等待指令。")
     
-    if success:
-        print("\n🎉 所有模型下载成功！")
-        print("现在可以运行主程序了。")
-    else:
-        print("\n💥 模型下载失败，请检查错误信息并重试。")
+    # --- 监听模式参数 ---
+    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})")
 
-if __name__ == "__main__":
-    main()
+    # --- 推理模式参数 ---
+    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("监听模式已停止。")
+    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)