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

import gradio as grad
from diffusers import StableDiffusionPipeline

tokenizer = AutoTokenizer.from_pretrained("shahp7575/gpt2-horoscopes")
model = AutoModelWithLMHead.from_pretrained("shahp7575/gpt2-horoscopes")

def fn(sign, cat):
    sign = "scorpio"

    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=4,
    )

    final_out = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
    starting_text = " ".join(final_out.split(" ")[4:])
    pipe = pipeline("text-generation", model="Gustavosta/MagicPrompt-Stable-Diffusion", tokenizer="gpt2")

    seed = random.randint(100, 1000000)
    set_seed(seed)

    response = pipe(starting_text, max_length=(len(starting_text) + random.randint(60, 90)), num_return_sequences=1)
    pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
    image = pipe(response[0]["generated_text"], num_inference_steps=5).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)