Add Anything v3 and AnalogDiffusion variants of SD (huggingface#685)

* base support for anythingv3

* add analogdiffusiont

* Update readme

* keep max len 77 till support for 64 added for variants

* lint fix

Author: PhaneeshB

shark/examples/shark_inference/stable_diffusion/README.md

@@ -42,3 +42,15 @@ unzip ~/.local/shark_tank/<your unet>/inputs.npz
 
 iree-benchmark-module --module_file=/path/to/output/vmfb --entry_function=forward --function_input=@arr_0.npy --function_input=1xf16 --function_input=@arr_2.npy --function_input=@arr_3.npy --function_input=@arr_4.npy  
 ```
+
+## Using other supported Stable Diffusion variants with SHARK:
+
+Currently we support the following fine-tuned versions of Stable Diffusion:
+- [AnythingV3](https://huggingface.co/Linaqruf/anything-v3.0)
+- [Analog Diffusion](https://huggingface.co/wavymulder/Analog-Diffusion)
+
+use the flag `--variant=` to specify the model to be used.
+
+```shell
+python .\shark\examples\shark_inference\stable_diffusion\main.py --variant=anythingv3 --max_length=77 --prompt="1girl, brown hair, green eyes, colorful, autumn, cumulonimbus clouds, lighting, blue sky, falling leaves, garden"
+```

shark/examples/shark_inference/stable_diffusion/main.py

@@ -50,7 +50,7 @@ def end_profiling(device):
     neg_prompt = args.negative_prompts
     height = 512  # default height of Stable Diffusion
     width = 512  # default width of Stable Diffusion
-    if args.version == "v2.1":
+    if args.version == "v2.1" and args.variant == "stablediffusion":
         height = 768
         width = 768
 
@@ -71,9 +71,9 @@ def end_profiling(device):
         sys.exit("prompts and negative prompts must be of same length")
 
     set_iree_runtime_flags()
+    clip = get_clip()
     unet = get_unet()
     vae = get_vae()
-    clip = get_clip()
     if args.dump_isa:
         dump_isas(args.dispatch_benchmarks_dir)
 

shark/examples/shark_inference/stable_diffusion/model_wrappers.py

@@ -12,6 +12,16 @@
 
 # clip has 2 variants of max length 77 or 64.
 model_clip_max_length = 64 if args.max_length == 64 else 77
+if args.variant != "stablediffusion":
+    model_clip_max_length = 77
+
+model_variant = {
+    "stablediffusion": "SD",
+    "anythingv3": "Linaqruf/anything-v3.0",
+    "dreamlike": "dreamlike-art/dreamlike-diffusion-1.0",
+    "openjourney": "prompthero/openjourney",
+    "analogdiffusion": "wavymulder/Analog-Diffusion",
+}
 
 model_input = {
     "v2.1": {
@@ -47,18 +57,32 @@
 }
 
 # revision param for from_pretrained defaults to "main" => fp32
-model_revision = "fp16" if args.precision == "fp16" else "main"
+model_revision = {
+    "stablediffusion": "fp16" if args.precision == "fp16" else "main",
+    "anythingv3": "diffusers",
+    "analogdiffusion": "main",
+}
 
 
 def get_clip_mlir(model_name="clip_text", extra_args=[]):
 
     text_encoder = CLIPTextModel.from_pretrained(
         "openai/clip-vit-large-patch14"
     )
-    if args.version != "v1.4":
+    if args.variant == "stablediffusion":
+        if args.version != "v1.4":
+            text_encoder = CLIPTextModel.from_pretrained(
+                model_config[args.version], subfolder="text_encoder"
+            )
+
+    elif args.variant in ["anythingv3", "analogdiffusion"]:
         text_encoder = CLIPTextModel.from_pretrained(
-            model_config[args.version], subfolder="text_encoder"
+            model_variant[args.variant],
+            subfolder="text_encoder",
+            revision=model_revision[args.variant],
         )
+    else:
+        raise (f"{args.variant} not yet added")
 
     class CLIPText(torch.nn.Module):
         def __init__(self):
@@ -83,9 +107,11 @@ class VaeModel(torch.nn.Module):
         def __init__(self):
             super().__init__()
             self.vae = AutoencoderKL.from_pretrained(
-                model_config[args.version],
+                model_config[args.version]
+                if args.variant == "stablediffusion"
+                else model_variant[args.variant],
                 subfolder="vae",
-                revision=model_revision,
+                revision=model_revision[args.variant],
             )
 
         def forward(self, input):
@@ -96,16 +122,27 @@ def forward(self, input):
             return x.round()
 
     vae = VaeModel()
-    if args.precision == "fp16":
-        vae = vae.half().cuda()
-        inputs = tuple(
-            [
-                inputs.half().cuda()
-                for inputs in model_input[args.version]["vae"]
-            ]
-        )
+    if args.variant == "stablediffusion":
+        if args.precision == "fp16":
+            vae = vae.half().cuda()
+            inputs = tuple(
+                [
+                    inputs.half().cuda()
+                    for inputs in model_input[args.version]["vae"]
+                ]
+            )
+        else:
+            inputs = model_input[args.version]["vae"]
+    elif args.variant in ["anythingv3", "analogdiffusion"]:
+        if args.precision == "fp16":
+            vae = vae.half().cuda()
+            inputs = tuple(
+                [inputs.half().cuda() for inputs in model_input["v1.4"]["vae"]]
+            )
+        else:
+            inputs = model_input["v1.4"]["vae"]
     else:
-        inputs = model_input[args.version]["vae"]
+        raise (f"{args.variant} not yet added")
 
     shark_vae = compile_through_fx(
         vae,
@@ -121,9 +158,11 @@ class UnetModel(torch.nn.Module):
         def __init__(self):
             super().__init__()
             self.unet = UNet2DConditionModel.from_pretrained(
-                model_config[args.version],
+                model_config[args.version]
+                if args.variant == "stablediffusion"
+                else model_variant[args.variant],
                 subfolder="unet",
-                revision=model_revision,
+                revision=model_revision[args.variant],
             )
             self.in_channels = self.unet.in_channels
             self.train(False)
@@ -141,16 +180,30 @@ def forward(self, latent, timestep, text_embedding, guidance_scale):
             return noise_pred
 
     unet = UnetModel()
-    if args.precision == "fp16":
-        unet = unet.half().cuda()
-        inputs = tuple(
-            [
-                inputs.half().cuda() if len(inputs.shape) != 0 else inputs
-                for inputs in model_input[args.version]["unet"]
-            ]
-        )
+    if args.variant == "stablediffusion":
+        if args.precision == "fp16":
+            unet = unet.half().cuda()
+            inputs = tuple(
+                [
+                    inputs.half().cuda() if len(inputs.shape) != 0 else inputs
+                    for inputs in model_input[args.version]["unet"]
+                ]
+            )
+        else:
+            inputs = model_input[args.version]["unet"]
+    elif args.variant in ["anythingv3", "analogdiffusion"]:
+        if args.precision == "fp16":
+            unet = unet.half().cuda()
+            inputs = tuple(
+                [
+                    inputs.half().cuda() if len(inputs.shape) != 0 else inputs
+                    for inputs in model_input["v1.4"]["unet"]
+                ]
+            )
+        else:
+            inputs = model_input["v1.4"]["unet"]
     else:
-        inputs = model_input[args.version]["unet"]
+        raise (f"{args.variant} is not yet added")
     shark_unet = compile_through_fx(
         unet,
         inputs,