force the device

app.py

@@ -119,6 +119,9 @@ class ImageGenerator:
             max_length=max_length,
             dtype=dtype,
         )
+        self.ae = self.ae.to(device=self.device, dtype=torch.float32)
+        self.dit = self.dit.to(device=self.device, dtype=dtype)
+        self.llm_encoder = self.llm_encoder.to(device=self.device, dtype=dtype)
 
     def prepare(self, prompt, img, ref_image, ref_image_raw):
         bs, _, h, w = img.shape
@@ -314,6 +317,7 @@ class ImageGenerator:
 
         ref_images_raw = self.load_image(ref_images_raw)
         ref_images_raw = ref_images_raw.to(self.device)
+        print(f'self.ae, self.dit device: {self.ae.device}, {self.dit.device}')
         ref_images = self.ae.encode(ref_images_raw.to(self.device) * 2 - 1)
 
         seed = int(seed)
@@ -398,7 +402,7 @@ def prepare_infer_func():
 
     return image_edit.generate_image
 
-@spaces.GPU
+@spaces.GPU(duration=240)
 def inference(prompt, ref_images, seed, size_level, infer_func=None):
     start_time = time.time()