code/generate.py

# code/generate.py
import subprocess, json, tempfile, os, pathlib, torch

REPO = pathlib.Path(__file__).resolve().parents[1]

def generate_word_image(cfg, device):
    """
    cfg   : plain dict or EasyDict
    device: torch.device  (ignored here; CLI handles CUDA)
    Returns absolute path of the rendered PNG.
    """
    with tempfile.TemporaryDirectory() as tmp:
        cfg_path = pathlib.Path(tmp) / "cfg.json"
        with open(cfg_path, "w") as f:
            json.dump(cfg, f)

        # Call the original CLI exactly like your bash script
        cmd = [
            "python", os.fspath(REPO / "code" / "main.py"),
            "--config", os.fspath(cfg_path),
        ]
        subprocess.check_call(cmd)

        # main.py saves into cfg['log_dir']/…/final.png – read it back
        out_png = next((REPO / cfg["log_dir"]).rglob("*.png"))
        return os.fspath(out_png)