viz_agent.py

OPENROUTER_API_KEY_DEFAULT = ""
import torch
import os, os.path as osp
import logging
import numpy as np
from lib_4d.camera import SimpleFovCamerasIndependent
from lib_4d.gs_static_model import StaticGaussian
from lib_4d.gs_ed_model import DynSCFGaussian
import imageio
from omegaconf import OmegaConf
from lib_4d.render_helper import GS_BACKEND
from lib_4d.figure_viz_helper import *

import random
import time 
import requests 
import json 
import cv2, base64

# Azure OpenAI optional import guard
try:
    from openai import AzureOpenAI as _AzureOpenAI  # type: ignore
except Exception:
    _AzureOpenAI = None

API_PROVIDER = os.getenv("VIZ_AGENT_API", "openrouter")  # overridden by CLI --api

logging.getLogger("imageio_ffmpeg").setLevel(logging.ERROR)

def gptv_reflection_prompt_selectbest(user_prompt, img_prompt, idea_transcript, listofimages, args):
    num_img = len(listofimages)
    transcript = [{ "role": "system", "content": [] }, {"role": "user", "content": []}]
    # System prompt
    transcript[0]["content"].append("You are a helpful assistant.\n\nYou are a judge to rank provided images. Below are %d images generated by an AI art generation model, indexed from 0 to %d."%(num_img,num_img-1))
    transcript[0]["content"].append("From scale 1 to 10, decide how similar each image is to the user imagined IDEA of the scene.")

    # ## Example & Query prompt
    # if args.select_fewshot:
    #     transcript[-1]["content"] = transcript[-1]["content"] + prepare_fewshot_selectbest(user_prompt, img_prompt, listofimages, args)

    transcript[-1]["content"] = transcript[-1]["content"] + idea_transcript
    for img_i in range(num_img):
        transcript[-1]["content"].append("%d. "%img_i)
        transcript[-1]["content"].append(load_img(listofimages[img_i]))

    transcript[-1]["content"].append("Let's think step by step. Check all aspects to see how well these images strictly follow the content in IDEA, including having correct object counts, attributes, entities, relationships, sizes, appearance, and all other descriptions in the IDEA. Then give a score for each input images. Finally, consider the scores and select the image with the best overall quality with image index 0 to %d wrapped with <START> and <END>. Only wrap single image index digits between <START> and <END>."%(num_img-1))

    response = gptv_query(transcript)
    if '<START>' not in response or '<END>' not in response: ## one format retry
        response = gptv_query(transcript, temp=0.1)
    if args.verbose:
        print('gptv_reflection_prompt_selectbest\n %s\n'%(response))
    if '<START>' not in response or '<END>' not in response:
        return random.randint(0,num_img-1), response
    prompts = response.split('<START>')[1]
    prompts = prompts.strip().split('<END>')[0]
    return int(prompts) if prompts.isdigit() else random.randint(0,num_img-1), response

