downloader.py

# ==============================================================================
# 统一入口和依赖项
# ==============================================================================
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}")


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

    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))
    
    if ui_guidance_scale is None:
        ui_guidance_scale = PIPELINE_CONFIG_YAML.get("first_pass", {}).get("guidance_scale", 1.0)
    
    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}")
            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.")
        
        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
        }
        
        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.")

    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() * 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")
    
    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}")
    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)}")
    
    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")
    
    # --- 模式选择 ---
    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)