Update app.py

app.py

@@ -105,6 +105,9 @@ import torch
 import time 
 import gc
 
+import torch.nn.functional as F
+from sageattention import sageattn
+
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
@@ -388,7 +391,6 @@ def uploadNote(prompt,num_inference_steps,guidance_scale,timestamp):
     return filename
     
 '''
-
 pyx = cyper.inline(code, fast_indexing=True, directives=dict(boundscheck=False, wraparound=False, language_level=3))
 
 @spaces.GPU(duration=40)
@@ -401,9 +403,14 @@ def generate_30(
     height: int = 768,
     guidance_scale: float = 4,
     num_inference_steps: int = 125,
+    sage: bool = False,
     use_resolution_binning: bool = True, 
     progress=gr.Progress(track_tqdm=True)
 ):
+    if Sage==True:
+        F.scaled_dot_product_attention = sageattn
+    if Sage==False:
+        F.scaled_dot_product_attention = F.scaled_dot_product_attention
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
     options = {
@@ -450,6 +457,7 @@ def generate_60(
     height: int = 768,
     guidance_scale: float = 4,
     num_inference_steps: int = 125,
+    sage: bool = False,
     use_resolution_binning: bool = True, 
     progress=gr.Progress(track_tqdm=True)
 ):
@@ -494,6 +502,7 @@ def generate_90(
     height: int = 768,
     guidance_scale: float = 4,
     num_inference_steps: int = 125,
+    sage: bool = False,
     use_resolution_binning: bool = True, 
     progress=gr.Progress(track_tqdm=True)
 ):
@@ -622,6 +631,15 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
                 step=10,
                 value=180,
             )
+            options = [True, False]
+            sage = gr.Radio(
+                show_label=True,
+                container=True,
+                interactive=True,
+                choices=options,
+                value=False,
+                label="Use SageAttention: ",
+            )
 
     gr.Examples(
         examples=examples,
@@ -651,6 +669,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            sage,
         ],
         outputs=[result],
     )
@@ -670,6 +689,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            sage,
         ],
         outputs=[result],
     )
@@ -689,6 +709,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            sage,
         ],
         outputs=[result],
     )