Upload app.py

app.py

@@ -62,17 +62,12 @@ class FluxEditor:
         # init all components
         self.t5 = load_t5(self.device, max_length=256 if self.name == "flux-schnell" else 512)
         self.clip = load_clip(self.device)
-        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, device='cuda')
+        self.ae = load_ae(self.name, device='cuda' if self.offload else self.device)
         self.t5.eval()
         self.clip.eval()
         self.ae.eval()
         self.model.eval()
-
-        if self.offload:
-            self.model.cpu()
-            torch.cuda.empty_cache()
-            self.ae.encoder.to(self.device)
     
     @torch.inference_mode()
     @spaces.GPU(duration=30)
@@ -111,11 +106,6 @@ class FluxEditor:
         t0 = time.perf_counter()
 
         opts.seed = None
-        if self.offload:
-            self.ae = self.ae.cpu()
-            torch.cuda.empty_cache()
-            self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
-
         #############inverse#######################
         info = {}
         info['feature'] = {}
@@ -129,12 +119,6 @@ class FluxEditor:
             inp_target = prepare(self.t5, self.clip, init_image, prompt=opts.target_prompt)
         timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(self.name != "flux-schnell"))
 
-        # offload TEs to CPU, load model to gpu
-        if self.offload:
-            self.t5, self.clip = self.t5.cpu(), self.clip.cpu()
-            torch.cuda.empty_cache()
-            self.model = self.model.to(self.device)
-
         # inversion initial noise
         with torch.no_grad():
             z, info = denoise(self.model, **inp, timesteps=timesteps, guidance=1, inverse=True, info=info)
@@ -146,12 +130,6 @@ class FluxEditor:
         # denoise initial noise
         x, _ = denoise(self.model, **inp_target, timesteps=timesteps, guidance=guidance, inverse=False, info=info)
 
-        # offload model, load autoencoder to gpu
-        if self.offload:
-            self.model.cpu()
-            torch.cuda.empty_cache()
-            self.ae.decoder.to(x.device)
-
         # decode latents to pixel space
         x = unpack(x.float(), opts.width, opts.height)