Update app.py

app.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import gradio as gr
+import spaces
+
 import argparse
 import inspect
 import os
@@ -1616,22 +1618,20 @@ if __name__ == "__main__":
     parser.add_argument('--experiment_name', default="AccDiffusion")
 
     args = parser.parse_args()
-
+    
+    pipe = AccDiffusionSDXLPipeline.from_pretrained(args.model_ckpt, torch_dtype=torch.float16).to("cuda")
+    generator = torch.Generator(device='cuda')
+    
     # GRADIO MODE
 
-    def infer(prompt, progress=gr.Progress(track_tqdm=True)):
+    @spaces.GPU()
+    def infer(prompt, num_inference_steps, guidance_scale, progress=gr.Progress(track_tqdm=True)):
         set_seed(args.seed)
         width,height = list(map(int, args.resolution.split(',')))
-        pipe = AccDiffusionSDXLPipeline.from_pretrained(args.model_ckpt, torch_dtype=torch.float16).to("cuda")
-        generator = torch.Generator(device='cuda')
-        generator = generator.manual_seed(args.seed)
         cross_attention_kwargs = {"edit_type": "visualize",
                                   "n_self_replace": 0.4,
                                   "n_cross_replace": {"default_": 1.0, "confetti": 0.8},
                                   }
-        
-        
-        
         seed = args.seed
         generator = generator.manual_seed(seed)
         
@@ -1644,8 +1644,8 @@ if __name__ == "__main__":
                       view_batch_size=args.view_batch_size, 
                       stride=args.stride,
                       cross_attention_kwargs=cross_attention_kwargs,
-                      num_inference_steps=args.num_inference_steps, 
-                      guidance_scale = 7.5, 
+                      num_inference_steps=num_inference_steps, 
+                      guidance_scale = guidance_scale, 
                       multi_guidance_scale = args.multi_guidance_scale,
                       cosine_scale_1=args.cosine_scale_1, 
                       cosine_scale_2=args.cosine_scale_2, 
@@ -1680,7 +1680,7 @@ if __name__ == "__main__":
                     <img src='https://img.shields.io/badge/Project-Page-blue'>
                 </a>
                 <a href='https://github.com/lzhxmu/AccDiffusion'>
-                    <img src='https://img.shields.io/badge/Code-blue'>
+                    <img src='https://img.shields.io/badge/Code-github-blue'>
                 </a>
                 <a href='https://arxiv.org/abs/2407.10738v1'>
                     <img src='https://img.shields.io/badge/Paper-Arxiv-red'>
@@ -1688,6 +1688,9 @@ if __name__ == "__main__":
             </div>
             """)
             prompt = gr.Textbox(label="Prompt")
+            with gr.Accordion("Advanced settings", open=False):
+                num_inference_steps = gr.Slider(label="Inference Steps", minimum=2, maximum=50, step=1, value=50)
+                guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=510, step=0.1, value=7.5)
             submit_btn = gr.Button("Submit")
             output_images = gr.Image(format="png")
             gr.Examples(
@@ -1700,7 +1703,7 @@ if __name__ == "__main__":
             )
         submit_btn.click(
             fn = infer,
-            inputs = [prompt],
+            inputs = [prompt, num_inference_steps, guidance_scale],
             outputs = [output_images],
             show_api=False
         )