CPU/GPU support, more options

app.py

@@ -1,26 +1,29 @@
 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:0")
+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, seed):
+def generate(prompt1, prompt2, prompt3, overlap, guidance_scale, steps, seed):
     """Mixture of Diffusers generation"""
+    tile_width = 640
+    tile_height = 640
     return pipeline(
-        canvas_height=640,
-        canvas_width=1408,
+        canvas_height=tile_height,
+        canvas_width=tile_width + (tile_width - overlap) * 2,
         regions=[
-            Text2ImageRegion(0, 640, 0, 640, guidance_scale=8,
+            Text2ImageRegion(0, tile_height, 0, tile_width, guidance_scale=guidance_scale,
                 prompt=prompt1),
-            Text2ImageRegion(0, 640, 384, 1024, guidance_scale=8,
+            Text2ImageRegion(0, tile_height, tile_width - overlap, tile_width - overlap + tile_width, guidance_scale=guidance_scale,
                 prompt=prompt2),
-            Text2ImageRegion(0, 640, 768, 1408, guidance_scale=8,
+            Text2ImageRegion(0, tile_height, (tile_width - overlap) * 2, (tile_width - overlap) * 2 + tile_width, guidance_scale=guidance_scale,
                 prompt=prompt3),
         ],
-        num_inference_steps=50,
+        num_inference_steps=steps,
         seed=seed,
     )["sample"][0]
 
@@ -31,6 +34,9 @@ demo = gr.Interface(
         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",
@@ -39,8 +45,13 @@ demo = gr.Interface(
             "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")