wan-22-comfy/client.py

#!/usr/bin/env python3
"""
Клиент для DaSiWa API Server.
Запускается на ТВОЁМ ПК. Отправляет подписанные запросы на сервер.

Использование:
    python client.py --server http://<ip>:5000 --image photo.png --prompt "woman dancing"
    python client.py --server http://<ip>:5000 --image start.png --last-image end.png --prompt "smooth transition"
"""

import argparse
import base64
import json
import os
import sys
import time

import requests

from hmac_auth import sign_request

# ============================================================================
# Конфигурация
# ============================================================================

KEYS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "keys.json")


def load_keys():
    if not os.path.exists(KEYS_FILE):
        print(f"❌ Файл ключей не найден: {KEYS_FILE}")
        print("   Запусти: python generate_keys.py")
        sys.exit(1)
    with open(KEYS_FILE, "r") as f:
        return json.load(f)


def image_to_base64(path: str) -> str:
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode()


def send_request(server_url: str, payload: dict, client_id: str, secret_key: str) -> dict:
    """Отправляет подписанный запрос на сервер."""
    body = json.dumps(payload).encode("utf-8")
    auth_headers = sign_request(body, secret_key, client_id)

    headers = {
        "Content-Type": "application/json",
        **auth_headers
    }

    response = requests.post(
        f"{server_url}/generate",
        data=body,
        headers=headers,
        timeout=600
    )

    return response.status_code, response.json()


def main():
    parser = argparse.ArgumentParser(description="DaSiWa API Client")
    parser.add_argument("--server", required=True, help="Server URL, e.g. http://1.2.3.4:5000")
    parser.add_argument("--image", required=True, help="Path to first frame image")
    parser.add_argument("--last-image", default=None, help="Path to last frame image (FLF2V mode)")
    parser.add_argument("--prompt", required=True, help="Text prompt")
    parser.add_argument("--negative-prompt", default=None, help="Negative prompt")
    parser.add_argument("--width", type=int, default=528)
    parser.add_argument("--height", type=int, default=768)
    parser.add_argument("--length", type=int, default=81, help="Frame count")
    parser.add_argument("--steps", type=int, default=4)
    parser.add_argument("--cfg", type=float, default=1.0)
    parser.add_argument("--seed", type=int, default=-1)
    parser.add_argument("--fps", type=int, default=16)
    parser.add_argument("--output", "-o", default="output.mp4", help="Output video path")
    args = parser.parse_args()

    keys = load_keys()

    # Формируем payload
    payload = {
        "prompt": args.prompt,
        "image_base64": image_to_base64(args.image),
        "width": args.width,
        "height": args.height,
        "length": args.length,
        "steps": args.steps,
        "cfg": args.cfg,
        "seed": args.seed,
        "fps": args.fps,
    }

    if args.negative_prompt:
        payload["negative_prompt"] = args.negative_prompt

    if args.last_image:
        payload["last_image_base64"] = image_to_base64(args.last_image)
        print(f"🎬 Режим: FLF2V (first + last frame)")
    else:
        print(f"🎬 Режим: I2V (image to video)")

    print(f"📐 {args.width}x{args.height}, {args.length} frames, {args.steps} steps")
    print(f"📤 Отправляю запрос на {args.server}...")

    start = time.time()
    status_code, result = send_request(
        args.server, payload, keys["client_id"], keys["secret_key"]
    )
    elapsed = time.time() - start

    if status_code != 200:
        print(f"❌ Ошибка {status_code}: {result.get('error', 'Unknown')}")
        if "detail" in result:
            print(f"   Детали: {result['detail']}")
        sys.exit(1)

    if "video" in result:
        video_bytes = base64.b64decode(result["video"])
        with open(args.output, "wb") as f:
            f.write(video_bytes)
        print(f"✅ Видео сохранено: {args.output} ({len(video_bytes) / 1024 / 1024:.1f} MB)")
        print(f"⏱  Время: {elapsed:.1f}s (сервер: {result.get('elapsed', '?')}s)")
        print(f"🌱 Seed: {result.get('seed', '?')}")
    else:
        print(f"❌ Ошибка: {result.get('error', 'No video in response')}")
        sys.exit(1)


if __name__ == "__main__":
    main()