"""
generuj-obrazy.py - seria obrazów z Z-Image Turbo przez lokalne ComfyUI, bez klikania (promptowy.com/z-image-turbo/)
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 Z-Image Turbo w folderach models
(skrypt instaluj-comfyui.ps1 pobiera je automatycznie). Tylko biblioteka standardowa Pythona 3.9+.

Przykłady:
  python generuj-obrazy.py "zdjęcie produktowe kubka ceramicznego na betonie"
  python generuj-obrazy.py "plakat z górami o świcie" --ile 4 --szer 1024 --wys 1024
  python generuj-obrazy.py --plik polecenia.txt          (jedno polecenie w linii)

Wynik: folder ./obrazy/ + plik obrazy/log.csv z czasem każdego obrazu i szczytem pamięci karty.
Uwaga: polskie litery na grafice (np. „ŻÓŁĆ”) model zwykle psuje - napisy dodawaj w edytorze.
"""
import argparse, csv, json, os, random, shutil, subprocess, sys, threading, time, urllib.parse, urllib.request, uuid

MODEL = {"unet": "z_image_turbo_bf16.safetensors", "clip": "qwen_3_4b.safetensors", "vae": "ae.safetensors"}


def graf(polecenie, szer, wys, kroki, seed, prefiks):
    # odpowiednik oficjalnego szablonu ComfyUI „Z-Image Turbo text to image”
    return {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": MODEL["unet"], "weight_dtype": "default"}},
        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": MODEL["clip"], "type": "lumina2", "device": "default"}},
        "3": {"class_type": "VAELoader", "inputs": {"vae_name": MODEL["vae"]}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": polecenie}},
        "5": {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["4", 0]}},
        "6": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["1", 0], "shift": 3.0}},
        "7": {"class_type": "EmptySD3LatentImage", "inputs": {"width": szer, "height": wys, "batch_size": 1}},
        "8": {"class_type": "KSampler", "inputs": {"model": ["6", 0], "seed": seed, "steps": kroki, "cfg": 1.0, "sampler_name": "res_multistep",
                                                    "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}},
    }


class PamiecKarty(threading.Thread):
    """Szczyt zajętej pamięci karty (nvidia-smi co 0,5 s). Bez karty Nvidii po prostu zwraca 0."""
    def __init__(self):
        super().__init__(daemon=True); self.stop = False; self.szczyt = 0
    def run(self):
        while not self.stop:
            try:
                o = subprocess.check_output(["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], text=True, timeout=5)
                self.szczyt = max(self.szczyt, int(o.split()[0]))
            except Exception:
                return
            time.sleep(0.5)


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=30))


def main():
    a = argparse.ArgumentParser(description="Seria obrazów z Z-Image Turbo przez lokalne ComfyUI")
    a.add_argument("polecenie", nargs="?", help="opis obrazu (najlepiej po angielsku)")
    a.add_argument("--plik", help="plik tekstowy z poleceniami, jedno w linii")
    a.add_argument("--ile", type=int, default=1, help="ile wariantów każdego polecenia (domyślnie 1)")
    a.add_argument("--szer", type=int, default=1344); a.add_argument("--wys", type=int, default=768)
    a.add_argument("--kroki", type=int, default=9, help="9 jak w moim teście, szablon ComfyUI ma 8")
    a.add_argument("--seed", type=int, default=None, help="stałe ziarno = powtarzalny wynik")
    a.add_argument("--serwer", default="http://127.0.0.1:8188")
    a.add_argument("--wyjscie", default="obrazy")
    x = a.parse_args()
    polecenia = [x.polecenie] if x.polecenie else []
    if x.plik:
        polecenia += [l.strip() for l in open(x.plik, encoding="utf-8") if l.strip() and not l.startswith("#")]
    if not polecenia:
        a.error("podaj polecenie albo --plik")
    try:
        stat = 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 i spróbuj ponownie.")
    os.makedirs(x.wyjscie, exist_ok=True)
    log = open(os.path.join(x.wyjscie, "log.csv"), "a", newline="", encoding="utf-8")
    w = csv.writer(log)
    nr = 0
    for p in polecenia:
        for i in range(x.ile):
            nr += 1
            seed = x.seed + i if x.seed is not None else random.randint(1, 2**31)   # kolejne warianty: ziarno +1
            prefiks = f"promptowy_{uuid.uuid4().hex[:8]}"
            mon = PamiecKarty(); mon.start(); t0 = time.time()
            pid = api(x.serwer, "/prompt", {"prompt": graf(p, x.szer, x.wys, 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
                time.sleep(0.5)
            czas = time.time() - t0; mon.stop = True
            wynik = h[pid]
            if wynik.get("status", {}).get("status_str") != "success":
                print(f"[{nr}] BŁĄD: {json.dumps(wynik.get('status', {}), ensure_ascii=False)[:400]}"); continue
            img = [o for v in wynik["outputs"].values() for o in v.get("images", [])][0]
            q = urllib.parse.urlencode({"filename": img["filename"], "subfolder": img.get("subfolder", ""), "type": img.get("type", "output")})
            cel = os.path.join(x.wyjscie, f"{nr:03d}_{seed}.png")
            with urllib.request.urlopen(f"{x.serwer}/view?{q}", timeout=60) as r, open(cel, "wb") as f:
                shutil.copyfileobj(r, f)
            w.writerow([time.strftime("%Y-%m-%d %H:%M:%S"), cel, x.szer, x.wys, x.kroki, seed, round(czas, 1), mon.szczyt, p]); log.flush()
            print(f"[{nr}] {cel}  {czas:.1f} s  pamięć karty do {mon.szczyt} MB")
    print("Gotowe. Pierwszy obraz trwa dłużej, bo model wczytuje się z dysku.")


if __name__ == "__main__":
    main()