from transformers import pipeline, set_seed
import gradio as grad
import random
import re

# Model initialization
gpt2_pipe = pipeline('text-generation', model='Gustavosta/MagicPrompt-Stable-Diffusion', tokenizer='gpt2')
with open("ideas.txt", "r") as f:
    line = f.readlines()

def generate(starting_text):
    for count in range(4):
        seed = random.randint(100, 1000000)
        set_seed(seed)

        if starting_text == "":
            starting_text: str = line[random.randrange(0, len(line))].replace("\n", "").lower().capitalize()
            starting_text: str = re.sub(r"[,:\-–.!;?_]", '', starting_text)
            print(starting_text)
    
        response = gpt2_pipe(starting_text, max_length=random.randint(60, 90), num_return_sequences=4)
        response_list = []
        for x in response:
            resp = x['generated_text'].strip()
            if resp != starting_text and len(resp) > (len(starting_text) + 4) and resp.endswith((":", "-", "—")) is False:
                response_list.append(resp+'\n')
    
        response_end = "\n".join(response_list)
        response_end = re.sub('[^ ]+\.[^ ]+','', response_end)
        response_end = response_end.replace("<", "").replace(">", "")

        if response_end != "":
            return response_end
        if count == 4:
            return response_end

# Example generation
examples = []
for x in range(8):
    examples.append(line[random.randrange(0, len(line))].replace("\n", "").lower().capitalize())

# Custom CSS
custom_css = """
.gradio-container {
    background: linear-gradient(145deg, #2a2a2a, #3a3a3a);
    border-radius: 20px;
    box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.37);
    backdrop-filter: blur(4px);
    -webkit-backdrop-filter: blur(4px);
    border: 1px solid rgba(255, 255, 255, 0.18);
}

.input-textbox, .output-textbox {
    background: rgba(255, 255, 255, 0.1) !important;
    border-radius: 15px !important;
    padding: 20px !important;
    box-shadow: 0 4px 15px rgba(0,0,0,0.2);
    transform: perspective(1000px) rotateX(2deg);
    transition: all 0.3s ease;
}

.input-textbox:hover, .output-textbox:hover {
    transform: perspective(1000px) rotateX(0deg);
    box-shadow: 0 6px 20px rgba(0,0,0,0.3);
}

.example-btn {
    background: linear-gradient(45deg, #6b46c1, #805ad5) !important;
    border-radius: 10px !important;
    border: none !important;
    color: white !important;
    transform: translateY(0);
    transition: all 0.2s ease;
}

.example-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(107, 70, 193, 0.4);
}

.title-text {
    font-size: 2.5em !important;
    background: linear-gradient(45deg, #6b46c1, #805ad5);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
}

.container {
    animation: float 6s ease-in-out infinite;
}

@keyframes float {
    0% { transform: translateY(0px); }
    50% { transform: translateY(-20px); }
    100% { transform: translateY(0px); }
}

.submit-btn {
    background: linear-gradient(45deg, #6b46c1, #805ad5) !important;
    border-radius: 15px !important;
    padding: 10px 20px !important;
    font-weight: bold !important;
    transform: translateY(0);
    transition: all 0.3s ease;
}

.submit-btn:hover {
    transform: translateY(-3px);
    box-shadow: 0 10px 20px rgba(107, 70, 193, 0.3);
}
"""

# Interface content
title = """
<div class="title-text">
    Stable Diffusion Prompt Generator
</div>
"""

description = """
<div style="padding: 20px; background: rgba(255,255,255,0.1); border-radius: 15px; margin: 10px 0;">
    This is a demo of the model series: "MagicPrompt", aimed at Stable Diffusion. 
    To use it, simply submit your text or click on one of the examples.
    <br><br>
    <a href="https://huggingface.co/Gustavosta/MagicPrompt-Stable-Diffusion" 
       style="color: #805ad5; text-decoration: none; font-weight: bold;">
       Learn more about the model →
    </a>
</div>
"""

article = """
<div class="container" style="text-align: center; padding: 20px;">
    <img src='https://visitor-badge.glitch.me/badge?page_id=_Stable_Diffusion' alt='visitor badge'/>
    <div style="margin-top: 20px; color: #805ad5;">
        ✨ Created with AI Magic ✨
    </div>
</div>
"""

# Gradio Interface
interface = grad.Interface(
    fn=generate,
    inputs=[
        grad.Textbox(
            lines=1, 
            label="Initial Text", 
            placeholder="Enter your prompt here...", 
            elem_classes="input-textbox"
        )
    ],
    outputs=[
        grad.Textbox(
            lines=4, 
            label="Generated Prompts", 
            elem_classes="output-textbox"
        )
    ],
    examples=examples,
    title=title,
    description=description,
    article=article,
    flagging_mode="never",
    cache_examples=False,
    css=custom_css,
    theme=grad.themes.Soft()  # 기본 제공되는 Soft 테마 사용
)

# Launch the interface
interface.launch(debug=True)