import gradio as gr
import os 
import sys
from pathlib import Path
import time
import accelerate



models = [
    "",
    "runwayml/stable-diffusion-v1-5",
    "CompVis/stable-diffusion-v1-4",
    "claudfuen/photorealistic-fuen-v1",
    "andite/anything-v4.0",
    "naclbit/trinart_stable_diffusion_v2",
    "nitrosocke/Arcane-Diffusion",
    "nitrosocke/archer-diffusion",
    "nitrosocke/elden-ring-diffusion",
    "nitrosocke/redshift-diffusion",
    "nitrosocke/spider-verse-diffusion", 
    "nitrosocke/mo-di-diffusion",
    "nitrosocke/classic-anim-diffusion",
    "dreamlike-art/dreamlike-photoreal-1.0",
    "dreamlike-art/dreamlike-photoreal-2.0",    
    "wavymulder/wavyfusion",
    "wavymulder/Analog-Diffusion",
    "prompthero/midjourney-v4-diffusion",
    "prompthero/openjourney",
    "dallinmackay/Van-Gogh-diffusion",
    "hakurei/waifu-diffusion",
    "DGSpitzer/Cyberpunk-Anime-Diffusion",
    "Fictiverse/Stable_Diffusion_BalloonArt_Model",
    "dallinmackay/Tron-Legacy-diffusion",
    "AstraliteHeart/pony-diffusion",
    "nousr/robo-diffusion",
    "Linaqruf/anything-v3",
    "Omnibus/maximum_diffusion_fast",
    "",
    
]
model_1=models[1]
model_2=models[2]
model_3=models[3]
model_4=models[4]

model_5=models[5]
model_6=models[6]
model_7=models[9]
model_8=models[13]

model_9=models[14]
model_10=models[15]
model_11=models[16]
model_12=models[17]

text_gen=gr.Interface.load("spaces/Omnibus/MagicPrompt-Stable-Diffusion_link",live=True, preprocess=True)

proc1=gr.Interface.load(f"models/{model_1}",live=False,preprocess=True, postprocess=False)
proc2=gr.Interface.load(f"models/{model_2}",live=False,preprocess=True, postprocess=False)
proc3=gr.Interface.load(f"models/{model_3}",live=False,preprocess=True, postprocess=False)
proc4=gr.Interface.load(f"models/{model_4}",live=False,preprocess=True, postprocess=False)

proc5=gr.Interface.load(f"models/{model_5}",live=False,preprocess=True, postprocess=False)
proc6=gr.Interface.load(f"models/{model_6}",live=False,preprocess=True, postprocess=False)
proc7=gr.Interface.load(f"models/{model_7}",live=False,preprocess=True, postprocess=False)
proc8=gr.Interface.load(f"models/{model_8}",live=False,preprocess=True, postprocess=False)

proc9=gr.Interface.load(f"models/{model_9}",live=False,preprocess=True, postprocess=False)
proc10=gr.Interface.load(f"models/{model_10}",live=False,preprocess=True, postprocess=False)
proc11=gr.Interface.load(f"models/{model_11}",live=False,preprocess=True, postprocess=False)
proc12=gr.Interface.load(f"models/{model_12}",live=False,preprocess=True, postprocess=False)


'''
proc1=gr.Interface.load("models/nitrosocke/Arcane-Diffusion", live=False, preprocess=True, postprocess=False)
proc2=gr.Interface.load("models/naclbit/trinart_stable_diffusion_v2", live=False, preprocess=True, postprocess=False)
proc3=gr.Interface.load("models/nitrosocke/redshift-diffusion", live=False, preprocess=True, postprocess=False)
proc4=gr.Interface.load("models/runwayml/stable-diffusion-v1-5", live=False, preprocess=True, postprocess=False)

proc5=gr.Interface.load("models/claudfuen/photorealistic-fuen-v1", live=False, preprocess=True, postprocess=False)
proc6=gr.Interface.load("models/CompVis/stable-diffusion-v1-4", live=False, preprocess=True, postprocess=False)
proc7=gr.Interface.load("models/Linaqruf/anything-v3.0", live=False, preprocess=True, postprocess=False)
proc8=gr.Interface.load("models/andite/anything-v4.0", live=False, preprocess=True, postprocess=False)

proc9=gr.Interface.load("models/dreamlike-art/dreamlike-photoreal-1.0", live=False, preprocess=True, postprocess=False)
proc10=gr.Interface.load("models/prompthero/openjourney", live=False, preprocess=True, postprocess=False)
proc11=gr.Interface.load("models/prompthero/midjourney-v4-diffusion", live=False, preprocess=True, postprocess=False)
proc12=gr.Interface.load("models/wavymulder/Analog-Diffusion", live=False, preprocess=True, postprocess=False)
'''

