"""
generuj-wideo.py - Wan 2.2 TI2V 5BとローカルComfyUIで動画を生成
作者: Piotr Olszewski。スクリプトのライセンス: MIT。2026年9月26日にComfyUI 0.37、RTX 4090で確認。

必要なもの: Wan 2.2 5Bのファイルを配置したComfyUI (http://127.0.0.1:8188)
(instaluj-comfyui.ps1 -Modele wan) と、MP4作成用のPATHが通ったffmpeg (winget install Gyan.FFmpeg)。
ffmpegがない場合はPNGフレームのみ出力します。Python 3.9以降の標準ライブラリーのみを使用します。

実行例:
  python generuj-wideo.py "a yellow tram passing a rainy street at dusk, static camera, realistic"
  python generuj-wideo.py "flat 2D animation of a cat walking across a desk" --sekundy 3 --szer 960 --wys 544

RTX 4090の場合: 1280x704の5秒動画は約4分（最初の動画は18 GBのファイル読み込み分だけ長くなります）。
モデルは固定カメラと背景の動きを得意とし、液体を注ぐ動作や仕草など、正確な物理表現は苦手です。
"""
import argparse, json, os, random, shutil, subprocess, sys, tempfile, time, urllib.parse, urllib.request, uuid

FPS = 24
NEG = "blurry, distorted, low quality, jitter, watermark, text, subtitles, static image, deformed hands, extra limbs"


def graf(polecenie, szer, wys, klatki, kroki, seed, prefiks):
    # ComfyUI公式の「Wan 2.2 5B text to video」テンプレートに相当
    return {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "wan2.2_ti2v_5B_fp16.safetensors", "weight_dtype": "default"}},
        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors", "type": "wan", "device": "default"}},
        "3": {"class_type": "VAELoader", "inputs": {"vae_name": "wan2.2_vae.safetensors"}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": polecenie}},
        "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": NEG}},
        "6": {"class_type": "ModelSamplingSD3", "inputs": {"model": ["1", 0], "shift": 8.0}},
        "7": {"class_type": "Wan22ImageToVideoLatent", "inputs": {"vae": ["3", 0], "width": szer, "height": wys, "length": klatki, "batch_size": 1}},
        "8": {"class_type": "KSampler", "inputs": {"model": ["6", 0], "seed": seed, "steps": kroki, "cfg": 5.0, "sampler_name": "uni_pc", "scheduler": "simple",
                                                    "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["7", 0], "denoise": 1.0}},
        "9": {"class_type": "VAEDecode", "inputs": {"samples": ["8", 0], "vae": ["3", 0]}},
        "10": {"class_type": "SaveImage", "inputs": {"images": ["9", 0], "filename_prefix": prefiks}},
    }


def api(serwer, sciezka, dane=None):
    req = urllib.request.Request(serwer + sciezka, data=json.dumps(dane).encode() if dane is not None else None,
                                 headers={"Content-Type": "application/json"})
    return json.load(urllib.request.urlopen(req, timeout=60))


def main():
    a = argparse.ArgumentParser(description="Wan 2.2 5Bで動画を生成（ローカルComfyUI）")
    a.add_argument("polecenie", help="シーンの説明（英語を推奨。カメラワークも指定してください）")
    a.add_argument("--sekundy", type=float, default=5.0, help="動画の長さ（初期値: 5秒）")
    a.add_argument("--szer", type=int, default=1280); a.add_argument("--wys", type=int, default=704)
    a.add_argument("--kroki", type=int, default=20)
    a.add_argument("--seed", type=int, default=None)
    a.add_argument("--serwer", default="http://127.0.0.1:8188")
    a.add_argument("--wyjscie", default="wideo")
    x = a.parse_args()
    if x.szer % 32 or x.wys % 32:
        a.error("幅と高さは32の倍数にしてください（例: 1280x704、960x544）")
    klatki = int(round(x.sekundy * FPS / 4)) * 4 + 1          # Wanが必要とする4n+1形式
    try:
        api(x.serwer, "/system_stats")
    except Exception as e:
        sys.exit(f"ComfyUI ({x.serwer}) に接続できません ({e})。start-comfyui.batを起動してください。")
    seed = x.seed if x.seed is not None else random.randint(1, 2**31)
    prefiks = f"wan_{uuid.uuid4().hex[:8]}"
    print(f"{klatki}フレーム（{klatki / FPS:.1f}秒）を{x.szer}x{x.wys}で生成します。シード: {seed}...")
    t0 = time.time()
    pid = api(x.serwer, "/prompt", {"prompt": graf(x.polecenie, x.szer, x.wys, klatki, x.kroki, seed, prefiks), "client_id": str(uuid.uuid4())})["prompt_id"]
    while True:
        h = api(x.serwer, f"/history/{pid}")
        if pid in h:
            break
        print(f"\r  処理中... {time.time() - t0:.0f}秒", end="", flush=True)
        time.sleep(2)
    print()
    wynik = h[pid]
    if wynik.get("status", {}).get("status_str") != "success":
        sys.exit("ComfyUIエラー: " + json.dumps(wynik.get("status", {}), ensure_ascii=False)[:600])
    imgs = sorted([o for v in wynik["outputs"].values() for o in v.get("images", [])], key=lambda o: o["filename"])
    os.makedirs(x.wyjscie, exist_ok=True)
    tmp = tempfile.mkdtemp(prefix="wan_")
    for i, img in enumerate(imgs):
        q = urllib.parse.urlencode({"filename": img["filename"], "subfolder": img.get("subfolder", ""), "type": img.get("type", "output")})
        with urllib.request.urlopen(f"{x.serwer}/view?{q}", timeout=60) as r, open(os.path.join(tmp, f"{i:04d}.png"), "wb") as f:
            shutil.copyfileobj(r, f)
    czas = time.time() - t0
    nazwa = os.path.join(x.wyjscie, f"wan_{time.strftime('%Y%m%d_%H%M%S')}_{seed}")
    if shutil.which("ffmpeg"):
        subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-framerate", str(FPS), "-i", os.path.join(tmp, "%04d.png"),
                        "-c:v", "libx264", "-crf", "20", "-pix_fmt", "yuv420p", "-movflags", "+faststart", nazwa + ".mp4"], check=True)
        shutil.rmtree(tmp, ignore_errors=True)
        print(f"完了: {nazwa}.mp4（{len(imgs)}フレーム、{czas:.0f}秒）")
    else:
        shutil.move(tmp, nazwa + "_klatki")
        print(f"完了: {nazwa}_klatki にPNGフレームを保存しました。MP4を作るにはffmpegをインストールしてください。")


if __name__ == "__main__":
    main()