import gradio as gr
import torch
from diffusers import StableDiffusionPipeline


#gr.Interface.load("models/bariscal/cbst_style")
#pipe = StableDiffusionPipeline.from_pretrained("bariscal/cbst_style", safety_checker=None) #, torch_dtype=torch.float16

# Create a PyTorch generator object
#generator = torch.Generator(device='cpu')


pipe = StableDiffusionPipeline.from_pretrained("bariscal/cbst_style", safety_checker=None)

def inference(prompt, negative_prompt, num_samples, height=512, width=512, num_inference_steps=50, guidance_scale=7.5):
    #with torch.inference_mode():
    return pipe(
        prompt, height=int(height), width=int(width),
        negative_prompt=negative_prompt,
        num_images_per_prompt=int(num_samples),
        num_inference_steps=int(num_inference_steps), guidance_scale=guidance_scale,
    ).images

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", value="portrait of human in cbst style")
            negative_prompt = gr.Textbox(label="Negative Prompt", value="")
            run = gr.Button(value="Generate")
        with gr.Column():
            num_samples = gr.Number(label="Number of Samples", value=2)
            guidance_scale = gr.Number(label="Guidance Scale", value=7.5)
    with gr.Row():
        height = gr.Number(label="Height", value=512, interactive = False)
        width = gr.Number(label="Width", value=512, interactive = False)
        num_inference_steps = gr.Slider(label="Steps", value=24)
    gallery = gr.Gallery()

    run.click(inference, inputs=[prompt, negative_prompt, num_samples, height, width, num_inference_steps, guidance_scale], outputs=gallery)

demo.launch(debug=True)