Update app.py

app.py

@@ -1,7 +1,5 @@
-# ==============================================================================
-# 统一入口和依赖项
-# ==============================================================================
 import torch
+
 import numpy as np
 import random
 import os
@@ -10,279 +8,48 @@ 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
+if latent_upsampler_instance: 
+    latent_upsampler_instance.to(target_inference_device)
 
-        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
+# --- Helper function for dimension calculation ---
+MIN_DIM_SLIDER = 256
+TARGET_FIXED_SIDE = 768
 
-async def watch_directory_for_uploads(dir_to_watch, card_id, secret, get_machine_id_func):
+def calculate_new_dimensions(orig_w, orig_h):
     """
-    监视指定目录中的新文件，并自动上传。
+    both are multiples of 32, and within [MIN_DIM_SLIDER, MAX_IMAGE_SIZE].
     """
-    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.")
+    if orig_w == 0 or orig_h == 0:
 
-    while True:
-        await asyncio.sleep(5) # 每5秒检查一次
-        try:
-            if not os.path.isdir(dir_to_watch):
-                continue
+        return int(TARGET_FIXED_SIDE), int(TARGET_FIXED_SIDE)
 
-            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)
+    if orig_w >= orig_h:
+        new_h = TARGET_FIXED_SIDE
+        aspect_ratio = orig_w / orig_h
+        new_w_ideal = new_h * aspect_ratio
 
-        except Exception as e:
-            logger.error(f"[Watcher] Error in file watching loop: {e}")
 
+        new_w = round(new_w_ideal / 32) * 32
 
-async def start_listener_mode(card_id, machine_secret, watch_dir):
-    """
-    启动监听模式的主函数。
-    """
-    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
+        new_w = max(MIN_DIM_SLIDER, min(new_w, MAX_IMAGE_SIZE))
+
+        new_h = max(MIN_DIM_SLIDER, min(new_h, MAX_IMAGE_SIZE)) 
+    else:
+        new_w = TARGET_FIXED_SIDE
+        aspect_ratio = orig_h / orig_w
+        new_h_ideal = new_w * aspect_ratio
 
-    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))
+        new_h = round(new_h_ideal / 32) * 32
 
-    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.")
+        new_h = max(MIN_DIM_SLIDER, min(new_h, MAX_IMAGE_SIZE))
 
-    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.")
+        new_w = max(MIN_DIM_SLIDER, min(new_w, MAX_IMAGE_SIZE))
 
+    return int(new_h), int(new_w)
 
 def generate(prompt, negative_prompt="worst quality, inconsistent motion, blurry, jittery, distorted", 
              input_image_filepath=None, input_video_filepath=None,
@@ -290,10 +57,18 @@ def generate(prompt, negative_prompt="worst quality, inconsistent motion, blurry
              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)
@@ -305,190 +80,281 @@ def generate(prompt, negative_prompt="worst quality, inconsistent motion, blurry
     target_frames_ideal = duration_ui * FPS
     target_frames_rounded = round(target_frames_ideal)
     if target_frames_rounded < 1: 
-        target_frames_rounded = 1
-    
-    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)
-    
-    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
+
+
     
     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}")
+            print(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}")
+            print(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
+    print(f"Moving models to {target_inference_device} for inference (if not already there)...")
+    
     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.")
+            raise RuntimeError("Spatial upscaler model not loaded or improve_texture not selected, cannot use multi-scale.")
         
         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)}
         
-        multi_scale_call_kwargs = {
-            **call_kwargs, "downscale_factor": PIPELINE_CONFIG_YAML["downscale_factor"],
-            "first_pass": first_pass_args, "second_pass": second_pass_args
-        }
+        first_pass_args = PIPELINE_CONFIG_YAML.get("first_pass", {}).copy()
+        first_pass_args["guidance_scale"] = float(ui_guidance_scale)
+
+        first_pass_args.pop("num_inference_steps", None)
+
+
+        second_pass_args = PIPELINE_CONFIG_YAML.get("second_pass", {}).copy()
+        second_pass_args["guidance_scale"] = float(ui_guidance_scale)
+
+        second_pass_args.pop("num_inference_steps", None)
+        
+        multi_scale_call_kwargs = call_kwargs.copy()
+        first_pass_config_from_yaml = PIPELINE_CONFIG_YAML.get("first_pass", {})
+
+        single_pass_call_kwargs["timesteps"] = first_pass_config_from_yaml.get("timesteps")
+        single_pass_call_kwargs["guidance_scale"] = float(ui_guidance_scale)
+        single_pass_call_kwargs["stg_scale"] = first_pass_config_from_yaml.get("stg_scale")
+        single_pass_call_kwargs["rescaling_scale"] = first_pass_config_from_yaml.get("rescaling_scale")
+        single_pass_call_kwargs["skip_block_list"] = first_pass_config_from_yaml.get("skip_block_list")
         
