modified: app.py

app.py

@@ -10,7 +10,7 @@ from einops import rearrange
 #from fire import Fire
 from PIL import ExifTags, Image
 from safetensors.torch import load_file, save_file
-
+import spaces
 
 import torch
 import torch.nn.functional as F
@@ -116,7 +116,7 @@ class FluxEditor:
         return img, gr_gallery
 
         
-
+    @spaces.GPU(duration=120)
     @torch.inference_mode()
     def generate_image(
         self,
@@ -239,6 +239,7 @@ class FluxEditor:
         return img, history_gallery
 
 
+    @spaces.GPU(duration=120)
     @torch.inference_mode()
     def edit(self, init_image, source_prompt, target_prompt, editing_strategy, denoise_strategy, num_steps, guidance, attn_guidance_start_block, inject_step, init_image_2=None):
         
@@ -531,16 +532,16 @@ def create_demo(model_name: str, device: str = "cuda" if torch.cuda.is_available
     return demo
 
 
-if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser(description="Flux")
-    parser.add_argument("--name", type=str, default="flux-dev", choices=list(configs.keys()), help="Model name")
-    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device to use")
-    parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use")
-    parser.add_argument("--share", action="store_true", help="Create a public link to your demo")
-    parser.add_argument("--port", type=int, default=9090)
-    args = parser.parse_args()
 
-    demo = create_demo(args.name, args.device, args.offload)
-    #demo.launch(server_name='0.0.0.0', share=args.share, server_port=args.port)
-    demo.launch(share=True)
+import argparse
+parser = argparse.ArgumentParser(description="Flux")
+parser.add_argument("--name", type=str, default="flux-dev", choices=list(configs.keys()), help="Model name")
+parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device to use")
+parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use")
+parser.add_argument("--share", action="store_true", help="Create a public link to your demo")
+parser.add_argument("--port", type=int, default=9090)
+args = parser.parse_args()
+
+demo = create_demo(args.name, args.device, args.offload)
+#demo.launch(server_name='0.0.0.0', share=args.share, server_port=args.port)
+demo.launch(share=True)