system/refinement_process.py

import sys
import os
import subprocess
from pathlib import Path
import json
import importlib
from enum import Enum
from bson import ObjectId
from PIL import Image
import random
import shutil
from loguru import logger
import threading
import time
import io

from utils.image import plot_image_grid
from utils.code import get_code_as_string

from tasksolver.event import *
from tasksolver.common import  Question
from tasksolver.exceptions import  CodeExecutionException
from tasksolver.agent import Agent
from tqdm import tqdm

from agents import EditCodeAgent, GeneralAgent, Agent


class TaskSetting(Enum):
    LIGHTING = 1
    MATERIAL = 2
    SHAPEKEY = 3
    GEONODES = 4
    PLACEMENT = 5

TASKSETTING2PROMPTMODULE = {TaskSetting.GEONODES: "geonodes",
                            TaskSetting.LIGHTING: "lighting",
                            TaskSetting.MATERIAL: "material",
                            TaskSetting.SHAPEKEY: "shapekey",
                            TaskSetting.PLACEMENT: "placement"}

def tree_branch(branching_factor:int, depth_index:int, question_to_agent:Question, agent:Agent,
                script_save:Path, render_save:Path, thoughtprocess_save:Path,
                blender_file:str, blender_script:str,
                iteration:int, config:dict):
    '''
    For a given question, generate a list of runnable modifications by think and act
    '''

    query_and_act_semaphore = threading.Semaphore(
        min(config["run_config"]["max_concurrent_rendering_processes"], 
            config["run_config"]["max_concurrent_generator_requests"]) )
    
    results = [None] * branching_factor     # each slot is a position for a proposed modification
    def thread(question_to_agent, idx, results):
        # Fill one spot in the list results with a potential code change
        # The code change must be runnable, as checked by agent.act
        with query_and_act_semaphore: 
            #logger.info(f"thread {idx} acquired semaphore lock.")
            done = False
            num_tries = 0 
            max_tries = 1
            while not done and num_tries < max_tries:
                num_tries += 1
                try:
                    # Generate the code by think, the whole trunk
                    p_ans = agent.think(question_to_agent, num_tokens=3000, agent_idx=(idx, depth_index))
                    if len(p_ans.code) == 0:
                        logger.warning(f"The following response didn't parse into any code:\n{idx, script_save}")
                        pass
                except Exception as e: # TODO  ratelimitexception
                    # print(e)
                    # logger.warning(f"thread {idx} LLM querying failed with error:\n{str(e)}") 
                    continue
                try:
                    # Execute the code by act
                    code_path, render_path = agent.act(p_ans, 
                                                    script_save=script_save, 
                                                    render_save=render_save, 
                                                    iteration=iteration,
                                                    blender_file=blender_file,
                                                    blender_script=blender_script,
                                                    config=config,
                                                    blender_step=blender_step)
                    done = True
                except CodeExecutionException:
                    # blender execution failed, count failure.
                    # code_exec_exceptions += 1 
                    #logger.warning(f"thread {idx} code execution failed.")
                    pass
            if not done: 
                code_path = None
                render_path = None 
            results[idx] = (code_path, render_path, 'placeholder')  # The in-place modification of results with the 3-tuple
        # logger.info(f"thread {idx} released semaphore lock.")

    # FOR DEBUGGING
    # thread(question_to_agent, 0, results)    

    # Run the think-act pipeline for branching_factor(width) times
    llm_threads = [threading.Thread(target=thread,
                     args=(question_to_agent, i, results))
                     for i in range(branching_factor)]
    for idx, x in enumerate(llm_threads):
        x.start()
        #logger.info(f"starting thread {idx}...")

    for x in llm_threads:
        x.join() # wait till they all finish
    #logger.info(f"joined all threads")
    
    # assert all([el is not None for el in results])
    clean_results = []
    for result in results:
        if result is None:
            clean_results.append((None, None, 'placeholder'))
        else:
            clean_results.append(result)
    
    results = clean_results
    return results


