Add attention slicing support (#1087)

apps/stable_diffusion/src/models/model_wrappers.py

@@ -226,6 +226,11 @@ class SharkifyStableDiffusionModel:
                 )
                 self.in_channels = self.unet.in_channels
                 self.train(False)
+                if(args.attention_slicing is not None and args.attention_slicing != "none"):
+                    if(args.attention_slicing.isdigit()):
+                        self.unet.set_attention_slice(int(args.attention_slicing))
+                    else:
+                        self.unet.set_attention_slice(args.attention_slicing)
 
             def forward(
                 self, latent, timestep, text_embedding, guidance_scale

apps/stable_diffusion/src/utils/stable_args.py

@@ -263,6 +263,13 @@ p.add_argument(
     help="Use the accelerate package to reduce cpu memory consumption",
 )
 
+p.add_argument(
+    "--attention_slicing",
+    type=str,
+    default="none",
+    help="Amount of attention slicing to use (one of 'max', 'auto', 'none', or an integer)",
+)
+
 ##############################################################################
 ### IREE - Vulkan supported flags
 ##############################################################################