ChatGPT_Maya.py

"""
    Just a tool that simply executes the code to maya after being generated by openai chatGPT.

    Unfortunatly, complicated requests don't work,
    but simple automation tasks works, specially if you communicate in a logical way and break down steps.

    To test out and play with it:

    - Maya needs openai installed: - in autodesk/maya/bin directory: run "mayapy -m pip install openai" with terminal.
    - API key is needed, so either directly add your key into the API variable below (Not recommended,
        your API keys should be kept private.) or put it in maya env under a variable called "OPENAI_API_KEY".

    Have fun!
    www.LouisRossouw.com
"""

import os
import openai
import maya.cmds as cmds




class ChatChatAI:

    def __init__(self):

        # UI
        self.window_title = 'ChatChatGPTEES playplay'
        self.UI_width_height = 303, 65

        # Colors
        self.buttonS_BG_col = (0, 1, 0.7)
        self.window_BG_col = (0.12, 0.12, 0.12)

        # OpenAI settings.
        self.START_PROMPT = "write maya python code with no instructions that: "
        self.ENGINE = "text-davinci-003"
        self.MAX_TOKENS = 1024
        self.NUMBERS = 1
        self.STOP = None
        self.TEMP = 0.5

        API = os.getenv("OPENAI_API_KEY")
        openai.api_key = API



    def return_openai(self):
        """ Returns generated data from openAI GPT. """

        self.progress_bar(amount=5)

        prompt = cmds.textField(self.promptField, query=True, text=True)
        completions = openai.Completion.create(
                            engine=self.ENGINE,
                            prompt=self.START_PROMPT + prompt,
                            max_tokens=self.MAX_TOKENS,
                            n=self.NUMBERS,
                            stop=self.STOP,
                            temperature=self.TEMP
                            )

        self.progress_bar(amount=5)
        self.progress_bar(amount=-10)

        response = completions.choices[0].text
        print(response)

        return response



    def push_button_save(self, generated_code):
        """ Save out python code to a default location. """

        response = self.return_openai()
        scripts_path = os.getenv("MAYA_SCRIPT_PATH").split(";")[0]

        if os.path.exists(scripts_path) != True:
            os.mkdir(scripts_path)

        with open(scripts_path + '/maya_ai_gengen.py', 'w') as f:
            f.write(response)

        os.startfile(scripts_path)



    def push_button_show(self, *args):
        """ Shows the generated code. """

        response = self.return_openai()
        self.show_code_ui(response)



    def push_button_run(self, *args):
        """ Executes the returned generated code from GPT. """

        response = f"{self.return_openai()}"

        # Execute returned generated code (that is a string.)
        exec(response)



    def progress_bar(self, amount):
        """ Updates Progress bar. """

        cmds.progressBar(self.progressControl, edit=True, step=amount)



    def show_code_ui(self, response):
        """ Builds basic UI with maya cmds. """

        cmds.window(
                    title=self.window_title,
                    widthHeight=(self.UI_width_height[0], 250),
                    bgc=self.window_BG_col)

        cmds.columnLayout(adjustableColumn=True)
        cmds.rowColumnLayout(numberOfColumns=2, columnAttach=(1, 'right', 0), columnWidth=[(10, 100), (2, 250)])

        self.textField = cmds.text(response, width=250, height=250, al="left")

        cmds.setParent('..')
        cmds.showWindow()



    def main_ui(self):
        """ Builds main UI with maya cmds. """

        window = cmds.window(title=self.window_title,
                            widthHeight=(self.UI_width_height[0], self.UI_width_height[1]),
                            bgc=self.window_BG_col,
                            tlb=True)

        cmds.columnLayout(adjustableColumn=True)
        cmds.rowColumnLayout(numberOfColumns=1,
                             columnAttach=(1, 'left', 0),
                             columnWidth=[(10, 100), (2, 250)])

        self.promptField = cmds.textField(width=1050, ann="What you would like to do")
        self.progressControl = cmds.progressBar(maxValue=10, width=300)

        cmds.setParent('..')

        cmds.gridLayout(numberOfColumns=3,
                        cellWidthHeight=(self.UI_width_height[0] / 3, 20))

        cmds.button(label='Run',
                    bgc=self.buttonS_BG_col,
                    command=self.push_button_run)

        cmds.button(label='Show Script',
                    bgc=self.buttonS_BG_col,
                    command=self.push_button_show)

        cmds.button(label='Save Script',
                    bgc=self.buttonS_BG_col,
                    command=self.push_button_save)

        cmds.setParent('..')

        cmds.showWindow()



if __name__ == "__main__":

    ChatChatAI().main_ui()