import gradio as gr
import torch
import numpy as np
from diffusers import DiffusionPipeline
from transformers import pipeline

pipe = pipeline('text-generation', model='daspartho/prompt-extend')

def extend_prompt(prompt):
    return pipe(prompt+',', num_return_sequences=1)[0]["generated_text"]

def text_it(inputs):
    return extend_prompt(inputs)


def load_pipeline(use_cuda):
    device = "cuda" if use_cuda and torch.cuda.is_available() else "cpu"

    if device == "cuda":
        torch.cuda.max_memory_allocated(device=device)
        torch.cuda.empty_cache()
        pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
        pipe.enable_xformers_memory_efficient_attention()
        pipe = pipe.to(device)
        torch.cuda.empty_cache()
    else:
        pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
        pipe = pipe.to(device)

    return pipe

def genie(prompt="sexy woman",  use_details=True,steps=2, seed=398231747038484200, use_cuda=False):
    pipe = load_pipeline(use_cuda)
    generator = np.random.seed(0) if seed == 0 else torch.manual_seed(seed)
    if use_details:
        extended_prompt = extend_prompt(prompt)
    else:
        extended_prompt=prompt
    int_image = pipe(prompt=extended_prompt, generator=generator, num_inference_steps=steps, guidance_scale=0.0).images[0]
    return int_image, extended_prompt


# Custom HTML for the interface
html_code = '''
    <style>
        body {
            background-color: #F0F0F0;
        }
    </style>
    <h1 style="color:black; text-align:center;">Stable Diffusion Turbo with GPT</h1>
'''

with gr.Blocks() as myface:
    gr.HTML(html_code)  # Add the custom HTML
    
    with gr.Row():
        input_text = gr.Textbox(label='Text prompt.', lines=1)
       
        with gr.Row():
            details_checkbox = gr.Checkbox(label="details", info="Generate Details?")
            steps_slider = gr.Slider(1, maximum=5, value=2, step=1, label='Number of Iterations')
            seed_slider = gr.Slider(minimum=0, step=1, maximum=999999999999999999, randomize=False, value=398231747038484200)
            cuda_checkbox = gr.Checkbox(label="cuda", info="Do you have cuda?")
        with gr.Row():
            generate_button = gr.Button("Generate")
    with gr.Row():
        output_image1 = gr.Image()
        output_image2 = gr.Image()
    with gr.Row():
        output_text1 = gr.Textbox(label="Generated Text", lines=2)
        output_text2 = gr.Textbox(label="Generated Text", lines=2)

    with gr.Row():
        output_image3 = gr.Image()
        output_image4 = gr.Image()
    with gr.Row():
        output_text3 = gr.Textbox(label="Generated Text", lines=2)
        output_text4 = gr.Textbox(label="Generated Text", lines=2)

    generate_button.click(genie, inputs=[input_text, details_checkbox, steps_slider, seed_slider, cuda_checkbox], outputs=[output_image1, output_text1], concurrency_limit=10)
    generate_button.click(genie, inputs=[input_text, details_checkbox, steps_slider, seed_slider, cuda_checkbox], outputs=[output_image2, output_text2], concurrency_limit=10)
    generate_button.click(genie, inputs=[input_text, details_checkbox, steps_slider, seed_slider, cuda_checkbox], outputs=[output_image3, output_text3], concurrency_limit=10)
    generate_button.click(genie, inputs=[input_text, details_checkbox, steps_slider, seed_slider, cuda_checkbox], outputs=[output_image4, output_text4], concurrency_limit=10)
myface.launch(inline=True, show_api=False, max_threads=200)