Update model.py

ldm/modules/diffusionmodules/model.py

@@ -210,25 +210,29 @@ def forward(self, x):
  h_ = torch.zeros_like(k, device=q.device)
 
  device_type = 'mps' if q.device.type == 'mps' else 'cuda'
-
- if device_type == 'mps':
- mem_free_total = psutil.virtual_memory().available
- else:
+ if device_type == 'cuda':
  stats = torch.cuda.memory_stats(q.device)
  mem_active = stats['active_bytes.all.current']
  mem_reserved = stats['reserved_bytes.all.current']
  mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
  mem_free_torch = mem_reserved - mem_active
  mem_free_total = mem_free_cuda + mem_free_torch
 
- tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * 4
- mem_required = tensor_size * 2.5
- steps = 1
+  tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * 4
+  mem_required = tensor_size * 2.5
+  steps = 1
 
- if mem_required > mem_free_total:
- steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
+ if mem_required > mem_free_total:
+ steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
+
+ slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
 
- slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
+ else:
+ if psutil.virtual_memory().available / (1024**3) < 12:
+ slice_size = 1
+ else:
+ slice_size = min(q.shape[1], math.floor(2**30 / (q.shape[0] * q.shape[1])))
+
  for i in range(0, q.shape[1], slice_size):
  end = i + slice_size