import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler


model_base = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16, use_safetensors=True, safety_checker=None,)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

lora_model_path = "./loralucy3/checkpoint-95000"
pipe.unet.load_attn_procs(lora_model_path)
pipe.to("cuda")


def greet(description,color,features,occasion,type_):
    
    prompt = 'white background '
    
    description = 'description:' + description.replace(' ', '-')

    color = ' color:' + ','.join(color)

    features = ' features:' + ','.join(features)

    occasion = ' occasion:' + ','.join(occasion)

    type_ = ' type:' + ','.join(type_)
   
    prompt += description + color + features +  occasion + type_

    print('prompt:',prompt)
    image = pipe(
        prompt,
        negative_prompt='deformed face,bad anatomy',
        width=312,
        height=512,
        num_inference_steps=100,
        guidance_scale=7.5,
        cross_attention_kwargs={"scale": 1.0}
        ).images[0]
    
    return image

iface = gr.Interface(fn=greet, 
                     inputs=[gr.Textbox(label='Description'),
                            gr.Dropdown(label='Color',choices=['Beige','Black','Blue','Brown','Green','Grey','Orange','Pink','Purple','Red','White','Yellow'],multiselect=True),
                            gr.Dropdown(label='Features',choices=['3/4-sleeve','Babydoll','Closed-Back','Corset','Crochet','Cutouts','Draped','Floral','Gloves','Halter','Lace','Long','Long-Sleeve','Midi','No-Slit','Off-The-Shoulder','One-Shoulder','Open-Back','Pockets','Print','Puff-Sleeve','Ruched','Satin','Sequins','Shimmer','Short','Short-Sleeve','Side-Slit','Square-Neck','Strapless','Sweetheart-Neck','Tight','V-Neck','Velvet','Wrap'],multiselect=True),
                            gr.Dropdown(label='Occasion',choices=['Homecoming','Casual','Wedding-Guest','Festival','Sorority','Day','Vacation','Summer','Pool-Party','Birthday','Date-Night','Party','Holiday','Winter-Formal','Valentines-Day','Prom','Graduation'],multiselect=True),
                            gr.Dropdown(label='Type',choices=['Mini-Dresses','Midi-Dresses','Maxi-Dresses','Two-Piece-Sets','Rompers','Jeans','Jumpsuits','Pants','Tops','Jumpers/Cardigans','Skirts','Shorts','Bodysuits','Swimwear'],multiselect=True),
                            ], 
                     outputs="image")
iface.launch()