1#!/usr/bin/env python3
2"""Kling v3 pro motion control — enhance OmniHuman output with better quality."""
3import argparse
4import json
5import os
6from pathlib import Path
7from urllib.request import urlopen
8 
9 
10def load_fal_key() -> str:
11    for name in ("FAL_KEY", "FAL_AI_KEY", "FAL_API_KEY"):
12        value = os.getenv(name)
13        if value:
14            return value.strip()
15    raise SystemExit("Missing fal.ai key. Set FAL_KEY or FAL_AI_KEY.")
16 
17 
18def main() -> None:
19    parser = argparse.ArgumentParser(description="Kling v3 pro motion control: image + motion video → enhanced video")
20    parser.add_argument("--image", required=True, help="Character reference image.")
21    parser.add_argument("--video", required=True, help="Motion reference video (e.g. OmniHuman output).")
22    parser.add_argument("--face-image", help="Frontal face image for identity preservation.")
23    parser.add_argument("--prompt", default="", help="Guidance prompt.")
24    parser.add_argument("--output", required=True, help="Output mp4 path.")
25    parser.add_argument("--json-out", help="Optional path to save raw fal response JSON.")
26    parser.add_argument("--model", default="fal-ai/kling-video/v3/pro/motion-control")
27    args = parser.parse_args()
28 
29    os.environ["FAL_KEY"] = load_fal_key()
30    import fal_client
31 
32    image_path = Path(args.image).expanduser().resolve()
33    video_path = Path(args.video).expanduser().resolve()
34    output_path = Path(args.output).expanduser().resolve()
35    output_path.parent.mkdir(parents=True, exist_ok=True)
36 
37    print(f"Uploading image: {image_path}")
38    image_url = fal_client.upload_file(str(image_path))
39    print(f"Uploading video: {video_path}")
40    video_url = fal_client.upload_file(str(video_path))
41 
42    payload = {
43        "image_url": image_url,
44        "video_url": video_url,
45        "character_orientation": "video",
46        "keep_original_sound": True,
47    }
48 
49    if args.prompt.strip():
50        payload["prompt"] = args.prompt.strip()
51 
52    if args.face_image:
53        face_path = Path(args.face_image).expanduser().resolve()
54        print(f"Uploading face image: {face_path}")
55        face_url = fal_client.upload_file(str(face_path))
56        payload["elements"] = [{"frontal_image_url": face_url, "reference_image_urls": [face_url]}]
57 
58    print(f"Submitting to {args.model}...")
59 
60    def _print_logs(update):
61        status = getattr(update, "status", None)
62        if status:
63            print(f"STATUS: {status}")
64 
65    result = fal_client.subscribe(
66        args.model,
67        arguments=payload,
68        with_logs=True,
69        on_queue_update=_print_logs,
70    )
71 
72    if args.json_out:
73        json_out = Path(args.json_out).expanduser().resolve()
74        json_out.parent.mkdir(parents=True, exist_ok=True)
75        json_out.write_text(json.dumps(result, ensure_ascii=False, indent=2))
76 
77    video = result.get("video") if isinstance(result, dict) else None
78    if not isinstance(video, dict) or not video.get("url"):
79        raise SystemExit(f"Unexpected result: {json.dumps(result, ensure_ascii=False)[:2000]}")
80 
81    print(f"Downloading: {video['url']}")
82    with urlopen(video["url"], timeout=300) as resp:
83        output_path.write_bytes(resp.read())
84    print(f"Saved: {output_path}")
85    print(json.dumps({"ok": True, "output": str(output_path)}))
86 
87 
88if __name__ == "__main__":
89    main()
90