import gradio as gr

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, use_auth_token=True).to("cuda:0")

def generate(prompt1, prompt2, prompt3, seed)
    """Mixture of Diffusers generation"""
    return pipeline(
        canvas_height=640,
        canvas_width=1408,
        regions=[
            Text2ImageRegion(0, 640, 0, 640, guidance_scale=8,
                prompt=prompt1),
            Text2ImageRegion(0, 640, 384, 1024, guidance_scale=8,
                prompt=prompt2),
            Text2ImageRegion(0, 640, 768, 1408, guidance_scale=8,
                prompt=prompt3),
        ],
        num_inference_steps=50,
        seed=seed,
    )["sample"][0]


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(lines=2, placeholder="Left prompt"),
        gr.Textbox(lines=2, placeholder="Center prompt"),
        gr.Textbox(lines=2, placeholder="Right prompt"),
        gr.Textbox(lines=1, placeholder="Random Seed"),
    ]
    outputs="image"
)
demo.launch()