reduce gpu usage

app.py

@@ -34,7 +34,7 @@ offload = False
 
 name = model_name
 is_schnell = model_name == "flux-schnell"
-resize_longside = 512
+resize_longside = 480
 save = False
 output_dir = 'outputs_gradio'
 
@@ -189,6 +189,7 @@ def edit(brush_canvas, source_prompt, inversion_guidance,
     if training_epochs != 0:
         t5.to('cpu')
         clip.to('cpu')
+        ae.to('cpu')
         torch.set_grad_enabled(True)
         inp_optim["img"] = z0
         _, info, _, _, trainable_noise_list = denoise_with_noise_optim(model,**inp_optim,token_ids=token_ids,source_mask=source_mask,training_steps=1,training_epochs=training_epochs,learning_rate=0.01,seed=seed,noise_scale=noise_scale,timesteps=timesteps,info=info,guidance=denoise_guidance)
@@ -208,7 +209,7 @@ def edit(brush_canvas, source_prompt, inversion_guidance,
     
      # decode latents to pixel space
     batch_x = unpack(x.float(), width,height)
-
+    ae.to(device)
     for x in batch_x:
         x = x.unsqueeze(0)
         
@@ -276,7 +277,7 @@ def create_demo(model_name: str):
 
         🎨 [<b>Examples</b>] Click our examples below, draw your mask and click the "Edit" button. <br>
         
-        🔔 [<b>Note</b>] Due to limited resources in spaces, we will resize image to <=512 longside. <br>
+        🔔 [<b>Note</b>] Due to limited resources in spaces, we will resize image to <=480 longside. <br>
         If you need high resolution for better quality, go to https://github.com/oyly16/LORE for more usage with your own resource. <br>
         """
     article = r"""

flux/sampling_lore.py

@@ -156,7 +156,6 @@ def denoise(
     guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
     
     step_list = []
-    attn_map_list = []
     for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])):
         t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
         info['t'] = t_prev if inverse else t_curr
@@ -200,10 +199,8 @@ def denoise(
 
         # return attnmaps L,1,512,N
         step_list.append(t_curr)
-        attn_map_list.append((attn_maps_mid+attn_maps)/2)
 
-    attn_map_list = torch.stack(attn_map_list)
-    return img, info, step_list, attn_map_list
+    return img, info, step_list, None
 
 selected_layers = range(8,44)
 
@@ -276,7 +273,6 @@ def denoise_with_noise_optim(
     guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
 
     step_list = []
-    attn_map_list = []
     trainable_noise_list = []
     for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])):
         if i >= training_steps:
@@ -353,13 +349,11 @@ def denoise_with_noise_optim(
             optimizer.step()
             print(f"Time {t_curr:.4f} Step {j+1}/{training_epochs}, Loss: {total_loss.item():.6f}")
             
-        attn_map_list.append(attn_maps.detach())
         step_list.append(t_curr)
         trainable_noise = trainable_noise.detach()
         trainable_noise_list.append(trainable_noise.clone())
 
-    attn_map_list = torch.stack(attn_map_list)
-    return img, info, step_list, attn_map_list, trainable_noise_list
+    return img, info, step_list, None, trainable_noise_list
 
 def unpack(x: Tensor, height: int, width: int) -> Tensor:
     return rearrange(