import gradio as gr
import numpy as np
import psutil

def create_map():
    return np.zeros(shape=(512, 512), dtype=np.uint8)+255

def get_system_memory():
    memory = psutil.virtual_memory()
    memory_percent = memory.percent
    memory_used = memory.used / (1024.0 ** 3)
    memory_total = memory.total / (1024.0 ** 3)
    return {"percent": f"{memory_percent}%", "used": f"{memory_used:.3f}GB", "total": f"{memory_total:.3f}GB"}



def create_demo_keypose(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Keypose)')
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(source='upload', type="numpy")
                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                with gr.Row():
                    type_in = gr.inputs.Radio(['Keypose', 'Image'], type="value", default='Image', label='Input Types\n (You can input an image or a keypose map)')
                    fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed to produce a fixed output)')
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the keypose to the result)", minimum=0, maximum=1, value=1, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
        ips = [input_img, type_in, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo

def create_demo_openpose(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Openpose)')
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(source='upload', type="numpy")
                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                with gr.Row():
                    type_in = gr.inputs.Radio(['Openpose', 'Image'], type="value", default='Image', label='Input Types\n (You can input an image or a openpose map)')
                    fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed to produce a fixed output)')
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the openpose to the result)", minimum=0, maximum=1, value=1, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
        ips = [input_img, type_in, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo

def create_demo_sketch(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Sketch)')
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(source='upload', type="numpy")
                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                with gr.Row():
                    type_in = gr.inputs.Radio(['Sketch', 'Image'], type="value", default='Image', label='Input Types\n (You can input an image or a sketch)')
                    color_back = gr.inputs.Radio(['White', 'Black'], type="value", default='Black', label='Color of the sketch background\n (Only work for sketch input)')
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the sketch to the result)", minimum=0, maximum=1, value=0.4, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed)')
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
            ips = [input_img, type_in, color_back, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo

def create_demo_color_sketch(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Color + Sketch)')
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    input_img_sketch = gr.Image(source='upload', type="numpy", label='Sketch guidance')
                    input_img_color = gr.Image(source='upload', type="numpy", label='Color guidance')
                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                type_in_color = gr.inputs.Radio(['ColorMap', 'Image'], type="value", default='Image', label='Input Types of Color\n (You can input an image or a color map)')
                with gr.Row():
                    type_in = gr.inputs.Radio(['Sketch', 'Image'], type="value", default='Image', label='Input Types of Sketch\n (You can input an image or a sketch)')
                    color_back = gr.inputs.Radio(['White', 'Black'], type="value", default='Black', label='Color of the sketch background\n (Only work for sketch input)')
                with gr.Row():
                    w_sketch = gr.Slider(label="Depth guidance weight", minimum=0, maximum=2, value=1.0, step=0.1)
                    w_color = gr.Slider(label="Color guidance weight", minimum=0, maximum=2, value=1.2, step=0.1)
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the sketch to the result)", minimum=0, maximum=1, value=0.4, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed)')
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=3, height='auto')
            ips = [input_img_sketch, input_img_color, type_in, type_in_color, w_sketch, w_color, color_back, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo

def create_demo_style_sketch(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Style + Sketch)')
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    input_img_sketch = gr.Image(source='upload', type="numpy", label='Sketch guidance')
                    input_img_style = gr.Image(source='upload', type="numpy", label='Style guidance')
                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                with gr.Row():
                    type_in = gr.inputs.Radio(['Sketch', 'Image'], type="value", default='Image', label='Input Types of Sketch\n (You can input an image or a sketch)')
                    color_back = gr.inputs.Radio(['White', 'Black'], type="value", default='Black', label='Color of the sketch background\n (Only work for sketch input)')
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the sketch to the result)", minimum=0, maximum=1, value=1, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed)')
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
            ips = [input_img_sketch, input_img_style, type_in, color_back, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo

def create_demo_color(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Color)')
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(source='upload', type="numpy", label='Color guidance')
                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                type_in_color = gr.inputs.Radio(['ColorMap', 'Image'], type="value", default='Image', label='Input Types of Color\n (You can input an image or a color map)')
                w_color = gr.Slider(label="Color guidance weight", minimum=0, maximum=2, value=1, step=0.1)
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the sketch to the result)", minimum=0, maximum=1, value=1, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed)')
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
            ips = [input_img, prompt, neg_prompt, pos_prompt, w_color, type_in_color, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo

def create_demo_seg(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Segmentation)')
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(source='upload', type="numpy")
                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                with gr.Row():
                    type_in = gr.inputs.Radio(['Segmentation', 'Image'], type="value", default='Image', label='You can input an image or a segmentation. If you choose to input a segmentation, it must correspond to the coco-stuff')
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the segmentation to the result)", minimum=0, maximum=1, value=1, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed)')
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
            ips = [input_img, type_in, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo

def create_demo_depth(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Depth)')
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(source='upload', type="numpy")
                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                with gr.Row():
                    type_in = gr.inputs.Radio(['Depth', 'Image'], type="value", default='Image', label='You can input an image or a depth map')
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the depth map to the result)", minimum=0, maximum=1, value=1, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed)')
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
            ips = [input_img, type_in, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo

def create_demo_depth_keypose(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Depth & Keypose)')
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    input_img_depth = gr.Image(source='upload', type="numpy", label='Depth guidance')
                    input_img_keypose = gr.Image(source='upload', type="numpy", label='Keypose guidance')

                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                with gr.Row():
                    type_in_depth = gr.inputs.Radio(['Depth', 'Image'], type="value", default='Image', label='You can input an image or a depth map')
                    type_in_keypose = gr.inputs.Radio(['Keypose', 'Image'], type="value", default='Image', label='You can input an image or a keypose map (mmpose style)')
                with gr.Row():
                    w_depth = gr.Slider(label="Depth guidance weight", minimum=0, maximum=2, value=1.0, step=0.1)
                    w_keypose = gr.Slider(label="Keypose guidance weight", minimum=0, maximum=2, value=1.5, step=0.1)
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the multi-guidance to the result)", minimum=0, maximum=1, value=1, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed)')
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=3, height='auto')
            ips = [input_img_depth, input_img_keypose, type_in_depth, type_in_keypose, w_depth, w_keypose, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo

def create_demo_draw(process):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown('## T2I-Adapter (Hand-free drawing)')
        with gr.Row():
            with gr.Column():
                create_button = gr.Button(label="Start", value='Hand-free drawing')
                input_img = gr.Image(source='upload', type="numpy",tool='sketch')
                create_button.click(fn=create_map, outputs=[input_img], queue=False)
                prompt = gr.Textbox(label="Prompt")
                neg_prompt = gr.Textbox(label="Negative Prompt",
                value='ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face')
                pos_prompt = gr.Textbox(label="Positive Prompt",
                value = 'crafted, elegant, meticulous, magnificent, maximum details, extremely hyper aesthetic, intricately detailed')
                run_button = gr.Button(label="Run")
                con_strength = gr.Slider(label="Controling Strength (The guidance strength of the sketch to the result)", minimum=0, maximum=1, value=0.4, step=0.1)
                scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                fix_sample = gr.inputs.Radio(['True', 'False'], type="value", default='False', label='Fix Sampling\n (Fix the random seed)')
                base_model = gr.inputs.Radio(['sd-v1-4.ckpt', 'anything-v4.0-pruned.ckpt'], type="value", default='sd-v1-4.ckpt', label='The base model you want to use')
            with gr.Column():
                result = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
            ips = [input_img, prompt, neg_prompt, pos_prompt, fix_sample, scale, con_strength, base_model]
        run_button.click(fn=process, inputs=ips, outputs=[result])
    return demo