def gptv_query(transcript=None, temp=0.):
    max_tokens = 1500 #512
    wait_time = 10

    provider = API_PROVIDER
    # Azure OpenAI GPT-4.1 path
    if provider == "xh-gpt4.1":
        endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "https://2024alextest.openai.azure.com/")
        deployment = os.getenv("AZURE_OPENAI_DEPLOYMENT", "gpt-4-1-2025-04-14")
        api_version = os.getenv("AZURE_OPENAI_API_VERSION", "2024-12-01-preview")
        subscription_key = os.getenv("AZURE_OPENAI_API_KEY", "")
        if _AzureOpenAI is None:
            raise RuntimeError("AzureOpenAI SDK not available. Please install openai>=1.42 and set AZURE_OPENAI_API_KEY.")
        client = _AzureOpenAI(api_version=api_version, azure_endpoint=endpoint, api_key=subscription_key)
        # Normalize transcript to OpenAI content parts format required by GPT-4.1
        def _to_content_parts(content):
            # Convert legacy strings/lists to array of content-part objects
            parts = []
            if isinstance(content, str):
                return [{"type": "text", "text": content}]
            if isinstance(content, list):
                for item in content:
                    if isinstance(item, str):
                        parts.append({"type": "text", "text": item})
                    elif isinstance(item, dict):
                        # Preserve objects that already look like content parts
                        if "type" in item:
                            parts.append(item)
                        elif "image_url" in item:
                            parts.append({"type": "image_url", "image_url": item.get("image_url")})
                        else:
                            parts.append({"type": "text", "text": json.dumps(item, ensure_ascii=False)})
                    else:
                        parts.append({"type": "text", "text": str(item)})
                return parts
            # Fallback stringify
            return [{"type": "text", "text": str(content)}]

        messages = []
        for m in (transcript or []):
            if not isinstance(m, dict):
                # Fallback: treat as user text
                messages.append({"role": "user", "content": [{"type": "text", "text": str(m)}]})
                continue
            role = m.get("role", "user")
            content = m.get("content", "")
            messages.append({"role": role, "content": _to_content_parts(content)})
        # Temperature/top_p mapping
        completion = None
        while completion is None:
            try:
                completion = client.chat.completions.create(
                    messages=messages,
                    max_completion_tokens=max(1, max_tokens),
                    temperature=float(temp),
                    top_p=0.5,
                    frequency_penalty=0.0,
                    presence_penalty=0.0,
                    model=deployment,
                )
            except Exception as e:
                print(e)
                time.sleep(wait_time)
                if float(temp) < 1.0:
                    temp = min(float(temp) + 0.2, 1.0)
                continue
        return completion.choices[0].message.content

    # Default: OpenRouter path
    api_key = os.getenv("OPENROUTER_API_KEY", OPENROUTER_API_KEY_DEFAULT)
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    # Optional headers for OpenRouter rankings/telemetry
    referer = os.getenv("OPENROUTER_HTTP_REFERER") or os.getenv("HTTP_REFERER")
    title = os.getenv("OPENROUTER_X_TITLE") or os.getenv("X_TITLE")
    if referer:
        headers["HTTP-Referer"] = referer
    if title:
        headers["X-Title"] = title
    data = {
        'model': 'qwen/qwen2.5-vl-72b-instruct:free',
        'max_tokens':max_tokens, 
        'temperature': temp,
        'top_p': 0.5,
        'messages':[]
    }
    if transcript is not None:
        data['messages'] = transcript

    response_text, retry, response_json = '', 0, None
    while len(response_text)<2:
        retry += 1
        try:
            response = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, data=json.dumps(data)) 
            response_json = response.json()
        except Exception as e:
            if random.random()<1: print(e)
            time.sleep(wait_time)
            continue
        if response.status_code != 200:
            print(response.headers,response.content)
            if random.random()<0.01: print(f"The response status code for is {response.status_code} (Not OK)")
            time.sleep(wait_time)
            data['temperature'] = min(data['temperature'] + 0.2, 1.0)
            continue
        if 'choices' not in response_json:
            time.sleep(wait_time)
            continue
        response_text = response_json["choices"][0]["message"]["content"]
    return response_json["choices"][0]["message"]["content"]

def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def load_img(image_path):
    base64_image = encode_image(image_path)
    image_meta = "data:image/png;base64" if 'png' in image_path else "data:image/jpeg;base64"
    img_dict = {
        "type": "image_url",
        "image_url": {
          "url": f"{image_meta},{base64_image}",
          "detail": "low"
        }
    }
    return img_dict

