import torch
from diffusers import DiffusionPipeline
import gradio as gr
import os
import spaces

model_list = os.environ.get("MODELS").split(",")
lora_list = os.environ.get("LORAS") # Not in use

@spaces.GPU
def generate(prompt, model):
    pipe = DiffusionPipeline.from_pretrained(
        model, 
        torch_dtype=torch.float16, 
        use_safetensors=True, 
    )
    pipe.to('cuda')
    
    negative_prompt = "nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"
    
    image = pipe(
        prompt, 
        negative_prompt=negative_prompt,
        width=832,
        height=1216, 
        guidance_scale=7,
        num_inference_steps=28
    ).images[0]
    
    return image


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            inp_prompt = gr.Textbox(label="Prompt")
            inp_model = gr.Dropdown(model_list, label="Select a model")            
            btn = gr.Button("Run")

        with gr.Column():
            out = gr.Image()

    btn.click(fn=generate, inputs=[inp_prompt, inp_model], outputs=out)

demo.launch()