apprentice.py

import os
import json
import itertools
from core.rag import RAG
from core.logger import Logger
from core.dreamer import Dreamer
from core.sampler import Sampler
from core.generator import Generator
from core.evaluator import Evaluator
from core.compiler import Compiler
from core.export import Exporter
from sceneprogexec import SceneProgExec
from concurrent.futures import ThreadPoolExecutor

class MaterialAgent:
    def __init__(self,
                 logger,
                 api_key,
                 repo_path = "dataset/yt-dataset-processed",
                 model_name="gpt-5",
                 reasoning_effort="medium",
                 num_procedures=5,
                 num_samples=10,
                 experiment_setup="default"):

        self.logger = logger
        self.api_key = api_key
        self.id = self.logger.id
        self.repo_path = repo_path
        self.model_name = model_name
        self.reasoning_effort = reasoning_effort
        self.num_procedures = num_procedures
        self.num_samples = num_samples
        self.exec = SceneProgExec()
        self.generator = Generator(model_name=self.model_name, reasoning_effort=self.reasoning_effort, experiment_setup=experiment_setup, api_key=self.api_key)
        self.rag = RAG(repo_path=self.repo_path, model_name=self.model_name, reasoning_effort=self.reasoning_effort, experiment_setup=experiment_setup, api_key=self.api_key)
        self.sampler = Sampler(id=self.id, model_name=self.model_name, reasoning_effort=self.reasoning_effort, api_key=self.api_key)
        self.compiler = Compiler(model_name=self.model_name, reasoning_effort=self.reasoning_effort, api_key=self.api_key)
        self.evaluator = Evaluator(id=self.id, model_name=self.model_name, reasoning_effort=self.reasoning_effort, api_key=self.api_key)
        self.dreamer = Dreamer(id=self.id, model_name=self.model_name, reasoning_effort=self.reasoning_effort, api_key=self.api_key)
        self.exporter = Exporter(id=self.id, model_name=self.model_name, reasoning_effort=self.reasoning_effort, api_key=self.api_key)


    def __call__(self, prompt, filepath=None):
        ## extract text and image parts from prompt
        text_prompt = prompt["text"]
        image_prompt = prompt["image"]
        initial_text = prompt["initial_text"]
        initial_code = prompt["initial_code"]
        self.logger.add_prompt(text_prompt, image_prompt, initial_text, initial_code)

        if initial_code:
            assert initial_text is not None, "Initial text must be provided along with initial code."
            print(f"******* Editing provided material: {initial_text} *******")
            if image_prompt is None:
                print("*******Dreaming target material image*******")
                image_prompt = self.dreamer(initial_code, initial_text, text_prompt)
            print("******Retrieving Video context for Initial Code*******")
            context = self.rag(initial_text, ignore_materials=[initial_text])
            _, context_editing, _ = context
            print("*******Reverse Compiling to get initial procedure*******")
            initial_procedure = self.compiler.reverse(initial_text, initial_code, context_editing)
        else:
            if image_prompt is None:
                print("*******Dreaming material image*******")
                image_prompt = self.dreamer.text2image(text_prompt)
            initial_procedure = None

        print("******Retrieving Video context*******")
        context = self.rag(text_prompt)
        self.logger.add_context(context)
        context, _, material_ref = context

        for i in range(5):
            ## generate a process
            print("*******Generating initial process*******")
            with ThreadPoolExecutor(max_workers=min(self.num_procedures, 20)) as executor:
                procedures = list(executor.map(
            lambda p: generate_procedure(self.generator, text_prompt, context, image_prompt, p),
            itertools.repeat(initial_procedure, self.num_procedures)
             ))
            print("*******Selecting best procedure*******")
            best_procedure = self.evaluator.select_procedure(self.logger, procedures, text_prompt, image_prompt)

            ## Turn it into code
            try:
                print("*******Converting procedure to code*******")
                program, doc = self.compiler(best_procedure)
                tmp = f"""
from blender.graph import MaterialGraph
{program}
graph.export('{self.id}')
    """
                self.exec(tmp)
                self.logger.add_best_initial_procedure(best_procedure)
                self.logger.add_initial_code_doc(program, doc, f"tmp/render{self.id}.png")
                break
            except Exception as e:
                ## try again if compilation or execution fails, up to 5 times
                print(f"Error during compilation or execution: {e}")
                pass

        print("*******Sampling programs from best procedure*******")
        sampled_programs = self.sampler(program, doc, text_prompt, image_prompt, n_samples=self.num_samples)
        sampled_programs.append(program)  # include the best procedure itself
        program = self.evaluator.select_program(self.logger, 1, sampled_programs, text_prompt, image_prompt)
        sampled_programs = self.sampler(program, doc, text_prompt, image_prompt, n_samples=self.num_samples)
        sampled_programs.append(program)  # include the best procedure itself
        program = self.evaluator.select_program(self.logger, 2, sampled_programs, text_prompt, image_prompt)
        sampled_programs = self.sampler(program, doc, text_prompt, image_prompt, n_samples=self.num_samples)
        sampled_programs.append(program)  # include the best procedure itself
        program = self.evaluator.select_program(self.logger, 3, sampled_programs, text_prompt, image_prompt)
        self.logger.save()
        export_code = self.exporter(
            best_procedure,
            program,
            filepath=f"tmp/{self.id}.blend" if filepath is None else filepath,
            doc=doc,
        )
        self.exec(export_code)
        # ------------------------------------------------------------
        # Save material references next to the output .blend
        # ------------------------------------------------------------
        if filepath is not None:
            output_dir = os.path.dirname(filepath)
        else:
            output_dir = "tmp"
        os.makedirs(output_dir, exist_ok=True)
        refs_path = os.path.join(output_dir, "references.json")
        with open(refs_path, "w") as f:
            json.dump(material_ref, f, indent=2)
        print(f"Saved material references to: {refs_path}")
        return best_procedure, program