def gptv_init_prompt(user_prompt, img_prompt, idea_transcript, args):
    transcript = [{ "role": "system", "content": [] }, {"role": "user", "content": []}]
    # System prompt
    transcript[0]["content"].append("You are a helpful assistant.\n\nInstruction: Given a user imagined IDEA of the scene, converting the IDEA into the following format. If user says they want to change color, then operation should be color_func and nothing else. If user says they want to extract an object, then operation should be extraction and nothing else. If user says they want to delete an object, then operation should be deletion and nothing else.\n")
    transcript[0]["content"].append("edit: \n ")
    transcript[0]["content"].append(', '.join(["objects:" + "["+"cow" , "grass", "ground", "sidewalk", "car", "road" + "] \n"]))
    transcript[0]["content"].append("operations:" + "deletion/extraction/color_func" + "\n")
    transcript[0]["content"].append("targets:" +  "car" + "\n" + "threshold: 0.191 \n")
    # transcript[0]["content"].append("- Each prompt should consist of a description of the scene followed by modifiers divided by commas.\n- The modifiers should alter the mood, style, lighting, and other aspects of the scene.\n- Multiple modifiers can be used to provide more specific details.\n- When generating prompts, reduce abstract psychological and emotional descriptions.\n- When generating prompts, explain images and unusual entities in IDEA with detailed descriptions of the scene.\n- Do not mention 'given image' in output, use detailed texts to describe the image in IDEA instead.\n- Generate diverse prompts.\n- Each prompt should have no more than 50 words.\n")
    # transcript[0]["content"].append("""function: lambda color: (color + torch.Tensor([adjustment]).to(color)) if isinstance(adjustment,(list, tuple)) else color * adjustment
    #                                 # This function adjusts the color tensor dynamically based on the provided adjustment.
    #                                 # It supports additive, subtractive and multiplicative adjustments, allowing for various transformations.
    #                                 # Examples:
    #                                 # Adjust Red: lambda color: (color + torch.Tensor([[R, 0, 0]]).to(color))  # R ∈ [-4, 4], positive adds red, negative removes red
    #                                 # Adjust Green: lambda color: (color + torch.Tensor([[0, G, 0]]).to(color))  # G ∈ [-4, 4], positive adds green, negative removes green
    #                                 # Adjust Blue: lambda color: (color + torch.Tensor([[0, 0, B]]).to(color))  # B ∈ [-4, 4], positive adds blue, negative removes blue
    #                                 # Increase Brightness: lambda color: (color + torch.Tensor([[V, V, V]]).to(color))  # V ∈ [-4, 4], positive brightens, negative darkens
    #                                 # Decrease Brightness: lambda color: color * S  # Scale factor S ∈ (0, 1] reduces brightness
    #                                 # Reduce a Specific Color: lambda color: (color - torch.Tensor([[R, G, B]]).to(color))  # Explicit subtraction of RGB values
    #                                 # Adjust Saturation: lambda color: color * S + (1 - S) * torch.mean(color, dim=-1, keepdim=True)  # S ∈ [0,1], lower reduces saturation
    #                                 # Convert to Grayscale (Black and white): lambda color: torch.mean(color, dim=-1, keepdim=True).expand_as(color)
    #                                 # Additional operations:
    #                                 # - Combine multiple adjustments for more complex effects
    #                                 # - Use negative values to reduce specific channels
    #                                 # - Use **negative values** to explicitly remove color components to make dark colors.
    #                                 # Example for Brown: lambda color: (color + torch.Tensor([[R, G, -B]]).to(color)) # Increase R & G, reduce complement B a lot
    #                                 # - Keep values within valid RGB ranges to avoid overflow
    #                                 # Feel free to experiment with different values for precise color control.""")
    
    transcript[0]["content"].append("""function: lambda color: torch.stack([color[:, 0].masked_fill(color[:, 0] >= -0.7, assign_values[0]),color[:, 1].masked_fill(color[:, 1] >= -0.7, assign_values[1]),color[:, 2].masked_fill(color[:, 2] >= -0.7, assign_values[2])], dim=1) if mode == "assign" else ((color + torch.Tensor([adjustment]).to(color)) if mode == "additive" else color * adjustment)
                                    # This function supports three types of transformations:
                                    # 1. **Direct Assignment (default, mode="assign")** - Overrides color values based on the fixed condition (color >= -0.7).
                                    # 2. **Additive/Subtractive (mode="additive")** - Adjusts RGB values by adding/subtracting a tensor.
                                    # 3. **Multiplicative Scaling (mode="multiplicative")** - Adjusts brightness/contrast by scaling color values.
                                    # The condition `color >= -0.7` is fixed for assignments to maintain color transformation stability.
                                    # `mode="assign"` (Used for Specific Colors) Example:
                                    # lambda color: torch.stack([color[:, 0].masked_fill(color[:, 0] >= -0.7, R), color[:, 1].masked_fill(color[:, 1] >= -0.7, G),color[:, 2].masked_fill(color[:, 2] >= -0.7, B)], dim=1)
                                    # R ∈ [-2, 4], -2 is black, 4 is bright red
                                    # G ∈ [-2, 4], -2 is black, 4 is bright green
                                    # B ∈ [-2, 4], -2 is black, 4 is bright blue
                                    # `mode="additive"` (Automatically Applied for "More/Less Color") Examples:
                                    # Adjust Red: lambda color: (color + torch.Tensor([[R, 0, 0]]).to(color))  # R ∈ [-4, 4], positive adds red, negative removes red
                                    # Adjust Green: lambda color: (color + torch.Tensor([[0, G, 0]]).to(color))  # G ∈ [-4, 4], positive adds green, negative removes green
                                    # Adjust Blue: lambda color: (color + torch.Tensor([[0, 0, B]]).to(color))  # B ∈ [-4, 4], positive adds blue, negative removes blue
                                    # Adjust Brightness: lambda color: (color + torch.Tensor([[V, V, V]]).to(color))  # V ∈ [-4, 4], positive brightens, negative darkens
                                    # Reduce a Specific Color: lambda color: (color - torch.Tensor([[R, G, B]]).to(color))  # Explicit RGB reduction
                                    # `mode="multiplicative"` Examples:
                                    # Darken: lambda color: color * S  # S ∈ [-1, 1]
                                    # Brighten: lambda color: color * S  # S > 1
                                    # Adjust Saturation: lambda color: color * S + (1 - S) * torch.mean(color, dim=-1, keepdim=True)  # S ∈ [0,1], lower reduces saturation
                                    # Convert to Grayscale (Black and white): lambda color: torch.mean(color, dim=-1, keepdim=True).expand_as(color) 
                                    # Additional Notes:
                                    # - **Default mode is `assign` for setting user-specified colors.**
                                    # - **Use `mode="additive"` for small color agjustment.**
                                    # - `"additive"` mode is **automatically used** for **"more/less [color]"**.
                                    # - **Use `mode="multiplicative"` for contrast and brightness control.**
                                    # - **Keep values within valid ranges to avoid unwanted artifacts.**
                                    # - **Mix different transformations for complex effects.**""")

    #transcript[0]["content"].append("function:lambda color: color[..., [2,1,0]]  # switch appropriate rgb channels(shs)\n")

    transcript[0]["content"] = ' '.join(transcript[0]["content"])
    ## Example & Query prompt
    print('transciprt content', transcript[-1]["content"])
    transcript[-1]["content"] = transcript[-1]["content"] + idea_transcript
    transcript[-1]["content"].append(
        "Based on the above information, you will write %d  prompts exactly about the IDEA follow the rules with different threshold values from %.3f to %.3f, and the potentially required color_func. Each prompt is wrapped with <START> and <END>.\n"
        % (args.num_prompt, float(getattr(args, "thr_min", 0.85)), float(getattr(args, "thr_max", 0.95)))
    )


    
    # tree 0.238 
    transcript[-1]["content"] = ' '.join(transcript[-1]["content"])

    print('transcript', transcript)
    response = gptv_query(transcript, temp=float(getattr(args, "gpt_temp", 0.0)))
    if '<START>' not in response or '<END>' not in response: ## one format retry
        response = gptv_query(transcript, temp=0.1)
    if args.verbose:
        print('gptv_init_prompt    IDEA: %s.\n %s\n'%(user_prompt,response))
    prompts = response.split('<START>')[1:]
    prompts = [x.strip().split('<END>')[0] for x in prompts]
    return prompts

