"""
generuj-wideo.py - klip z Wan 2.2 TI2V 5B przez lokalne ComfyUI (promptowy.com/wan-2-2/)
Autor: Piotr Olszewski, promptowy.com. Licencja skryptu: MIT. Sprawdzone 26.09.2026: ComfyUI 0.37, RTX 4090.

Wymaga: działającego ComfyUI (http://127.0.0.1:8188) z plikami Wan 2.2 5B
(instaluj-comfyui.ps1 -Modele wan) oraz ffmpeg w PATH do złożenia MP4 (winget install Gyan.FFmpeg).
Bez ffmpeg zostają same klatki PNG. Tylko biblioteka standardowa Pythona 3.9+.

Przykłady:
  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

Na RTX 4090: 5 s w 1280x704 to ok. 4 minuty (pierwszy klip dłużej - wczytanie 18 GB z dysku).
Model najlepiej radzi sobie ze statyczną kamerą i ruchem w tle, słabo z precyzyjną fizyką (nalewanie, gesty).
"""
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):
    # odpowiednik oficjalnego szablonu 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="Klip wideo z Wan 2.2 5B przez lokalne ComfyUI")
    a.add_argument("polecenie", help="opis ujęcia (najlepiej po angielsku, z opisem kamery)")
    a.add_argument("--sekundy", type=float, default=5.0, help="długość klipu (domyślnie 5 s)")
    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("szerokość i wysokość muszą być wielokrotnością 32 (np. 1280x704, 960x544)")
    klatki = int(round(x.sekundy * FPS / 4)) * 4 + 1          # Wan wymaga długości 4n+1
    try:
        api(x.serwer, "/system_stats")
    except Exception as e:
        sys.exit(f"Nie mogę połączyć się z ComfyUI pod {x.serwer} ({e}). Uruchom 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"Generuję {klatki} klatek ({klatki / FPS:.1f} s) w {x.szer}x{x.wys}, ziarno {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  trwa... {time.time() - t0:.0f} s", end="", flush=True)
        time.sleep(2)
    print()
    wynik = h[pid]
    if wynik.get("status", {}).get("status_str") != "success":
        sys.exit("Błąd 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"Gotowe: {nazwa}.mp4 ({len(imgs)} klatek, {czas:.0f} s)")
    else:
        shutil.move(tmp, nazwa + "_klatki")
        print(f"Gotowe: klatki PNG w {nazwa}_klatki ({czas:.0f} s). Zainstaluj ffmpeg, żeby dostać MP4.")


if __name__ == "__main__":
    main()