import gradio

class Model:
    def __init__(self, name, path="", prefix=""):
        self.name = name
        self.path = path
        self.prefix = prefix

models = [
   Model("Marvel","models/ItsJayQz/Marvel_WhatIf_Diffusion", "whatif style"), 
   Model("Portrait plus", "models/wavymulder/portraitplus", "portrait+ style"),
   Model("CF25", "models/gsdf/Counterfeit-V2.5", "anime style"),
   Model("vintedois", "models/22h/vintedois-diffusion-v0-1", "vintedois style"),
   Model("dreamlike", "models/dreamlike-art/dreamlike-diffusion-1.0","dreamlike style"),
   Model("GTA5","models/ItsJayQz/GTA5_Artwork_Diffusion", "GTA5 style")
]

model1=[]
model2=[]
model3=[]

for i in range(len(models)):
    model3.append(models[i].name)
    model1.append(gradio.Interface.load(models[i].path))
    model2.append(models[i].prefix)

def process1(prompt, modelSelected):
    if (modelSelected==''):
        modelSelected = "Marvel"
    model_idx=model3.index(modelSelected)
    prompt+=", in "+model2[model_idx]
    image_return = model1[model_idx](prompt)
    return image_return

sandbox = gradio.Interface(fn=process1, 
                        inputs=[gradio.Textbox(label="Enter Prompt:"),  gradio.Dropdown(model3)],
                        outputs=[gradio.Image(label="Produced Image")], 
                        title='Text to Image', 
                        examples=[
                        ["Viggo Mortensen Gryffindor wizard, Hogwart University, castle tower background", "Portrait plus"],
                        ["Male pirate, eye patch, black hat, sailing boat, ocean in background", "CF25"],
                        ["Portrait close up, Elvis Presley, concert hall in the background", "GTA5"],
                        ["Marvel Blackwidow portrait close up. building city background", "Marvel"],
                        ["Benedict Cumberbatch wizard of black magic, Hogwart University, castle tower background", "vintedois"],
                        ["A white rabbit wizard, Hogwart University, Castle in the background", "dreamlike"]
                        ])
                                  
sandbox.queue(concurrency_count=20).launch()