def load_model_cfg(cfg, log_dir, device=torch.device("cuda")):

    # get cfg
    if log_dir.endswith("/"):
        log_dir = log_dir[:-1]
    if isinstance(cfg, str):
        cfg = OmegaConf.load(cfg)
        OmegaConf.set_readonly(cfg, True)

    dataset_mode = getattr(cfg, "dataset_mode", "iphone")
    max_sph_order = getattr(cfg, "max_sph_order", 1)
    logging.info(f"Dataset mode: {dataset_mode}")

    ######################################################################
    ######################################################################

    d_model_ckpt = torch.load(osp.join(log_dir, "finetune_d_model.pth"), weights_only=True)
    d_model = DynSCFGaussian.load_from_ckpt(d_model_ckpt, device=device)
    # * load static model and camera again, because it's also finetuned
    saved_cam = torch.load(osp.join(log_dir, "finetune_s_model_cam.pth"), weights_only=True)
    cams: SimpleFovCamerasIndependent = SimpleFovCamerasIndependent(
        T=len(saved_cam["q_wc"]),
        fovdeg_init=40.0,  # dummy init
    )
    cams.load_state_dict(saved_cam, strict=True)

    s_model = StaticGaussian(
        load_fn=osp.join(log_dir, "finetune_s_model.pth"),
        max_sph_order=max_sph_order,
    ).to(device)

    cams.to(device)
    cams.eval()
    d_model.to(device)
    d_model.eval()
    s_model.to(device)
    s_model.eval()

    # feature_head
    feature_config = osp.join(log_dir, "feature_config.yaml")
    if osp.exists(feature_config):
        with open(feature_config, "r") as f:
            feature_config = OmegaConf.load(f)
            head_config = feature_config["Head"]

        feature_head = Feature_heads(head_config).to(device)

        feature_head_ckpt_path = osp.join(log_dir, "finetune_semantic_heads.pth")
        feature_head_state = torch.load(feature_head_ckpt_path, weights_only=True)
        feature_head.load_state_dict(feature_head_state)
        feature_head.eval()

    else:
        print(f"feature head config {feature_config} not found")
        feature_head = None

    return cfg, d_model, s_model, cams, feature_head


