sd2 oneflow compile

examples/text_to_image_sdxl_fp16.py

@@ -41,8 +41,7 @@
         os.environ["ONEFLOW_MLIR_ENABLE_ROUND_TRIP"] = "1"
         os.environ["ONEFLOW_MLIR_FUSE_FORWARD_OPS"] = "1"
         os.environ["ONEFLOW_MLIR_FUSE_OPS_WITH_BACKWARD_IMPL"] = "1"
-        # Open this will raise error
-        # os.environ["ONEFLOW_MLIR_GROUP_MATMUL"] = "1"
+        os.environ["ONEFLOW_MLIR_GROUP_MATMUL"] = "1"
         os.environ["ONEFLOW_MLIR_PREFER_NHWC"] = "1"
         os.environ["ONEFLOW_KERNEL_ENABLE_FUSED_CONV_BIAS"] = "1"
         os.environ["ONEFLOW_KERNEL_ENABLE_FUSED_LINEAR"] = "1"
@@ -56,5 +55,5 @@
 pipe.to("cuda")
 
 for i in range(3):
-    image = pipe(prompt=args.prompt).images[0]
+    image = pipe(prompt=args.prompt, height=96, width=128, num_inference_steps=50).images[0]
     image.save(f"{i}-{args.saved_image}")

examples/text_to_image_sdxl_fp16_base_and_refiner_with_oneflow_compile.py

@@ -0,0 +1,92 @@
+import os
+import argparse
+# cv2 must be imported before diffusers and oneflow to avlid error: AttributeError: module 'cv2.gapi' has no attribute 'wip'
+# Maybe bacause oneflow use a lower version of cv2
+import cv2
+import oneflow as flow
+import torch
+# oneflow_compile should be imported before importing any diffusers
+from onediff.infer_compiler import oneflow_compile 
+from diffusers import DiffusionPipeline
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--base", type=str, default="/share_nfs/hf_models/stable-diffusion-xl-base-1.0"
+)
+parser.add_argument(
+    "--refiner", type=str, default="stabilityai/stable-diffusion-xl-refiner-1.0"
+)
+parser.add_argument("--variant", type=str, default="fp16")
+parser.add_argument(
+    "--prompt",
+    type=str,
+    default="street style, detailed, raw photo, woman, face, shot on CineStill 800T",
+)
+parser.add_argument("--n_steps", type=int, default=30)
+parser.add_argument("--saved_image", type=str, required=False, default="xl-refiner-out.png")
+parser.add_argument("--seed", type=int, default=1)
+parser.add_argument("--compile", action=argparse.BooleanOptionalAction)
+parser.add_argument("--save", action=argparse.BooleanOptionalAction)
+parser.add_argument("--load", action=argparse.BooleanOptionalAction)
+parser.add_argument("--file", type=str, required=False, default="unet_compiled")
+cmd_args = parser.parse_args()
+
+# Normal SDXL
+seed = torch.manual_seed(cmd_args.seed)
+# SDXL base: StableDiffusionXLPipeline
+base = DiffusionPipeline.from_pretrained(
+    cmd_args.base,
+    torch_dtype=torch.float16,
+    variant=cmd_args.variant,
+    use_safetensors=True,
+)
+base.to("cuda")
+# SDXL refiner: StableDiffusionXLImg2ImgPipeline
+refiner = DiffusionPipeline.from_pretrained(
+    cmd_args.refiner,
+    text_encoder_2=base.text_encoder_2,
+    vae=base.vae,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant=cmd_args.variant,
+)
+refiner.to("cuda")
+
+# Compile unet with oneflow
+if cmd_args.compile:
+    base.unet = oneflow_compile(base.unet)
+    refiner.unet = oneflow_compile(refiner.unet)
+    print("unet is compiled to oneflow.")
+    if cmd_args.load:
+        # Load compiled unet with oneflow
+        print("loading graphs...")
+        base.unet._graph_load("base_" + cmd_args.file)
+        refiner.unet._graph_load("refiner_" + cmd_args.file)
+
+# Normal SDXL call
+sizes = [1024, 896, 768]
+#sizes = [1024]
+for h in sizes:
+    for w in sizes:
+        for i in range(3):
+            image = base(
+                prompt=cmd_args.prompt,
+                height=h,
+                width=w,
+                generator=seed,
+                num_inference_steps=cmd_args.n_steps,
+                output_type="latent",
+            ).images
+            image = refiner(
+                prompt=cmd_args.prompt,
+                generator=seed,
+                num_inference_steps=cmd_args.n_steps,
+                image=image,
+            ).images[0]
+            image.save(f"h{h}-w{w}-i{i}-{cmd_args.saved_image}")
+
+# Save compiled unet with oneflow
+if cmd_args.compile and cmd_args.save:
+    print("saving graphs...")
+    base.unet._graph_save("base_" + cmd_args.file)
+    refiner.unet._graph_save("refiner_" + cmd_args.file)

