modify text

app.py

@@ -115,17 +115,17 @@ def binary_mask_eval(args):
         ratio=args.ratio,
     )
 
-    # Print mask sparsity info
-    threshold = None if args.binary else args.lambda_threshold
-    threshold = None if args.scope is not None else threshold
-    name = ["ff", "attn"]
-    for n, hooker in zip(name, hookers):
-        total_num_heads, num_activate_heads, mask_sparsity = calculate_mask_sparsity(hooker, threshold)
-        print(f"model: {args.model}, {n} masking: {args.masking}")
-        print(
-            f"total num heads: {total_num_heads},"
-            + f"num activate heads: {num_activate_heads}, mask sparsity: {mask_sparsity}"
-        )
+    # # Print mask sparsity info
+    # threshold = None if args.binary else args.lambda_threshold
+    # threshold = None if args.scope is not None else threshold
+    # name = ["ff", "attn"]
+    # for n, hooker in zip(name, hookers):
+    #     total_num_heads, num_activate_heads, mask_sparsity = calculate_mask_sparsity(hooker, threshold)
+    #     print(f"model: {args.model}, {n} masking: {args.masking}")
+    #     print(
+    #         f"total num heads: {total_num_heads},"
+    #         + f"num activate heads: {num_activate_heads}, mask sparsity: {mask_sparsity}"
+    #     )
 
     # Prune the model
     pruned_pipe = prune_model(mask_pipe, hookers)
@@ -175,7 +175,7 @@ def create_demo():
                 This demo is currently being developed and may not be fully functional. More models and pruning ratios will be supported soon.
                 The current pruned model checkpoint is not optimal and does not provide the best performance.
                 
-                **Note: Please first initialize the model before generating images.**
+                **Note: Please first initialize the model before generating images. This may take up to 5 minutes as it is on CPU.**
                 """
             )
         with gr.Row():
@@ -183,6 +183,12 @@ def create_demo():
             pruning_ratio = gr.Dropdown(choices=["20%"], value="20%", label="Pruning Ratio", scale=1.2)
             status_label = gr.HighlightedText(label="Model Status", value=[("Model Not Initialized", "red")], scale=1)
             prune_btn = gr.Button("Initialize Original and Pruned Models", variant="primary", scale=1)
+        with gr.Row():
+            gr.Markdown(
+                """
+                **Generate images with the original model and the pruned model. May take up to 1 minute due to dynamic allocation of GPU.**
+                """
+            )
         with gr.Row():
             prompt = gr.Textbox(label="Prompt", value="A clock tower floating in a sea of clouds", scale=3)
             seed = gr.Number(label="Seed", value=44, precision=0, scale=1)