@torch.no_grad()
def viz_main(
    save_dir,
    log_dir,
    cfg_fn,
    N=1,
    H=480,
    W=854,
    move_angle_deg=4.5, #10.0,
    H_3d=960,
    W_3d=960,
    fov_3d=70,
    gptv_prompts=None,
    args=None,
    round_best=None,
):
    os.makedirs(save_dir, exist_ok = True)

    cfg, d_model, s_model, cams, feature_head = load_model_cfg(cfg_fn, log_dir)

    rel_focal_3d = 1.0 / np.tan(np.deg2rad(fov_3d) / 2.0)

    key_steps = [cams.T // 2, cams.T - 1, 0, cams.T // 4, 3 * cams.T // 4][:N]

    # * Get pose
    global_pose_list = get_global_3D_cam_T_cw(
        s_model,
        d_model,
        cams,
        H,
        W,
        cams.T // 2,
        back_ratio=0.5,
        up_ratio=0.2,
    )
    global_pose_list = global_pose_list[None].expand(cams.T, -1, -1)
    training_pose_list = [cams.T_cw(t) for t in range(cams.T)]

    # * #############################################################################
    # viz 3D
    save_fn_prefix = osp.join(save_dir, f"3D_moving")
    viz_single_2d_video_agent(
        H_3d,
        W_3d,
        cams,
        s_model,
        d_model,
        save_fn_prefix,
        global_pose_list,
        rel_focal=rel_focal_3d,
        gptv_prompts=gptv_prompts,
        args=args,
        round_best=round_best,
        feature_head=feature_head
    )

    # save_fn_prefix = osp.join(save_dir, f"3D_moving_node")
    # viz_single_2d_node_video(
    #     H_3d,
    #     W_3d,
    #     cams,
    #     s_model,
    #     d_model,
    #     save_fn_prefix,
    #     global_pose_list,
    #     rel_focal=rel_focal_3d,
    # )

    # save_fn_prefix = osp.join(save_dir, f"3D_moving_flow")
    # viz_single_2d_flow_video(
    #     H_3d,
    #     W_3d,
    #     cams,
    #     s_model,
    #     d_model,
    #     save_fn_prefix,
    #     global_pose_list,
    #     rel_focal=rel_focal_3d,
    # )

    # # flow
    # save_fn_prefix = osp.join(save_dir, f"training_moving_flow")
    # viz_single_2d_flow_video(
    #     H, W, cams, s_model, d_model, save_fn_prefix, training_pose_list
    # )
    # # node
    # save_fn_prefix = osp.join(save_dir, f"training_moving_node")
    # viz_single_2d_node_video(
    #     H, W, cams, s_model, d_model, save_fn_prefix, training_pose_list
    # )
    # # rgb
    # save_fn_prefix = osp.join(save_dir, f"training_moving")
    # viz_single_2d_video(
    #     H, W, cams, s_model, d_model, save_fn_prefix, training_pose_list, gptv_prompts
    # )

    # * #############################################################################
    # key_time_step = cams.T // 2
    for key_time_step in key_steps:
        fixed_pose_list = [cams.T_cw(key_time_step) for _ in range(cams.T)]
        round_pose_list = get_move_around_cam_T_cw(
            s_model,
            d_model,
            cams,
            H,
            W,
            np.deg2rad(move_angle_deg),
            total_steps=cams.T,  # cams.T
            center_id=key_time_step,
        )

    #     # viz flow
    #     save_fn_prefix = osp.join(save_dir, f"{key_time_step}_fixed_moving_flow")
    #     logging.info(f"Vizing fixed_moving_flow")
    #     viz_single_2d_flow_video(
    #         H, W, cams, s_model, d_model, save_fn_prefix, fixed_pose_list
    #     )
    #     logging.info(f"Vizing round_moving_flow")
    #     save_fn_prefix = osp.join(save_dir, f"{key_time_step}_round_moving_flow")
    #     viz_single_2d_flow_video(
    #         H, W, cams, s_model, d_model, save_fn_prefix, round_pose_list
    #     )
    #     # Viz node

    #     save_fn_prefix = osp.join(save_dir, f"{key_time_step}_round_moving_node")
    #     logging.info(f"Vizing round_moving_node")
    #     viz_single_2d_node_video(
    #         H, W, cams, s_model, d_model, save_fn_prefix, round_pose_list
    #     )
    #     save_fn_prefix = osp.join(save_dir, f"{key_time_step}_round_freezing_node")
    #     logging.info(f"Vizing round_freezing_node")
    #     viz_single_2d_node_video(
    #         H,
    #         W,
    #         cams,
    #         s_model,
    #         d_model,
    #         save_fn_prefix,
    #         round_pose_list,
    #         model_t=key_time_step,
    #     )
    #     save_fn_prefix = osp.join(save_dir, f"{key_time_step}_fixed_moving_node")
    #     logging.info(f"Vizing fixed_moving_node")
    #     viz_single_2d_node_video(
    #         H, W, cams, s_model, d_model, save_fn_prefix, fixed_pose_list
    #     )
        # Viz rgb
        save_fn_prefix = osp.join(save_dir, f"{key_time_step}_round_moving")
        logging.info(f"Vizing round_moving")
        viz_single_2d_video_agent(
            H, W, cams, s_model, d_model, save_fn_prefix, round_pose_list, args=args, gptv_prompts=gptv_prompts, round_best=round_best, feature_head=feature_head
        )
    #     save_fn_prefix = osp.join(save_dir, f"{key_time_step}_round_freezing")
    #     logging.info(f"Vizing round_freezing")
    #     viz_single_2d_video(
    #         H,
    #         W,
    #         cams,
    #         s_model,
    #         d_model,
    #         save_fn_prefix,
    #         round_pose_list,
    #         model_t=key_time_step,
    #     )
    #     save_fn_prefix = osp.join(save_dir, f"{key_time_step}_fixed_moving")
    #     logging.info(f"Vizing fixed_moving")
    #     viz_single_2d_video(
    #         H, W, cams, s_model, d_model, save_fn_prefix, fixed_pose_list
    #     )

    return


if __name__ == "__main__":

    # test_main(
    #     # saved_dir="./data/iphone_1x_dev/block/log/native_iphone_base.yaml20240507_155311",
    #     # data_root="./data/iphone_1x_dev/block/",
    #     saved_dir="./data/iphone_1x_dev/paper-windmill/log/native_iphone_base.yaml20240507_175337",
    #     data_root="./data/iphone_1x_dev/paper-windmill/",
    #     #
    #     cfg="./configs/iphone/iphone_base.yaml",
    #     device=torch.device("cuda"),
    #     tto_flag=True,
    # )

    import argparse

    args = argparse.ArgumentParser()
    args.add_argument("--config", "-c", type=str, required=True)
    args.add_argument("--root", "-r", type=str, required=True)
    # args.add_argument("--save", "-s", type=str, required=True)
    args.add_argument("--N", "-n", type=int, default=1)
    args.add_argument("--H", type=int, default=480)
    args.add_argument("--W", type=int, default=480)
    args.add_argument("--user_prompt", type=str, default="change the dog color to blue.")
    args.add_argument("--output_root", type=str, default="output")
    # "There are car, street and building objects and we want to change the car color to red."
    #"There are cow, grass, ground, bucket, and fence objects and we want to change the cow color to black and white dairy cow."
    #"There are train, track, road, tree, grass, and building objects and we want to delete the train"
    #"There are swan, water, bush, and wall objects and we want to change the color of swan to purple."
    #"There are dog, wall, window, and door objects and we want to change the dog color to darkest purple."
    #"There are dog, wall, window, and door objects and we want to change the dog color to look like Clifford."
    # Tunable parameters
    args.add_argument("--num_prompt", type=int, default=10)
    args.add_argument("--verbose", action="store_true")
    args.add_argument("--foldername", type=str, default="agentic_edit")
    # threshold scheduling for GPT prompt generation
    args.add_argument("--thr_min", type=float, default=0.85)
    args.add_argument("--thr_max", type=float, default=0.95)
    # temperature for LLM sampling used in selection and init
    args.add_argument("--gpt_temp", type=float, default=0.0)
    # API provider
    args.add_argument("--api", type=str, default="openrouter")
    args = args.parse_args()
    # propagate api provider to module global
    API_PROVIDER = getattr(args, "api", "openrouter") or "openrouter"
    os.environ["VIZ_AGENT_API"] = API_PROVIDER

    args.save = os.path.join(args.root,"editing_output")
    args.semantic_head_path = os.path.join(args.root,"finetune_semantic_heads.pth")


    os.makedirs(os.path.join(args.output_root, args.foldername), exist_ok=True)
    os.makedirs(os.path.join(args.output_root, args.foldername, 'iter'), exist_ok=True)
    os.makedirs(os.path.join(args.output_root, args.foldername, 'round1'), exist_ok=True)
    os.makedirs(os.path.join(args.output_root, args.foldername, 'iter_best'), exist_ok=True)
    os.makedirs(os.path.join(args.output_root, args.foldername, 'tmp'), exist_ok=True)

    # OPENROUTER_API_KEY can be provided via env; otherwise defaults are used in gptv_query

    # user_prompt, img_prompt = "change the cow color to purple", None
    user_prompt = args.user_prompt
    img_prompt = None
    prompt_list = user_prompt.split('<IMG>')
    user_prompt = user_prompt.split('<IMG>')[0] ## legacy, for naming use only
    idea_transcript = []
    for ii in range(len(prompt_list)):
        if ii == 0:
            idea_transcript.append("IDEA: %s."%prompt_list[0])
        elif ii%2==1:
            idea_transcript.append(load_img(prompt_list[ii]))
        elif ii%2==0:
            idea_transcript.append("%s"%prompt_list[ii])
    idea_transcript.append("End of IDEA.\n")

    idea_transcript = [idea_transcript[0] + idea_transcript[1]]
    print('idea transcirpt ===', idea_transcript)

    ### GPTV prompting iter
    current_prompts, prompt_history, select_history, image_history, reflection_history, bestidx_history = [],[],[],[],[],[]

    _sanitized_prompt = user_prompt.replace(' ','').replace('.','')
    os.makedirs(os.path.join(args.output_root, args.foldername, 'tmp', _sanitized_prompt), exist_ok=True)

    # Predefine for linters
    gptv_prompts = []
    rounds = 0

    for rounds in range(1):
       
        print('ROUND %d:\n'%rounds)
        ###### new rounds' prompt (init/revision)
        if rounds == 0:
            gptv_prompts = gptv_init_prompt(user_prompt, None, idea_transcript, args)
        current_prompts = gptv_prompts
        ###### t2i generation

        print('gptv_prompts', gptv_prompts)
        usable_n = min(args.num_prompt, len(gptv_prompts))
        for ii in range(usable_n):
            print('ii', ii)
            if ii >= len(gptv_prompts):
                break
            if gptv_prompts[ii] is None:
                continue
            if isinstance(gptv_prompts[ii], str) and len(gptv_prompts[ii].strip()) == 0:
                continue
            print(gptv_prompts[ii])

            args.save_name = os.path.join(args.output_root, args.foldername, 'tmp', _sanitized_prompt, f"{rounds}_{ii}_0.png")

            viz_main(
                args.save,
                args.root,
                args.config,
                N=1,
                H=480,
                W=480,
                gptv_prompts=gptv_prompts[ii],
                args=args
            )

    round_best, select_response = gptv_reflection_prompt_selectbest(
        user_prompt,
        img_prompt,
        idea_transcript,
        [os.path.join(args.output_root, args.foldername, 'tmp', _sanitized_prompt, f"{rounds}_{ii}_0.png") for ii in range(args.num_prompt)],
        args,
    )

    print('round best', round_best)

    viz_main(
        args.save,
        args.root,
        args.config,
        N=1,
        H=480,
        W=480,
        gptv_prompts=gptv_prompts[round_best] if isinstance(round_best, int) and 0 <= round_best < len(gptv_prompts) else (gptv_prompts[0] if len(gptv_prompts) > 0 else "no editing"),
        args=args,
        round_best=round_best,
    )
    # save_dir = "./debug/viz_debug7"
    # log_dir = "data/davis_dev/train/log/native_davis.2.yaml_dep=zoe_gt_cam=False_20240516_070323/"
    # cfg_fn = "./configs/wild/davis.2.yaml"
    # main(save_dir, log_dir, cfg_fn)