import gradio as gr
import torch

from diffusers import LMSDiscreteScheduler
from mixdiff import StableDiffusionCanvasPipeline, Text2ImageRegion

# Creater scheduler and model (similar to StableDiffusionPipeline)
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
pipeline = StableDiffusionCanvasPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", scheduler=scheduler).to("cuda" if torch.cuda.is_available() else "cpu")

def generate(prompt1, prompt2, prompt3, gc1, gc2, gc3, overlap, steps, seed):
    """Mixture of Diffusers generation"""
    tile_width = 640
    tile_height = 640
    return pipeline(
        canvas_height=tile_height,
        canvas_width=tile_width + (tile_width - overlap) * 2,
        regions=[
            Text2ImageRegion(0, tile_height, 0, tile_width, guidance_scale=gc1,
                prompt=prompt1),
            Text2ImageRegion(0, tile_height, tile_width - overlap, tile_width - overlap + tile_width, guidance_scale=gc2,
                prompt=prompt2),
            Text2ImageRegion(0, tile_height, (tile_width - overlap) * 2, (tile_width - overlap) * 2 + tile_width, guidance_scale=gc3,
                prompt=prompt3),
        ],
        num_inference_steps=steps,
        seed=seed,
    )["sample"][0]

with gr.Blocks(title="Mixture of Diffusers") as demo:
    gr.Markdown("# Mixture of Diffusers")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Left region")
            left_prompt = gr.Textbox(lines=2, label="Prompt")
            left_gs = gr.Slider(minimum=0, maximum=15, value=8, step=1, label="Guidance scale")
        with gr.Column(scale=1):
            gr.Markdown("### Center region")
            center_prompt = gr.Textbox(lines=2, label="Prompt")
            center_gs = gr.Slider(minimum=0, maximum=15, value=8, step=1, label="Guidance scale")
        with gr.Column(scale=1):
            gr.Markdown("### Right region")
            right_prompt = gr.Textbox(lines=2, label="Prompt")
            right_gs = gr.Slider(minimum=0, maximum=15, value=8, step=1, label="Guidance scale")
    gr.Markdown("### General parameters")
    with gr.Row():
        with gr.Column(scale=1):
            overlap = gr.Slider(minimum=128, maximum=320, value=256, step=8, label="Overlap between diffusion regions")
        with gr.Column(scale=1):
            steps = gr.Slider(minimum=1, maximum=50, value=15, step=1, label="Number of diffusion steps")
        with gr.Column(scale=1):
            seed = gr.Number(value=12345, precision=0, label="Random seed")
    with gr.Row():
        button = gr.Button(value="Generate")
    with gr.Row():
        output = gr.Image(label="Generated image")
    with gr.Row():
        gr.Examples(
            examples=[
                [
                    "A charming house in the countryside, by jakub rozalski, sunset lighting, elegant, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",
                    "A dirt road in the countryside crossing pastures, by jakub rozalski, sunset lighting, elegant, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",
                    "An old and rusty giant robot lying on a dirt road, by jakub rozalski, dark sunset lighting, elegant, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",
                    8, 8, 8,
                    256,
                    50,
                    7178915308
                ],
            ],
            inputs=[left_prompt, center_prompt, right_prompt, left_gs, center_gs, right_gs, overlap, steps, seed],
            # outputs=output,
            # fn=generate,
            # cache_examples=True
        )
    
    button.click(
        generate, 
        inputs=[left_prompt, center_prompt, right_prompt, left_gs, center_gs, right_gs, overlap, steps, seed],
        outputs=output
    )

demo.launch(server_name="0.0.0.0")