import gradio as gr
#import peft
import transformers
import os
import re
import json


device = "cpu"
is_peft = False
model_id = os.environ.get("MODEL_ID") or "treadon/prompt-fungineer-355M"
auth_token = os.environ.get("HUB_TOKEN") or True

print(f"Using model {model_id}.")

if auth_token != True:
    print("Using auth token.")

model = transformers.AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True,use_auth_token=auth_token)
tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2")


def format_prompt(prompt, enhancers=True, inspiration=False, negative_prompt=False):
    try:
        pattern = r"(BRF:|POS:|ENH:|INS:|NEG:) (.*?)(?= (BRF:|POS:|ENH:|INS:|NEG:)|$)"
        matches = re.findall(pattern, prompt)
        vals = {key: value.strip() for key, value,ex in matches}
        result = vals["POS:"]
        if enhancers:
            result += " " + vals["ENH:"]
        if inspiration:
            result += " " + vals["INS:"]
        if negative_prompt:
            result += "\n\n--no " + vals["NEG:"]

        return result
    except Exception as e:
        return "Failed to generate prompt."

    
def generate_text(prompt, extra=False, top_k=100, top_p=0.95, temperature=0.85, enhancers = True, inpspiration = False , negative_prompt = False):
    
    if not prompt.startswith("BRF:"):
        prompt = "BRF: " + prompt

    if not extra:
        prompt = prompt + " POS:"

    model.eval()
    # SOFT SAMPLE
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    samples = []
    try:
        for i in range(1):
            outputs = model.generate(**inputs, max_length=256, do_sample=True, top_k=top_k, top_p=top_p, temperature=temperature, num_return_sequences=4, pad_token_id=tokenizer.eos_token_id)
            for output in outputs:
                sample = tokenizer.decode(output, skip_special_tokens=True)
                sample = format_prompt(sample, enhancers, inpspiration, negative_prompt)
                samples.append(sample)
    except Exception as e:
        print(e)

    return samples


with gr.Blocks() as fungineer:
    with gr.Row():
        gr.Markdown("""# Midjourney / Dalle 2 / Stable Diffusion Prompt Generator
This is the 355M parameter model.  There is also a 7B parameter model that is much better but far slower (access coming soon).
Just enter a basic prompt and the fungineering model will use its wildest imagination to expand the prompt in detail.""")
    with gr.Row():
        with gr.Column():
            base_prompt = gr.Textbox(lines=5, label="Base Prompt", placeholder="An astronaut in space", info="Enter a very simple prompt that will be fungineered into something exciting!")
            extra = gr.Checkbox(value=True, label="Extra Fungineer Imagination", info="If checked, the model will be allowed to go wild with its imagination.")
            with gr.Accordion("Advanced Generation Settings", open=False):
                top_k = gr.Slider( minimum=10, maximum=1000, value=100, label="Top K", info="Top K sampling")
                top_p = gr.Slider( minimum=0.1, maximum=1, value=0.95, step=0.01, label="Top P", info="Top P sampling")
                temperature = gr.Slider( minimum=0.1, maximum=1.2, value=0.85, step=0.01, label="Temperature", info="Temperature sampling.  Higher values will make the model more creative")

            with gr.Accordion("Advanced Output Settings", open=False):
                enh = gr.Checkbox(value=True, label="Enhancers", info="Add image meta information such as lens type, shuffter speed, camera model, etc.")
                insp = gr.Checkbox(value=False, label="Inpsiration", info="Include inspirational photographers that are known for this type of photography.  Sometimes random people will appear here, needs more training.")
                neg = gr.Checkbox(value=False, label="Negative Prompt", info="Include a negative prompt, more often used in Stable Diffusion.  If you're a Stable Diffusion user, chances are you already have a better negative prompt you like to use.")

        with gr.Column():
            outputs = [
                gr.Textbox(lines=2, label="Fungineered Text 1"),
                gr.Textbox(lines=2, label="Fungineered Text 2"),
                gr.Textbox(lines=2, label="Fungineered Text 3"),
                gr.Textbox(lines=2, label="Fungineered Text 4"),
            ]

    inputs = [base_prompt, extra, top_k, top_p, temperature, enh, insp, neg]


    submit = gr.Button(label="Fungineer",variant="primary")
    submit.click(generate_text, inputs=inputs, outputs=outputs)

    examples = []
    with open("examples.json") as f:
        examples = json.load(f)

    for i, example in enumerate(examples):
        with gr.Tab(f"Example {i+1}"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown(f"### Base Prompt")
                    gr.Image(value=f"{example['base']['src']}")
                    gr.Markdown(f"{example['base']['prompt']}")
                with gr.Column():
                    gr.Markdown(f"### 355M Prompt Fungineered")
                    gr.Image(value=f"{example['355M']['src']}")
                    gr.Markdown(f"{example['355M']['prompt']}")
                with gr.Column():
                    gr.Markdown(f"### 7B Prompt Fungineered")
                    gr.Markdown(f"Coming Soon!")


fungineer.launch(enable_queue=True)