import gradio as gr
from share_btn import community_icon_html, loading_icon_html, share_js
import random
import re

import torch
from transformers import AutoModelWithLMHead, AutoTokenizer, pipeline, set_seed
from optimum.intel.openvino import OVStableDiffusionPipeline


horoscope_model_id = "shahp7575/gpt2-horoscopes"
tokenizer = AutoTokenizer.from_pretrained(horoscope_model_id)
model = AutoModelWithLMHead.from_pretrained(horoscope_model_id)
text_generation_pipe = pipeline("text-generation", model="Gustavosta/MagicPrompt-Stable-Diffusion", tokenizer="gpt2")
stable_diffusion_pipe = OVStableDiffusionPipeline.from_pretrained("echarlaix/stable-diffusion-v1-5-openvino", revision="fp16", compile=False)
height = 128
width = 128
stable_diffusion_pipe.reshape(batch_size=1, height=height, width=width, num_images_per_prompt=1)
stable_diffusion_pipe.compile()

def fn(sign, cat):
    prompt = f"<|category|> {cat} <|horoscope|> {sign}"
    prompt_encoded = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)

    sample_outputs = model.generate(
        prompt_encoded,
        do_sample=True,
        top_k=40,
        max_length=300,
        top_p=0.95,
        temperature=0.95,
        num_beams=4,
        num_return_sequences=1,
    )

    final_out = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
    starting_text = " ".join(final_out.split(" ")[4:])
    seed = random.randint(100, 1000000)
    set_seed(seed)
    response = text_generation_pipe(starting_text + " " + sign + " art", max_length=(len(starting_text) + random.randint(60, 90)), num_return_sequences=1)
    image = stable_diffusion_pipe(response[0]["generated_text"], height=height, width=width, num_inference_steps=30).images[0]

    return [image, starting_text]


block = gr.Blocks(css="./css.css")

with block:
    with gr.Group():
        with gr.Box():
            with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
                text = gr.Dropdown(
                    label="Star Sign",
                    choices=["Aries", "Taurus","Gemini", "Cancer", "Leo", "Virgo", "Libra", "Scorpio", "Sagittarius", "Capricorn", "Aquarius", "Pisces"],
                    show_label=True,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    elem_id="prompt-text-input",
                ).style(
                    border=(True, False, True, True),
                    rounded=(True, False, False, True),
                    container=False,
                )

                text2 = gr.Dropdown(
                    choices=["Love", "Career", "Wellness"],
                    label="Category",
                    show_label=True,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    elem_id="prompt-text-input",
                ).style(
                    border=(True, True, True, True),
                    rounded=(True, False, False, True),
                    container=False,
                )

                btn = gr.Button("Generate image").style(
                    margin=False,
                    rounded=(False, True, True, False),
                    full_width=False,
                )

        gallery = gr.Image(
            interactive=False,
            label="Generated images", show_label=False, elem_id="gallery"
        ).style(grid=[2], height="auto")
        text = gr.Textbox("Text")

        with gr.Group(elem_id="container-advanced-btns"):
            with gr.Group(elem_id="share-btn-container"):
                community_icon = gr.HTML(community_icon_html)
                loading_icon = gr.HTML(loading_icon_html)
                share_button = gr.Button("Share to community", elem_id="share-btn")

        btn.click(fn=fn, inputs=[text, text2], outputs=[gallery, text])
        share_button.click(
            None,
            [],
            [],
            _js=share_js,
        )
        

block.queue(concurrency_count=40, max_size=20).launch(max_threads=150)