examples/text_to_image_sdxl_fp16_with_oneflow_compile.py

@@ -0,0 +1,58 @@
+import os
+import argparse
+# cv2 must be imported before diffusers and oneflow to avlid error: AttributeError: module 'cv2.gapi' has no attribute 'wip'
+# Maybe bacause oneflow use a lower version of cv2
+import cv2
+import oneflow as flow
+import torch
+# oneflow_compile should be imported before importing any diffusers
+from onediff.infer_compiler import oneflow_compile 
+from diffusers import StableDiffusionXLPipeline
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--model", type=str, default="/share_nfs/hf_models/stable-diffusion-xl-base-1.0"
+)
+parser.add_argument("--variant", type=str, default="fp16")
+parser.add_argument(
+    "--prompt",
+    type=str,
+    default="street style, detailed, raw photo, woman, face, shot on CineStill 800T",
+)
+parser.add_argument("--saved_image", type=str, required=False, default="xl-base-out.png")
+parser.add_argument("--seed", type=int, default=1)
+parser.add_argument("--compile", action=argparse.BooleanOptionalAction)
+parser.add_argument("--save", action=argparse.BooleanOptionalAction)
+parser.add_argument("--load", action=argparse.BooleanOptionalAction)
+parser.add_argument("--file", type=str, required=False, default="unet_compiled")
+cmd_args = parser.parse_args()
+
+# Normal SDXL
+torch.manual_seed(cmd_args.seed)
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    cmd_args.model, torch_dtype=torch.float16, variant=cmd_args.variant, use_safetensors=True
+)
+pipe.to("cuda")
+
+# Compile unet with oneflow
+if cmd_args.compile:
+    pipe.unet = oneflow_compile(pipe.unet)
+    print("unet is compiled to oneflow.")
+    if cmd_args.load:
+        # Load compiled unet with oneflow
+        print("loading graphs...")
+        pipe.unet._graph_load(cmd_args.file)
+
+# Normal SDXL call
+sizes = [1024, 896, 768]
+#sizes = [1024]
+for h in sizes:
+    for w in sizes:
+        for i in range(3):
+            image = pipe(prompt=cmd_args.prompt, height=h, width=w, num_inference_steps=30).images[0]
+            image.save(f"h{h}-w{w}-i{i}-{cmd_args.saved_image}")
+
+# Save compiled unet with oneflow
+if cmd_args.compile and cmd_args.save:
+    print("saving graphs...")
+    pipe.unet._graph_save(cmd_args.file)

examples/unet_torch_interplay.py

@@ -1,4 +1,5 @@
 import os
+import random
 
 os.environ["ONEFLOW_MLIR_CSE"] = "1"
 os.environ["ONEFLOW_MLIR_ENABLE_INFERENCE_OPTIMIZATION"] = "1"
@@ -20,6 +21,9 @@
 os.environ["ONEFLOW_LINEAR_EMBEDDING_SKIP_INIT"] = "1"
 
 import click
+# cv2 must be imported before diffusers and oneflow to avlid error: AttributeError: module 'cv2.gapi' has no attribute 'wip'
+# Maybe bacause oneflow use a lower version of cv2
+import cv2
 import oneflow as flow
 from tqdm import tqdm
 from dataclasses import dataclass, fields
@@ -45,13 +49,13 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         flow.mock_torch.disable()
 
 
-def get_unet(token, _model_id):
+def get_unet(token, _model_id, revision):
     from diffusers import UNet2DConditionModel
 
     unet = UNet2DConditionModel.from_pretrained(
         _model_id,
         use_auth_token=token,
-        revision="fp16",
+        revision=revision,
         torch_dtype=flow.float16,
         subfolder="unet",
     )
@@ -67,14 +71,15 @@ def __init__(self, unet):
         self.unet = unet
         self.config.enable_cudnn_conv_heuristic_search_algo(False)
         self.config.allow_fuse_add_to_output(True)
+        self.debug(0)
 
-    def build(self, latent_model_input, t, text_embeddings):
+    def build(self, latent_model_input, t, text_embeddings, added_cond_kwargs=None):
         text_embeddings = flow._C.amp_white_identity(text_embeddings)
         return self.unet(
-            latent_model_input, t, encoder_hidden_states=text_embeddings
+            latent_model_input, t, encoder_hidden_states=text_embeddings, added_cond_kwargs=added_cond_kwargs,
         ).sample
 
-    def warmup_with_arg(self, arg_meta_of_sizes):
+    def warmup_with_arg(self, arg_meta_of_sizes, added):
         for arg_metas in arg_meta_of_sizes:
             print(f"warmup {arg_metas=}")
             arg_tensors = [
@@ -87,7 +92,7 @@ def warmup_with_arg(self, arg_meta_of_sizes):
                     dtype=arg_metas.gettype("cross_attention_dim"),
                 ).to("cuda"),
             ]
