modified: app.py

app.py

@@ -63,8 +63,8 @@ class FluxEditor:
         # init all components
         self.clip = load_clip(self.device)
         self.t5 = load_t5(self.device, max_length=256 if self.name == "flux-schnell" else 512)
-        self.model = load_flow_model(self.name, device="cpu" if self.offload else self.device)
-        self.ae = load_ae(self.name, device="cpu" if self.offload else self.device)
+        self.model = load_flow_model(self.name, self.device)
+        self.ae = load_ae(self.name, self.device)
         self.t5.eval()
         self.clip.eval()
         self.ae.eval()
@@ -144,7 +144,7 @@ class FluxEditor:
             init_image = torch.nn.functional.interpolate(init_image, (height, width))
             if self.offload:
                 self.ae.encoder.to(self.device)
-            init_image = self.ae.encode(init_image.to())
+            init_image = self.ae.encode(init_image)
             if self.offload:
                 self.ae = self.ae.cpu()
                 torch.cuda.empty_cache()