def get_top_candidate(candidates, target, judge, depth_index, task_setting:TaskSetting, config:dict, 
                            target_description=None, use_vision=True,):

    prompting_submodule = importlib.import_module("prompting."+TASKSETTING2PROMPTMODULE[task_setting])
    craft_eval_question = getattr(prompting_submodule, "craft_eval_question")
    evaluation_semaphore = threading.Semaphore(config["run_config"]["max_concurrent_evaluation_requests"])

    def competition_thread(candidate1, candidate2, results, index, target_image=target):
        # Takes in two candidates and return left or right

        with evaluation_semaphore:
            # randomize the ordering
            done = False
            num_tries = 0
            max_tries = 3

            while not done and num_tries < max_tries:
                num_tries += 1
                order = random.sample([0,1], 2)

                left_code = get_code_as_string([candidate1[0], candidate2[0]][order[0]])
                left_img_file = [candidate1[1], candidate2[1]][order[0]]
                left_img = Image.open(left_img_file)

                right_code = get_code_as_string([candidate1[0], candidate2[0]][order[1]])
                right_img_file = [candidate1[1], candidate2[1]][order[1]]
                right_img = Image.open(right_img_file)

                assert left_img is not None
                assert right_img is not None
                assert left_code is not None
                assert right_code is not None
                if target_description is None:
                    print("target description is None -- intended?")

                question_to_critic = craft_eval_question(
                                        target_image=target_image,
                                        left_image=left_img,
                                        right_image=right_img,
                                        left_code=left_code,
                                        right_code=right_code,
                                        target_description=target_description,
                                        use_vision=use_vision)
                
                if target_description is None and target is None:
                    raise ValueError("No target provided to the competition_thread, either textual or image")

                try: 
                    p_ans = judge.think(question_to_critic, num_tokens=1000, agent_idx=index)
                    done = True
                except Exception as e: # TODO  ratelimitexception
                    # logger.warning(f"Sleep for 30s, {str(e)}")
                    continue

            if done:
                if p_ans.data == "left":
                    winner_index = order[0]
                elif p_ans.data == "right":
                    winner_index = order[1]
                
                # return a tuple of (winner, (left_img, right_img), raw answer(left or right), question )
                results[index] = ([candidate1, candidate2][winner_index], 
                                (left_img_file, right_img_file),
                                p_ans.raw, 
                                question_to_critic)   

    # assert len(candidates)%2 == 0, "Number candidates should be even, otherwise not handled."
    assert len(candidates) > 0, "the candidate list is empty"
     
    odd_one_out = None
    if len(candidates)%2 == 1: # number of candidates is odd
        odd_one_out = candidates[-1] 
        candidates = candidates[:-1]

    num_candidates = len(candidates)
    assert num_candidates % 2 == 0

    results = [None]*(num_candidates//2)

    max_tries = 1
    num_tries = 0
    done = False
    while not done and num_tries < max_tries:
        num_tries += 1 

        if num_tries > 1:
            #logger.warning(f"reattempt # {num_tries} for competition between candidates")
            pass

        # results: each entry is (winner, (left_img, right_img), raw answer(left or right), question) 
        competition_threads = [threading.Thread(target=competition_thread,
                args=(candidates[2*i], candidates[2*i+1], results, i, target))
                for i in range(num_candidates//2)]
        
        for x in competition_threads:
            x.start()
        for x in competition_threads:
            x.join()             

        # Pick the winners from the list results
        winners = [winner[0] for winner in results if winner is not None]
        if len(winners) == 0 and num_candidates//2 > 0:
            continue
        done = True

        intermediates = [{"left": winner[1][0],
                        "right": winner[1][1],
                        "winner": winner[0][1],
                        "inbound_question": str(winner[3]),
                        "thought_string": winner[2]} 
                        for winner in results if winner is not None]

    if not(len(winners) > 0 or num_candidates//2 == 0):
        # raise ValueError("All comparisons between samples seem to have failed.")
        print("All comparisons between samples seem to have failed.")
    
    if odd_one_out is not None:
        winners += [odd_one_out]
    
    try:
        intermediates
    except NameError:
        intermediates = []

    if len(winners) > 1:
        winner, _intermediates = get_top_candidate(winners, target, judge, depth_index, config=config,
                    target_description=target_description, task_setting=task_setting,
                    use_vision=use_vision) 
        return winner, intermediates + _intermediates
    else:
        if not winners:
            logger.info('Randomly sampled winner is generated.')
            winners.append(random.choice(candidates))
        return winners[0], intermediates # the only winner


def make_if_nonexistent(folder):
    if not os.path.exists(folder):
        os.makedirs(folder)
    return folder

# def blender_step(config, blender_file, blender_script, script_path, render_path, 
#                 verify_render_path=True):

#     '''
#     Generate a rendered image with given script_path at render_path
#     '''

#     if verify_render_path  and os.path.isfile(render_path):
#         raise ValueError(f"verify_render_path is True but {render_path} already exists before blender process.")

#     assert blender_file is not None and blender_script is not None
    
#     # Enter the blender code
#     command = [config["run_config"]["blender_command"], "--background", blender_file, 
#                     "--python", blender_script, 
#                     "--", script_path, render_path]
#     command = ' '.join(command)
#     command_run = subprocess.run(command, shell=True, check=True)
    
#     if verify_render_path  and not os.path.isfile(render_path):
#         logger.warning(f"The following bpy script didn't run correctly in blender:{script_path}")
#         raise CodeExecutionException 

#     return None


def merge_images_in_directory(directory, saved_to_local=True, merge_dir_into_image=True):
    '''
    Merge all images in the given directory into a single image.
    '''
    # Get a list of image paths
    image_paths = [os.path.join(directory, f) for f in os.listdir(directory) if f.endswith(('png', 'jpg', 'jpeg', 'webp'))]

    if not image_paths:
        new_image = None
        images=[]
        shutil.rmtree(directory)
        raise CodeExecutionException
    else:
        # Open images and get their sizes
        images = [Image.open(img) for img in image_paths]
        widths, heights = zip(*(i.size for i in images))

        # Calculate total size for the final image
        total_width = sum(widths)
        max_height = max(heights)

        # Create a new blank image with the calculated size
        new_image = Image.new('RGB', (total_width, max_height))

    # Paste all images into the new image
    x_offset = 0
    for img in images:
        new_image.paste(img, (x_offset, 0))
        x_offset += img.width

    if saved_to_local:
        # Save the final image to local
        if not merge_dir_into_image:    # Preserve the dir, adding new image to the dir
            merged_image_path = os.path.join(directory, 'merged_image.png')
        else:   # Delete the dir, and save the merged image as the name of the dir
            shutil.rmtree(directory)
            merged_image_path = directory

        if new_image:
            new_image.save(merged_image_path)
            print(f"Merged image saved to {merged_image_path}")
        return new_image, merged_image_path
    else:
        return new_image, None



def blender_step(infinigen_installation_path, blender_file_path, blender_render_script_path, script_path, render_dir, merge_all_renders=False, replace_if_overlap=True, merge_dir_into_image=False):

    '''
    Generate a rendered image with given script_path at render_dir.

    Inputs:
        blender_file_path: file path to the .blend base file
        blender_render_script_path: file path to the render script of blender scene
        script_path: file path to the script we want to render
        render_dir: dir path to save the rendered images
        merge_all_renders[optional]: True will merge all images in render_dir
        replace_if_overlap[optional]: False will skip if the render_dir exists and is non-empty, and True will proceed replace every overlapping render 
    '''

    def is_directory_empty(directory_path):
        # Check if the directory exists and is indeed a directory
        if not os.path.isdir(directory_path):
            raise ValueError(f"{directory_path} is not a valid directory path.")
    
        # List the contents of the directory
        return len(os.listdir(directory_path)) == 0

    assert blender_file_path is not None and blender_render_script_path is not None

    if replace_if_overlap:  # Just overwrite the files
        os.makedirs(render_dir, exist_ok=True)
    else:   
        if os.path.isdir(render_dir) and not is_directory_empty(render_dir): # If such dir already exists and is non-empty, skip
            return None 

        os.makedirs(render_dir, exist_ok=True)

    print('blender_render_script_path: ', blender_render_script_path)
    print('script_path: ', script_path)
    print('render_dir: ', render_dir)

    # Enter the blender code
    command = [infinigen_installation_path, "--background", blender_file_path, 
                    "--python", blender_render_script_path, 
                    "--", script_path, render_dir]
    command = ' '.join(command)
    command_run = subprocess.run(command, shell=True, check=True)

    # if is_directory_empty(render_dir):
    #     print(f"The following bpy script didn't run correctly in blender:{script_path}")
    #     return False
    #     # raise CodeExecutionException 
    # else:
    #     if merge_all_renders:
    merge_images_in_directory(render_dir, saved_to_local=True, merge_dir_into_image=merge_dir_into_image)
    return os.path.exists(render_dir)

    return True


def refinement(config, credentials, breadth, depth, blender_file, blender_script, 
                init_code, method_variation, output_folder, overwrite=True):        
    
    run_config = config["run_config"]

    # define task type and import relevant prompting modules 
    task_type = getattr(TaskSetting, config["task"]["type"].upper())
    prompting_submodule = importlib.import_module("prompting."+TASKSETTING2PROMPTMODULE[task_type])

    
    # import task definitions 
    code_editing_task = getattr(prompting_submodule, "code_editing_task")
    parameter_search_task = getattr(prompting_submodule, "parameter_search_task")
    pruning_task = getattr(prompting_submodule, "pruning_task")
    craft_tuner_question = getattr(prompting_submodule, "craft_tuner_question")
    craft_leap_question =  getattr(prompting_submodule, "craft_leap_question")


    # Read in one line of target description(instruction)
    target_description = None
    if config["input"]["text_prompt"] is not None:
        target_description = config["input"]["text_prompt"]

    assert method_variation in ("tune_leap", "leap", "tune")

    # Create all the dirs for results
    script_save = Path(output_folder)/Path("scripts/")
    render_save = Path(output_folder)/Path("renders/")
    thoughtprocess_save = Path(output_folder)/Path("thought_process/")
    make_if_nonexistent(output_folder)
    make_if_nonexistent(script_save)
    make_if_nonexistent(render_save)
    make_if_nonexistent(thoughtprocess_save)
            
    init_render_file = os.path.join(output_folder, "init_render.png")      # The original blender rendered image
    target_code = config["input"]["target_code"]       # Target bpy code
    target_render_file = config["input"]["input_image"]     # Dalle generated pseudo-target image based on the text file
    
    assert init_code is not None
    if not os.path.exists(init_render_file):
        # blender_step(config, blender_file, blender_script, init_code, init_render_file, verify_render_path=True)        
        blender_step(config["run_config"]["blender_command"], blender_file, blender_script, init_code, init_render_file, merge_all_renders=True, merge_dir_into_image=True)

    init_image = Image.open(init_render_file)      # Keep an record of original image


    if target_render_file is not None:      # If provided with a path to ideal target image
        if target_code is not None and not os.path.exists(target_render_file):  # If target_code is also provided and no image provided
            # blender_step(config, blender_file, blender_script, target_code, target_render_file,  verify_render_path=True)  # Render and overwrite the dalle generated images
            blender_step(config["run_config"]["blender_command"], blender_file, blender_script, target_code, target_render_file, merge_all_renders=True, merge_dir_into_image=True)

        target_image = Image.open(target_render_file)      
    else:
        target_image = None

    # create agents 
    thinker_is_visual = None
    evaluator_is_visual = None
    
    # Initialize the thinker (code editor and parameter searcher)
    thinker_class_type = GeneralAgent if config["run_config"]["edit_style"] == "rewrite_code" else EditCodeAgent 

    # if run_config["edit_generator_type"] in ("gpt-4o", "o1", "o1-mini", "gpt-4-turbo", 'gpt-4o-mini', 'o3-mini', 'gpt-4-turbo'): 
    #     # Default!
    #     param_tuner = thinker_class_type (credentials["openai"], parameter_search_task, vision_model=run_config["edit_generator_type"])
    #     agent = thinker_class_type (credentials["openai"], code_editing_task, vision_model=run_config["edit_generator_type"])
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] in ("claude-3-5-sonnet-latest", "claude-3-5-haiku-latest", "claude-3-opus-latest"):
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (credentials["claude"], parameter_search_task, vision_model=run_config["edit_generator_type"])
    #     agent = thinker_class_type (credentials["claude"], code_editing_task, run_config["edit_generator_type"])
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] in ('gemini-pro', 'gemini-pro-vision', 'gemini-2.0-flash', 'gemini-1.5-flash', 'gemini-1.5-pro'):
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (credentials["gemini"], parameter_search_task, vision_model=run_config["edit_generator_type"])
    #     agent = thinker_class_type (credentials["gemini"], code_editing_task, run_config["edit_generator_type"])
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] == "Intern":
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (None, parameter_search_task, vision_model='OpenGVLab/InternVL2-8B')
    #     agent = param_tuner
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] == "InternLlama":
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (None, parameter_search_task, vision_model=run_config["edit_generator_type"])
    #     agent = param_tuner
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] == "Qwen":
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (None, parameter_search_task, vision_model='Qwen/Qwen2-VL-7B-Instruct-AWQ')
    #     agent = param_tuner
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] == "QwenLlama":
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (None, parameter_search_task, vision_model=run_config["edit_generator_type"])
    #     agent = param_tuner
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] == "MiniCPM":
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (None, parameter_search_task, vision_model='openbmb/MiniCPM-V-2_6-int4')
    #     agent = param_tuner
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] == "MiniCPMLlama":
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (None, parameter_search_task, vision_model=run_config["edit_generator_type"])
    #     agent = param_tuner
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] == "Phi":
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (None, parameter_search_task, vision_model='microsoft/Phi-3.5-vision-instruct')
    #     agent = param_tuner
    #     thinker_is_visual = True

    # elif run_config["edit_generator_type"] == "PhiLlama":
    #     # tasks stay the same.
    #     param_tuner = thinker_class_type (None, parameter_search_task, vision_model=run_config["edit_generator_type"])
    #     agent = param_tuner
    #     thinker_is_visual = True
    # else:
    #     param_tuner = thinker_class_type (None, parameter_search_task, vision_model=run_config["edit_generator_type"])

    #     # TODO: If your model requires API, you can add credentials["your_model"]
    #     # param_tuner = thinker_class_type (credentials["your_model"], parameter_search_task, vision_model=run_config["edit_generator_type"])
    #     agent = param_tuner
    #     thinker_is_visual = True

    param_tuner = thinker_class_type(credentials, parameter_search_task, vision_model=run_config["edit_generator_type"])
    print('Parameter tuner initialized')
    agent = param_tuner
    thinker_is_visual = True

    # # Initialize the state evaluator (pruning task)
    # if run_config["state_evaluator_type"] in ("gpt-4o", "o1", "o1-mini", "gpt-4-turbo", "gpt-4o-mini", 'o3-mini', 'gpt-4-turbo'): 
    #     judge = GeneralAgent(credentials, pruning_task, vision_model=run_config["state_evaluator_type"])
    #     evaluator_is_visual = True

    # elif run_config["state_evaluator_type"] in ("claude-3-5-sonnet-latest", "claude-3-5-haiku-latest", "claude-3-opus-latest", "claude-3-5-sonnet-20240620"):
    #     judge = GeneralAgent(credentials, pruning_task, vision_model=run_config["state_evaluator_type"])
    #     evaluator_is_visual = True

    # elif run_config["state_evaluator_type"] in ('gemini-pro', 'gemini-pro-vision', 'gemini-2.0-flash', 'gemini-1.5-flash', 'gemini-1.5-pro'):
    #     judge = GeneralAgent(credentials, pruning_task, vision_model=run_config["state_evaluator_type"])
    #     evaluator_is_visual = True

    # elif run_config["state_evaluator_type"] in ("Intern", "InternLlama"):
    #     judge = GeneralAgent(credentials, pruning_task, vision_model=run_config["state_evaluator_type"])
    #     evaluator_is_visual = True

    # elif run_config["state_evaluator_type"] in ("MiniCPM", "MiniCPMLlama"):
    #     judge = GeneralAgent(None, pruning_task, vision_model='openbmb/MiniCPM-V-2_6-int4')
    #     evaluator_is_visual = True

    # elif run_config["state_evaluator_type"] in ("Qwen", "QwenLlama"):
    #     judge = GeneralAgent(None, pruning_task, vision_model='Qwen/Qwen2-VL-7B-Instruct-AWQ')
    #     evaluator_is_visual = True

    # elif run_config["state_evaluator_type"] in ("Phi", "PhiLlama"):
    #     judge = GeneralAgent(None, pruning_task, vision_model='microsoft/Phi-3.5-vision-instruct')
    #     evaluator_is_visual = True

    # else:
    #     judge = GeneralAgent(None, pruning_task, vision_model=run_config["state_evaluator_type"])
    #     evaluator_is_visual = True

    judge = GeneralAgent(credentials, pruning_task, vision_model=run_config["state_evaluator_type"])
    evaluator_is_visual = True

    # raise ValueError(f"Invalid evaluator: {run_config['state_evaluator_type']}")

    assert thinker_is_visual is not None and evaluator_is_visual is not None, "Remember to assign these variables above."
    print(evaluator_is_visual, thinker_is_visual)
    if not evaluator_is_visual or not thinker_is_visual:
        assert target_description is not None, f"When either the {run_config['edit_generator_type']} or {run_config['state_evaluator_type']} is not visual, make sure that you have the target description."

    

    # start of simulation
    code_path = init_code       # original starter code
    render_path = init_render_file      # path of original rendered image

    intermediary_outputs = []

    for i in tqdm(range(depth)):       # Tree depth
        if not overwrite:
            if os.path.exists(thoughtprocess_save/f"iteration_{i}.json"):
                with open(thoughtprocess_save/f"iteration_{i}.json", "r") as f: 
                    process_json = json.load(f)

                # If there is a winner among the last layer of branches,
                # keep them since we will compare that to the output of this layer
                if "winner_code" in process_json[-1] and "winner_image" in process_json[-1]:
                    code_path = process_json[-1]["winner_code"]
                    render_path = process_json[-1]["winner_image"]
                    intermediary_outputs.append({
                                    'code_path': code_path, 
                                    'render_path': render_path,
                                    "iteration": i})
                    logger.info(f"Thought process loaded for iteration {i}")
                    continue 
                else:
                    pass # start at this iteration! 

        process_json = []
        
        if (method_variation in ('tune_leap',) and i%2 == 0) or method_variation in ('tune',):
        # The case of tune

            # Craft a question based on the image and text input
            tuner_question = craft_tuner_question(
                blender_init_code_str=get_code_as_string(code_path),
                init_image=Image.open(render_path), 
                target_image=target_image,
                target_description=target_description,
                use_vision=thinker_is_visual)
            
            logger.info(f"tuner_question_formed")

            # Run think-act on the agent by `breadth` times
            # results is a list of length `breadth`, one runnable modification on each entry
            # Each entry is (code_path, render_path, p_ans.raw), code_path is the path to the modified bpy script, 
            # render_path the resulting rendered image, and parsed answer
            results = tree_branch(breadth, i, tuner_question, 
                                        agent=param_tuner,
                                        script_save=script_save,
                                        render_save=render_save,
                                        thoughtprocess_save=process_json,
                                        blender_file=blender_file, 
                                        blender_script=blender_script,
                                        iteration=i, config=config)
            
            logger.info(f"Runnable modifications generated for iteration {i}/{depth-1}(0-indexed) of depth")

            results = [el for el in results if el[0] is not None]       # Take out the code_path

            # Register all the potential modifications to the json file
            process_json.append(
                {
                    "phase": "explode_options_TUNE",
                    "iteration": i,
                    "inbound_question": str(tuner_question),
                    "choices_image": [res[1] for res in results],
                    "choices_code": [res[0] for res in results],
                    "thought_strings": [res[2] for res in results]
                }   
            )

            results.append((code_path, render_path, 'placeholder'))

            # Get the top candidate by state evaluator
            if len(results) > 1:

                # top_candidate is a (code, image) pair
                top_candidate, intermediates = get_top_candidate(results,
                                target_image, judge, i, config=config, 
                                target_description=target_description, 
                                task_setting=task_type,
                                use_vision=evaluator_is_visual)
                process_json.append(
                    {
                        "phase": "selection",
                        "choices_image": [res[1] for res in results],
                        "choices_code": [res[0] for res in results],
                        "winner_image": top_candidate[1],
                        "winner_code": top_candidate[0],
                        "decision_process": intermediates
                    }
                )
            else:
                top_candidate = (code_path, render_path)
                process_json.append(
                    {
                        "phase": "selection",
                        "choices_image": [res[1] for res in results],
                        "choices_code": [res[0] for res in results],
                        "winner_image": top_candidate[1],
                        "winner_code": top_candidate[0],
                        "decision_process": None
                    }
                )
            
        elif (method_variation in ('tune_leap',) and i%2 == 1) or method_variation in ('leap',):
            # Similar to tune mode, but takes bigger step for the agent to think

            # we leap
            #  Craft a question based on the image and text input
            print(f'Here 689: {code_path}')
            question_to_agent = craft_leap_question(
                blender_init_code_str = get_code_as_string(code_path),
                init_image = Image.open(render_path),
                target_image=target_image,
                target_description=target_description,
                use_vision=thinker_is_visual) 

            results = tree_branch(breadth, i, question_to_agent, 
                                        agent=agent,
                                        script_save=script_save,
                                        render_save=render_save,
                                        thoughtprocess_save=process_json,
                                        blender_file=blender_file, 
                                        blender_script=blender_script,
                                        iteration=i, config=config)
            logger.info(f"Runnable modifications generated for iteration {i}/{depth} of depth")

            results = [el for el in results if el[0] is not None]
            process_json.append(
                {
                    "phase": "explode_options_LEAP",
                    "iteration": i,
                    "inbound_question": str(question_to_agent),
                    "choices_image": [res[1] for res in results],
                    "choices_code": [res[0] for res in results],
                    "thought_strings": [res[2] for res in results]
                }   
            )

            results.append((code_path, render_path, 'placeholder'))

            if len(results) > 1:
                top_candidate, intermediates = get_top_candidate(results, 
                                target_image, judge, i, config=config, 
                                target_description=target_description,
                                task_setting=task_type,
                                use_vision=evaluator_is_visual)
                process_json.append(
                    {
                        "phase": "selection",
                        "choices_image": [res[1] for res in results],
                        "choices_code": [res[0] for res in results],
                        "winner_image": top_candidate[1],
                        "winner_code": top_candidate[0],
                        "iteration": i,
                        "decision_process": intermediates
                    }
                )
            else:
                top_candidate = (code_path, render_path)
                process_json.append(
                    {
                        "phase": "selection",
                        "choices_image": [res[1] for res in results],
                        "choices_code": [res[0] for res in results],
                        "winner_image": top_candidate[1],
                        "winner_code": top_candidate[0],
                        "decision_process": None
                    }
                )
                
        # if config["run_config"]["enable_hypothesis_reversion"]:          # Compares the current one with the best of last layer
        #     if len(results) > 0:
        #         top_candidate_before = top_candidate
        #         top_candidate, intermediates = get_top_candidate([top_candidate, (code_path, render_path)], 
        #                         target_image, judge, config=config, target_description=target_description,
        #                         task_setting=task_type,
        #                         use_vision=evaluator_is_visual)
        #         process_json.append(
        #             {
        #                 "phase": "selection",
        #                 "choices_image": [top_candidate_before[1], render_path],
        #                 "choices_code": [top_candidate_before[0], code_path],
        #                 "winner_image": top_candidate[1],
        #                 "winner_code": top_candidate[0],
        #                 "iteration": i,
        #                 "decision_process": intermediates
        #             }
        #         ) 

        #         # If last layer's output is better, revert.
        #         if top_candidate != top_candidate_before:
        #             #logger.info("Check shows that old sample is better. Rebasing to old sample.")
        #             pass
        
        logger.info(f"Top candidated picked for iteration {i}/{depth-1}(0-indexed) of depth. Code:{top_candidate[0]}, image:{top_candidate[1]}")

        with open(thoughtprocess_save/f"iteration_{i}.json", "w") as f:
            json.dump(process_json, f, indent=4)
            logger.info(f"Thought process saved for iteration {i}")

        # set new, best so far
        code_path, render_path = top_candidate[0], top_candidate[1] 
        # add to final output
        intermediary_outputs.append({'code_path': code_path, 
                                'render_path': render_path,
                                "iteration": i})


    fig = plot_image_grid([(Image.open(el["render_path"]) if el is not None else None) 
                for el in intermediary_outputs], 
                rows=1, cols=len(intermediary_outputs))
    fig.savefig(str(output_folder/"best_of.png"))