-        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}")
+
+        single_pass_call_kwargs.pop("num_inference_steps", None) 
+        single_pass_call_kwargs.pop("first_pass", None) 
+        single_pass_call_kwargs.pop("second_pass", None)
         result_images_tensor = pipeline_instance(**single_pass_call_kwargs).images
 
     if result_images_tensor is None:
-        raise RuntimeError("Generation failed, result tensor is None.")
+        raise RuntimeError("Generation failed.")
 
     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
-    
-    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()
+
+    video_np = np.clip(video_np, 0, 1) 
+    video_np = (video_np * 255).astype(np.uint8)
 
-    video_np = (result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).clip(0, 255).astype(np.uint8)
-    
-    # 使用随机数确保文件名几乎不重复
     timestamp = random.randint(10000, 99999)
-    output_video_path = os.path.join(output_dir, f"output_{timestamp}_{seed_ui}.mp4")
+    output_video_path = f"output_{timestamp}.mp4"
+
     
     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)
-    
-    logger.info(f"Video saved successfully to: {output_video_path}")
+        with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], macro_block_size=1) as video_writer:
+            for frame_idx in range(video_np.shape[0]):
+
+                video_writer.append_data(video_np[frame_idx])
+                if frame_idx % 10 == 0:
+                    print(f"Saving frame {frame_idx + 1}/{video_np.shape[0]}")
+    except Exception as e:
+        print(f"Error saving video with macro_block_size=1: {e}")
+        try:
+            with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], format='FFMPEG', codec='libx264', quality=8) as video_writer:
+                for frame_idx in range(video_np.shape[0]):
+
+                    video_writer.append_data(video_np[frame_idx])
+                    if frame_idx % 10 == 0:
+                        print(f"Saving frame {frame_idx + 1}/{video_np.shape[0]} (fallback)")
+        except Exception as e2:
+            print(f"Fallback video saving error: {e2}")
+            raise RuntimeError(f"Failed to save video: {e2}")
+            
     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)}")
+def main():
+    parser = argparse.ArgumentParser(description="LTX Video Generation from Command Line")
+    parser.add_argument("--prompt", required=True, help="Text prompt for video generation")
+    parser.add_argument("--negative-prompt", default="worst quality, inconsistent motion, blurry, jittery, distorted", 
+                        help="Negative prompt")
+    parser.add_argument("--mode", choices=["text-to-video", "image-to-video", "video-to-video"], 
+                        default="text-to-video", help="Generation mode")
+    parser.add_argument("--input-image", help="Input image path for image-to-video mode")
+    parser.add_argument("--input-video", help="Input video path for video-to-video mode")
+    parser.add_argument("--duration", type=float, default=2.0, help="Video duration in seconds (0.3-8.5)")
+    parser.add_argument("--height", type=int, default=512, help="Video height (must be divisible by 32)")
+    parser.add_argument("--width", type=int, default=704, help="Video width (must be divisible by 32)")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--randomize-seed", action="store_true", help="Use random seed")
+    parser.add_argument("--guidance-scale", type=float, help="Guidance scale for generation")
+    parser.add_argument("--no-improve-texture", action="store_true", help="Disable texture improvement (faster)")
+    parser.add_argument("--frames-to-use", type=int, default=9, help="Frames to use from input video (for video-to-video)")
+
+
+
+    
+    args = parser.parse_args()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+    
+    # Validate parameters
+    if args.mode == "image-to-video" and not args.input_image:
+        print("Error: --input-image is required for image-to-video mode")
+        return
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+    
+    if args.mode == "video-to-video" and not args.input_video:
+        print("Error: --input-video is required for video-to-video mode")
+        return
+
+    
+    # Ensure dimensions are divisible by 32
+    args.height = ((args.height - 1) // 32 + 1) * 32
+    args.width = ((args.width - 1) // 32 + 1) * 32
+    
+    print(f"Starting video generation...")
+    print(f"Prompt: {args.prompt}")
+    print(f"Mode: {args.mode}")
+    print(f"Duration: {args.duration}s")
+    print(f"Resolution: {args.width}x{args.height}")
     
     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
+            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}")
+        
+        print(f"\nVideo generation completed!")
+        print(f"Output saved to: {output_path}")
+        print(f"Used seed: {used_seed}")
         
     except Exception as e:
-        logger.error(f"❌ Error during generation: {e}", exc_info=True)
-        sys.exit(1)
-
+        print(f"Error during generation: {e}")
+        raise
 
-# ==============================================================================
-# 主入口和参数解析
-# ==============================================================================
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="LTX Video Generation and Server Client")
+    if os.path.exists(models_dir) and os.path.isdir(models_dir):
+        print(f"Model directory: {Path(models_dir).resolve()}")
     
-    # --- 模式选择 ---
-    group = parser.add_argument_group('运行模式')
-    group.add_argument("--listen", action="store_true", help="以监听模式运行，连接到服务器等待指令。")
-    
-    # --- 监听模式参数 ---
-    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("监听模式已停止。")
-    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)
+    main()