def generate_procedure(generator, text_prompt, context, image_prompt=None, initial_procedure=None):
    procedure = generator(text_prompt, context, image_prompt=image_prompt, initial_procedure=initial_procedure)
    return procedure

# ============================================================
# CLI ENTRYPOINT
# ============================================================

import argparse


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-p",
        "--prompt",
        type=str,
        required=True,
        help="Material text prompt",
    )

    parser.add_argument(
        "-i",
        "--image",
        type=str,
        default=None,
        help="Optional reference image path",
    )

    parser.add_argument(
        "-o",
        "--output",
        type=str,
        required=True,
        help="Output .blend filepath",
    )

    parser.add_argument(
        "--model",
        type=str,
        default="gpt-5",
    )

    parser.add_argument(
        "-k",
        "--api_key",
        type=str,
        default=None,
        help="OpenAI API key (overrides the OPENAI_API_KEY environment variable)",
    )

    parser.add_argument(
        "--initial_text",
        type=str,
        default=None,
        help="Name/description of the source material to edit (use with --initial_code)",
    )

    parser.add_argument(
        "--initial_code",
        type=str,
        default=None,
        help="Path to an existing program.py to edit from",
    )

    args = parser.parse_args()

    # --------------------------------------------------------
    # Resolve API key from environment
    # --------------------------------------------------------

    api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
    if not api_key:
        raise EnvironmentError(
            "No OpenAI API key provided. Pass -k sk-... or "
            "export OPENAI_API_KEY=sk-... before running."
        )
    # Key routing:
    #   * The generation pipeline (generator/compiler/sampler/evaluator/dreamer/
    #     exporter/rag) uses `api_key` explicitly — on the server this is the
    #     VISITOR's key (passed via -k), so the visitor pays for the expensive work.
    #   * Deep, cheap node-level ops (blender/base.py embeddings + signature,
    #     graph.py agent) read OPENAI_API_KEY from the environment. We only set it
    #     as a fallback so we never clobber a server-provided key: on the server,
    #     OPENAI_API_KEY stays the operator's own key for those cheap ops; for the
    #     local CLI (no env key) it falls back to the key you passed.
    os.environ.setdefault("OPENAI_API_KEY", api_key)

    # --------------------------------------------------------
    # Create unique logger ID from output filename
    # --------------------------------------------------------

    output_name = os.path.splitext(
        os.path.basename(args.output)
    )[0]

    logger = Logger(desc=output_name)

    # --------------------------------------------------------
    # Construct agent
    # --------------------------------------------------------

    ma = MaterialAgent(
        logger,
        api_key=api_key,
        repo_path="dataset/yt-dataset-processed",
        model_name=args.model,
    )

    # --------------------------------------------------------
    # Build prompt object
    # --------------------------------------------------------

    initial_code = None
    if args.initial_code:
        with open(args.initial_code, "r") as f:
            initial_code = f.read()

    prompt = {
        "text": args.prompt,
        "image": args.image,
        "initial_text": args.initial_text,
        "initial_code": initial_code,
    }

    # --------------------------------------------------------
    # Run
    # --------------------------------------------------------

    procedure, program = ma(
        prompt,
        filepath=args.output,
    )

    print("Completed.")
    print("Output:", args.output)


if __name__ == "__main__":
    main()