Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import cv2
 import torch
@@ -12,29 +13,26 @@ import os
 os.environ["CUDA_VISIBLE_DEVICES"]="0"
 
 title = "MoMA"
-description = "This model has to run on GPU. By default, we load the model with 4-bit quantization to make it fit in smaller hardwares."
+description = "This model has to run on GPU. By default, we load the model with 4-bit quantization to make it fit in smaller hardware."
+
+seed_everything(0)
+args = parse_args()
+#load MoMA from HuggingFace. Auto download
+model = MoMA_main_modal(args).to(args.device, dtype=torch.float16)
 
 def MoMA_demo(rgb, subject, prompt, strength, seed):
-    seed = int(seed) if seed else 0
-    try:
-        seed = int(seed)
-    except ValueError:
-        seed = 0
-    seed = seed if not seed == 0 else np.random.randint(0,1000)
-    print(f"Seed: {seed}")
-    
     with torch.no_grad():
         generated_image = model.generate_images(rgb, subject, prompt, strength=strength, seed=seed)
     return generated_image
 
+@spaces.GPU
 def inference(rgb, subject, prompt, strength, seed):
+    seed = int(seed) if seed else 0
+    seed = seed if not seed == 0 else np.random.randint(0,1000)
+    
     result = MoMA_demo(rgb, subject, prompt, strength, seed)
     return result
 
-seed_everything(0)
-args = parse_args()
-#load MoMA from HuggingFace. Auto download
-model = MoMA_main_modal(args).to(args.device, dtype=torch.float16)
 
 gr.Interface(
     inference,