Add conditions to force use --import_mlir (huggingface#1028)

Author: yzhang93

apps/stable_diffusion/src/utils/sd_annotation.py

@@ -91,7 +91,12 @@ def load_lower_configs():
             spec = spec.split("-")[0]
 
     if args.annotation_model == "vae":
-        config_name = f"{args.annotation_model}_{args.precision}_{device}.json"
+        if not spec or spec in ["rdna3", "sm_80"]:
+            config_name = (
+                f"{args.annotation_model}_{args.precision}_{device}.json"
+            )
+        else:
+            config_name = f"{args.annotation_model}_{args.precision}_{device}_{spec}.json"
     else:
         if not spec or spec in ["rdna3", "sm_80"]:
             config_name = f"{args.annotation_model}_{version}_{args.precision}_{device}.json"

apps/stable_diffusion/src/utils/utils.py

@@ -240,7 +240,12 @@ def set_init_device_flags():
 
     # Use tuned models in the case of fp16, vulkan rdna3 or cuda sm devices.
     if (
-        args.ckpt_loc != ""
+        args.hf_model_id
+        in [
+            "runwayml/stable-diffusion-inpainting",
+            "stabilityai/stable-diffusion-2-inpainting",
+        ]
+        or args.ckpt_loc != ""
         or args.precision != "fp16"
         or args.height != 512
         or args.width != 512
@@ -288,6 +293,27 @@ def set_init_device_flags():
     elif args.height != 512 or args.width != 512 or args.batch_size != 1:
         args.import_mlir = True
 
+    elif args.use_tuned and args.hf_model_id in [
+        "dreamlike-art/dreamlike-diffusion-1.0",
+        "prompthero/openjourney",
+        "stabilityai/stable-diffusion-2-1",
+    ]:
+        args.import_mlir = True
+
+    elif (
+        args.use_tuned
+        and "vulkan" in args.device
+        and "rdna2" in args.iree_vulkan_target_triple
+    ):
+        args.import_mlir = True
+
+    elif (
+        args.use_tuned
+        and "cuda" in args.device
+        and get_cuda_sm_cc() == "sm_89"
+    ):
+        args.import_mlir = True
+
 
 # Utility to get list of devices available.
 def get_available_devices():