import torch
import spaces
import gradio as gr

from diffusers import StableDiffusionPipeline


repo = "Tramac/style-portrait-v1-5"
# Ensure model and scheduler are initialized in GPU-enabled function
if torch.cuda.is_available():
    pipeline = StableDiffusionPipeline.from_pretrained(
        repo,
        torch_dtype=torch.float16,
        variant="fp16"
    ).to("cuda")

@spaces.GPU(enable_queue=True)
def generate_image(prompt, neg_prompt):
    results = pipe(
        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(share=True)