feat: add web ui for core ml stable diffusion

Author: soulteary

python_coreml_stable_diffusion/web.py

@@ -0,0 +1,97 @@
+
+import python_coreml_stable_diffusion.pipeline as pipeline
+
+import gradio as gr
+from diffusers import StableDiffusionPipeline
+
+def init(args):
+    pipeline.logger.info("Initializing PyTorch pipe for reference configuration")
+    pytorch_pipe = StableDiffusionPipeline.from_pretrained(args.model_version,
+                                                           use_auth_token=True)
+
+    user_specified_scheduler = None
+    if args.scheduler is not None:
+        user_specified_scheduler = pipeline.SCHEDULER_MAP[
+            args.scheduler].from_config(pytorch_pipe.scheduler.config)
+
+    coreml_pipe = pipeline.get_coreml_pipe(pytorch_pipe=pytorch_pipe,
+                                  mlpackages_dir=args.i,
+                                  model_version=args.model_version,
+                                  compute_unit=args.compute_unit,
+                                  scheduler_override=user_specified_scheduler)
+
+
+    def infer(prompt, steps):
+        pipeline.logger.info("Beginning image generation.")
+        image = coreml_pipe(
+            prompt=prompt,
+            height=coreml_pipe.height,
+            width=coreml_pipe.width,
+            num_inference_steps=steps,
+        )
+        images = []
+        images.append(image["images"][0])
+        return images
+
+
+    demo = gr.Blocks()
+
+    with demo:
+        gr.Markdown(
+            "<center><h1>Core ML Stable Diffusion</h1>Run Stable Diffusion on Apple Silicon with Core ML</center>")
+        with gr.Group():
+            with gr.Box():
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Row():
+                            text = gr.Textbox(
+                                label="Prompt",
+                                lines=11,
+                                placeholder="Enter your prompt",
+                            )
+                        with gr.Row():
+                            btn = gr.Button("Generate image")
+                        with gr.Row():
+                            steps = gr.Slider(label="Steps", minimum=1,
+                                            maximum=50, value=10, step=1)
+                    with gr.Column():
+                        gallery = gr.Gallery(
+                            label="Generated image", elem_id="gallery"
+                        )
+
+            text.submit(infer, inputs=[text, steps], outputs=gallery)
+            btn.click(infer, inputs=[text, steps], outputs=gallery)
+
+    demo.launch(debug=True, server_name="0.0.0.0")
+
+
+if __name__ == "__main__":
+    parser = pipeline.argparse.ArgumentParser()
+
+    parser.add_argument(
+        "-i",
+        required=True,
+        help=("Path to input directory with the .mlpackage files generated by "
+              "python_coreml_stable_diffusion.torch2coreml"))
+    parser.add_argument(
+        "--model-version",
+        default="CompVis/stable-diffusion-v1-4",
+        help=
+        ("The pre-trained model checkpoint and configuration to restore. "
+         "For available versions: https://huggingface.co/models?search=stable-diffusion"
+         ))
+    parser.add_argument(
+        "--compute-unit",
+        choices=pipeline.get_available_compute_units(),
+        default="ALL",
+        help=("The compute units to be used when executing Core ML models. "
+              f"Options: {pipeline.get_available_compute_units()}"))
+    parser.add_argument(
+        "--scheduler",
+        choices=tuple(pipeline.SCHEDULER_MAP.keys()),
+        default=None,
+        help=("The scheduler to use for running the reverse diffusion process. "
+             "If not specified, the default scheduler from the diffusers pipeline is utilized"))
+
+    args = parser.parse_args()
+    init(args)

requirements.txt

@@ -2,4 +2,5 @@ coremltools
 diffusers[torch]
 torch
 transformers
-scipy
+scipy
+gradio

setup.py

@@ -19,6 +19,7 @@
         "torch",
         "transformers",
         "scipy",
+        "gradio",
     ],
     packages=find_packages(),
     classifiers=[