-            self(*arg_tensors)  # build and warmup
+            self(*arg_tensors, added)  # build and warmup
 
     def warmup_with_load(self, file_path):
         state_dict = flow.load(file_path)
@@ -131,7 +136,6 @@ def get_arg_meta_of_sizes(
         for j in resolution_scales
     ]
 
-
 @click.command()
 @click.option("--token")
 @click.option("--repeat", default=100)
@@ -140,37 +144,48 @@ def get_arg_meta_of_sizes(
 @click.option("--load", is_flag=True)
 @click.option("--file", type=str, default="./unet_graphs")
 @click.option("--model_id", type=str, default="runwayml/stable-diffusion-v1-5")
-def benchmark(token, repeat, sync_interval, save, load, file, model_id):
+@click.option("--revision", type=str, default="fp16")
+def benchmark(token, repeat, sync_interval, save, load, file, model_id, revision):
     RESOLUTION_SCALES = [2, 1, 0]
     BATCH_SIZES = [2]
     # TODO: reproduce bug caused by changing batch
     # BATCH_SIZES = [4, 2]
 
-    num_channels = 4
     # create a mocked unet graph
+    # unet mock should be placed before importing any diffusers
     with MockCtx():
-        unet = get_unet(token, model_id)
+        unet = get_unet(token, model_id, revision)
         unet_graph = UNetGraphWithCache(unet)
-        cross_attention_dim = unet.config["cross_attention_dim"]
-        warmup_meta_of_sizes = get_arg_meta_of_sizes(
-            batch_sizes=BATCH_SIZES,
-            resolution_scales=RESOLUTION_SCALES,
-            num_channels=num_channels,
-            cross_attention_dim=cross_attention_dim,
-        )
-        for (i, m) in enumerate(warmup_meta_of_sizes):
-            print(f"warmup case #{i + 1}:", m)
-        if load == True:
-            print("loading graphs...")
-            unet_graph.warmup_with_load(file)
-        else:
-            print("warmup with arguments...")
-            unet_graph.warmup_with_arg(warmup_meta_of_sizes)
 
-    # generate inputs with torch
+    num_channels = 4
+    cross_attention_dim = unet.config["cross_attention_dim"]
     from diffusers.utils import floats_tensor
     import torch
+    if model_id == "stabilityai/stable-diffusion-xl-base-1.0":
+        # sdxl needed
+        add_text_embeds = flow.utils.tensor.from_torch(floats_tensor((2, 1280)).to("cuda").to(torch.float16))
+        add_time_ids = flow.utils.tensor.from_torch(floats_tensor((2, 6)).to("cuda").to(torch.float16))
+        added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+    else:
+        added_cond_kwargs = None
+
+    warmup_meta_of_sizes = get_arg_meta_of_sizes(
+        batch_sizes=BATCH_SIZES,
+        resolution_scales=RESOLUTION_SCALES,
+        num_channels=num_channels,
+        cross_attention_dim=cross_attention_dim,
+    )
+    for (i, m) in enumerate(warmup_meta_of_sizes):
+        print(f"warmup case #{i + 1}:", m)
+
+    if load == True:
+        print("loading graphs...")
+        unet_graph.warmup_with_load(file)
+    else:
+        print("warmup with arguments...")
+        unet_graph.warmup_with_arg(warmup_meta_of_sizes, added_cond_kwargs)
 
+    # generate inputs with torch
     time_step = torch.tensor([10]).to("cuda")
     encoder_hidden_states_of_sizes = {
         batch_size: floats_tensor((batch_size, 77, cross_attention_dim))
@@ -187,6 +202,7 @@ def benchmark(token, repeat, sync_interval, save, load, file, model_id):
         k: flow.utils.tensor.from_torch(v)
         for k, v in encoder_hidden_states_of_sizes.items()
     }
+
     # convert to oneflow tensors
     time_step = flow.utils.tensor.from_torch(time_step)
 
@@ -199,7 +215,7 @@ def benchmark(token, repeat, sync_interval, save, load, file, model_id):
 
         noise = random.choice(noise_of_sizes)
         encoder_hidden_states = encoder_hidden_states_of_sizes[noise.shape[0]]
-        out = unet_graph(noise, time_step, encoder_hidden_states)
+        out = unet_graph(noise, time_step, encoder_hidden_states, added_cond_kwargs)
         # convert to torch tensors
         out = flow.utils.tensor.to_torch(out)
         if r == repeat - 1 or r % sync_interval == 0:

src/onediff/infer_compiler/__init__.py

@@ -1,6 +1,7 @@
 import os
 import torch
 import oneflow as flow
+from .with_oneflow_compile import oneflow_compile 
 from .with_fx_interpreter import OneFlowInterpreter
 from .with_fx_graph import fx_node_tranform