[SD][web] Add 64 max_length support in SD web

Signed-Off-by: Gaurav Shukla <[email protected]>

Author: Shukla-Gaurav

web/models/stable_diffusion/model_wrappers.py

@@ -10,34 +10,37 @@
     "v1.4": "CompVis/stable-diffusion-v1-4",
 }
 
+# clip has 2 variants of max length 77 or 64.
+model_clip_max_length = 64 if args.max_length == 64 else 77
+
 model_input = {
     "v2.1": {
-        "clip": (torch.randint(1, 2, (2, 77)),),
+        "clip": (torch.randint(1, 2, (2, model_clip_max_length)),),
         "vae": (torch.randn(1, 4, 96, 96),),
         "unet": (
             torch.randn(1, 4, 96, 96),  # latents
             torch.tensor([1]).to(torch.float32),  # timestep
-            torch.randn(2, 77, 1024),  # embedding
+            torch.randn(2, model_clip_max_length, 1024),  # embedding
             torch.tensor(1).to(torch.float32),  # guidance_scale
         ),
     },
     "v2.1base": {
-        "clip": (torch.randint(1, 2, (2, 77)),),
+        "clip": (torch.randint(1, 2, (2, model_clip_max_length)),),
         "vae": (torch.randn(1, 4, 64, 64),),
         "unet": (
             torch.randn(1, 4, 64, 64),  # latents
             torch.tensor([1]).to(torch.float32),  # timestep
-            torch.randn(2, 77, 1024),  # embedding
+            torch.randn(2, model_clip_max_length, 1024),  # embedding
             torch.tensor(1).to(torch.float32),  # guidance_scale
         ),
     },
     "v1.4": {
-        "clip": (torch.randint(1, 2, (2, 77)),),
+        "clip": (torch.randint(1, 2, (2, model_clip_max_length)),),
         "vae": (torch.randn(1, 4, 64, 64),),
         "unet": (
             torch.randn(1, 4, 64, 64),
             torch.tensor([1]).to(torch.float32),  # timestep
-            torch.randn(2, 77, 768),
+            torch.randn(2, model_clip_max_length, 768),
             torch.tensor(1).to(torch.float32),
         ),
     },

web/models/stable_diffusion/opt_params.py

@@ -46,7 +46,10 @@ def get_unet():
             bucket = "gs://shark_tank/stable_diffusion"
             model_name = "unet_8dec_fp16"
             if args.version == "v2.1base":
-                model_name = "unet2base_8dec_fp16"
+                if args.max_length == 64:
+                    model_name = "unet_19dec_v2p1base_fp16_64"
+                else:
+                    model_name = "unet2base_8dec_fp16"
             if args.version == "v2.1":
                 model_name = "unet2_14dec_fp16"
             iree_flags += [
@@ -149,7 +152,10 @@ def get_clip():
     bucket = "gs://shark_tank/stable_diffusion"
     model_name = "clip_18dec_fp32"
     if args.version == "v2.1base":
-        model_name = "clip2base_18dec_fp32"
+        if args.max_length == 64:
+            model_name = "clip_19dec_v2p1base_fp32_64"
+        else:
+            model_name = "clip2base_18dec_fp32"
     if args.version == "v2.1":
         model_name = "clip2_18dec_fp32"
     iree_flags += [