app.py

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, overlap, guidance_scale, 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=guidance_scale,
                prompt=prompt1),
            Text2ImageRegion(0, tile_height, tile_width - overlap, tile_width - overlap + tile_width, guidance_scale=guidance_scale,
                prompt=prompt2),
            Text2ImageRegion(0, tile_height, (tile_width - overlap) * 2, (tile_width - overlap) * 2 + tile_width, guidance_scale=guidance_scale,
                prompt=prompt3),
        ],
        num_inference_steps=steps,
        seed=seed,
    )["sample"][0]


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(lines=2, label="Left region prompt"),
        gr.Textbox(lines=2, label="Center region prompt"),
        gr.Textbox(lines=2, label="Right region prompt"),
        gr.Slider(minimum=128, maximum=320, value=256, step=8, label="Overlap between diffusion regions"),
        gr.Slider(minimum=0, maximum=15, value=8, step=1, label="Global guidance scale"),
        gr.Slider(minimum=1, maximum=50, value=15, step=1, label="Number of diffusion steps"),
        gr.Number(value=12345, precision=0),
    ],
    outputs="image",
    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",
            256,
            8,
            50,
            7178915308
        ],
    ],
    title="Mixture of Diffusers",
    article="",
)
demo.launch(server_name="0.0.0.0")