import torch
import spaces
import gradio as gr

from diffusers import StableDiffusionPipeline


repo = "Tramac/style-portrait-v1-5"
pipeline = StableDiffusionPipeline.from_pretrained(
    repo,
    torch_dtype=torch.float16,
)
# pipeline.enable_model_cpu_offload()

@spaces.GPU(enable_queue=True)
def generate_image(prompt, neg_prompt):
    results = pipeline(
        prompt=prompt,
        negative_prompt=neg_prompt,
        height=832,
        width=512,
        num_inference_steps=20,
    )
    return results.images[0]



with gr.Blocks() as demo:
    gr.HTML("<h1><center>Style Portrait</center></h1>")
    gr.HTML("<p><center>text-to-image generation</center></p><p><center><a href='https://huggingface.co/Tramac/style-portrait-v1-5'></a></center></p>")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label='Enter your prompt', scale=8)
            neg_prompt = gr.Textbox(label='Enter your negative prompt', scale=8)
            submit = gr.Button(scale=1, variant='primary')

        with gr.Column():
            img = gr.Image(label='Style-Portrait Generated Image')

    prompt.submit(
        fn=generate_image,
        inputs=[prompt, neg_prompt],
        outputs=img,
    )
    submit.click(
        fn=generate_image,
        inputs=[prompt, neg_prompt],
        outputs=img,
    )

demo.queue().launch()