MagicPrompt script for AUTOMATIC1111's Stable Diffusion GUI v1.2

import math
import os
import sys
import traceback

import modules.scripts as scripts
import gradio as gr

from modules.processing import Processed, process_images
from PIL import Image
from modules.shared import opts, cmd_opts, state

from aitextgen import aitextgen


class Script(scripts.Script):
    def title(self):
        return "Prompts from Gustavosta MagicPrompt"

    def ui(self, is_img2img):
        gpt2model = gr.Textbox(label='GPT-2 Model Directory', value="MagicPrompt-Stable-Diffusion")
        prompt_length = gr.Slider(label='Max Length', value=100, minimum=1, maximum=150, step=1)
        temperature_value = gr.Slider(label='Temperature', value=0.7, minimum=0.1, maximum=2, step=0.05)
        return [gpt2model, prompt_length, temperature_value]

    def run(self, p, gpt2model, prompt_length, temperature_value):
        ai = aitextgen(model_folder=gpt2model, tokenizer_file=gpt2model + "/tokenizer.json")
        p.prompt = ai.generate_one(prompt=p.prompt, max_length=prompt_length, temperature=temperature_value)
        p.prompt = p.prompt.translate({ord(i): None for i in '[:](){}\/'})
        print(f"Generated prompt: {p.prompt}")
        p.do_not_save_grid = True
        proc = process_images(p)
        return proc