Update app.py

app.py

@@ -61,7 +61,7 @@ def load_models(
     dit_path=None,
     ae_path=None,
     qwen2vl_model_path=None,
-    device="cuda",
+    device="cpu",
     max_length=256,
     dtype=torch.bfloat16,
 ):
@@ -118,7 +118,7 @@ class ImageGenerator:
         dit_path=None,
         ae_path=None,
         qwen2vl_model_path=None,
-        device="cuda",
+        device="cpu",
         max_length=640,
         dtype=torch.bfloat16,
     ) -> None:
@@ -135,9 +135,9 @@ class ImageGenerator:
         self.llm_encoder = self.llm_encoder.to(device=self.device, dtype=dtype)
     
     def to_cuda(self):
-        self.ae.to(device='cuda', dtype=torch.float32)
-        self.dit.to(device='cuda', dtype=torch.bfloat16)
-        self.llm_encoder.to(device='cuda', dtype=torch.bfloat16)
+        self.ae.to(device='cpu', dtype=torch.float32)
+        self.dit.to(device='cpu', dtype=torch.bfloat16)
+        self.llm_encoder.to(device='cpu', dtype=torch.bfloat16)
 
     def prepare(self, prompt, img, ref_image, ref_image_raw):
         bs, _, h, w = img.shape