def get_prompts(prompt_text):
    return text_gen(prompt_text)
def send_it1(inputs,proc1=proc1):
        output1=proc1(inputs)
        return(output1)
def send_it2(inputs,proc2=proc2):
        output2=proc2(inputs)
        return(output2)
def send_it3(inputs,proc3=proc3):
        output3=proc3(inputs)
        return(output3)
def send_it4(inputs,proc4=proc4):
        output4=proc4(inputs)
        return(output4)

def send_it5(inputs,proc5=proc5):
        output5=proc5(inputs)
        return(output5)
def send_it6(inputs,proc6=proc6):
        output6=proc6(inputs)
        return(output6)
def send_it7(inputs,proc7=proc7):
        output7=proc7(inputs)
        return(output7)
def send_it8(inputs,proc8=proc8):
        output8=proc8(inputs)
        return(output8)

def send_it9(inputs,proc9=proc9):
        output9=proc9(inputs)
        return(output9)
def send_it10(inputs,proc10=proc10):
        output10=proc10(inputs)
        return(output10)
def send_it11(inputs,proc11=proc11):
        output11=proc11(inputs)
        return(output11)
def send_it12(inputs,proc12=proc12):
        output12=proc12(inputs)
        return(output12)

def main():
    with gr.Blocks(batch=False) as myface:
        with gr.Row():
            with gr.Tab("Title"):
                    gr.HTML("""<title>Maximum Diffusion</title><div style="text-align: center; max-width: 1500px; margin: 0 auto;">
                            <div style="display: inline-flex;align-items: center;gap: 0.8rem;font-size: 1.75rem;margin-bottom: 10px;line-height: 1em;">
                            <h1 style="font-weight: 900;font-size:30; margin-top: 9px;">Maximum Diffusion</h1></div>
                            <br><br><p style="margin-bottom: 1px;font-size: 15;font-weight: 100;line-height: 1.5em;">
                            Text to Image Model Comparison Space - CPU</p>
                            <br><p style="margin-bottom: 0px;font-size: 15;font-weight: 100;line-height: 1.5em;">
                            This space becomes quickly encumbered when even 1 model of 12 fails to load, or error out</p>
                            <p style="margin-bottom: 0px;font-size: 15;font-weight: 100;line-height: 1.5em;">
                            
                            Any tips on how to "time-out" Gradio.Interface models are welcome in the "Community" button above<br>
                            <br></div>""")
            with gr.Tab("Description"):
                gr.HTML("""<div style="text-align:center;">
                <h4>Enter your Prompt into the "Short Prompt" box and click "Magic Prompt" to load a prettified version of your prompt<br>
                When you are satisfied with the prompt that is in the "Text to Image" box, click "Launch" to load the Models.<br><br>
                Images load faster with a simpler prompt.<br>  
                Most images should load within 2 minutes.<br>
                Some models become stuck on certain prompts, and refreshing the page seems to fix it.<br><br>
                Not responsible for content, use at your own risk.
                </h4></div>""")
            with gr.Tab("DIY"):
                gr.HTML("""<div style="text-align:Left;">
                <h4>Copy/Paste this code in your app.py file<br><br>
                import gradio as gr<br>
                max_d=gr.Interface.load("spaces/Omnibus/maximum_diffusion")<br>
                max_d.launch()<br>                    
                </h4></div>""")                    
            with gr.Tab("Credits"):
                with gr.Row():
                    gr.Column()
                    with gr.Column(style="text-align:left;"):
                        gr.HTML("""
                                <div style="vertical-align:center">
                                  <br>
                                  <p>I learned everything I know from:
                                  <p><a href="https://huggingface.co/spaces/anzorq/finetuned_diffusion">Finetuned Diffusion</a></p>
                                  <p><a href="https://huggingface.co/spaces/Gustavosta/MagicPrompt-Stable-Diffusion">Magic Prompt Stable Diffusion</a></p>
                                  <p><a href="https://huggingface.co/spaces/huggingface-projects/magic-diffusion">Magic Diffusion</a></p>
                                  <p>Models by <a href="https://huggingface.co/Gustavosta">@Gustavosta</a>, <a href="https://twitter.com/haruu1367">@haruu1367</a>, <a href="https://twitter.com/DGSpitzer">@Helixngc7293</a>, <a href="https://twitter.com/dal_mack">@dal_mack</a>, <a href="https://twitter.com/prompthero">@prompthero</a> and others.</p>
                                </div>
                                """)
                    gr.Column()
            with gr.Tab("Tools"):
                        with gr.Tab("View"):
                          with gr.Row():  
                            with gr.Column(style="width=50%, height=70%"):
                                    gr.Pil(label="Crop")
                            with gr.Column(style="width=50%, height=70%"):
                                    gr.Pil(label="Crop")
                        with gr.Tab("Draw"):
                          with gr.Row():
                            with gr.Column(style="width=50%, height=70%"):
                                    gr.Pil(label="Crop")
                            with gr.Column(style="width=50%, height=70%"):
                                    gr.Pil(label="Draw")
                                    gr.ImagePaint(label="Draw")
                        with gr.Tab("Text"):
                            with gr.Row():
                                with gr.Column(scale=50):
                                    gr.Textbox(label="", lines=8, interactive=True)            
                                with gr.Column(scale=50):
                                    gr.Textbox(label="", lines=8, interactive=True)
                        with gr.Tab("Color Picker"):
                            with gr.Row():
                                with gr.Column(scale=50):
                                    gr.ColorPicker(label="Color", interactive=True)            
                                with gr.Column(scale=50):
                                    gr.ImagePaint(label="Draw", interactive=True)                                
        
        with gr.Row():
                input_text=gr.Textbox(label="Short Prompt")
                see_prompts=gr.Button("Magic Prompt")
        with gr.Row():
                prompt=gr.Textbox(label="Text to Image")
                run=gr.Button("Launch")
        with gr.Row():
                output1=gr.Image(label=(f"{model_1}"))
                output2=gr.Image(label=(f"{model_2}"))
                output3=gr.Image(label=(f"{model_3}"))
                output4=gr.Image(label=(f"{model_4}"))
        with gr.Row():
                output5=gr.Image(label=(f"{model_5}"))
                output6=gr.Image(label=(f"{model_6}"))
                output7=gr.Image(label=(f"{model_7}"))
                output8=gr.Image(label=(f"{model_8}"))
        with gr.Row():
                output9=gr.Image(label=(f"{model_9}"))
                output10=gr.Image(label=(f"{model_10}"))
                output11=gr.Image(label=(f"{model_11}"))
                output12=gr.Image(label=(f"{model_12}"))
        
        see_prompts.click(get_prompts, inputs=[input_text], outputs=[prompt])
    
        run.click(send_it1, inputs=[prompt], outputs=[output1])
        run.click(send_it2, inputs=[prompt], outputs=[output2])
        run.click(send_it3, inputs=[prompt], outputs=[output3])
        run.click(send_it4, inputs=[prompt], outputs=[output4])
        run.click(send_it5, inputs=[prompt], outputs=[output5])
        run.click(send_it6, inputs=[prompt], outputs=[output6])
        run.click(send_it7, inputs=[prompt], outputs=[output7])
        run.click(send_it8, inputs=[prompt], outputs=[output8])
        run.click(send_it9, inputs=[prompt], outputs=[output9])
        run.click(send_it10, inputs=[prompt], outputs=[output10])
        run.click(send_it11, inputs=[prompt], outputs=[output11])
        run.click(send_it12, inputs=[prompt], outputs=[output12])

    myface.queue(concurrency_count=600,status_update_rate=1)
    myface.launch(enable_queue=True,inline=True,max_threads=600)
if __name__ == "__main__":
    main()