From 8707f88c07e50c35a6ad536f3e2d230856e51a97 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Mon, 14 Nov 2022 20:35:47 +0530
Subject: [PATCH] Show mem free info

---
 ui/index.html                  |  2 +-
 ui/media/js/main.js            | 11 ++++++++---
 ui/sd_internal/task_manager.py | 21 ++++++++++++++++++---
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/ui/index.html b/ui/index.html
index 19857c5c..39d7b32c 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -330,7 +330,7 @@
 <script src="media/js/inpainting-editor.js?v=1"></script>
 <script src="media/js/image-modifiers.js?v=6"></script>
 <script src="media/js/auto-save.js?v=8"></script>
-<script src="media/js/main.js?v=18"></script>
+<script src="media/js/main.js?v=19"></script>
 <script src="media/js/themes.js?v=4"></script>
 <script src="media/js/dnd.js?v=9"></script>
 <script>
diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index 6bc50574..f0a24712 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -190,10 +190,15 @@ function playSound() {
     }
 }
 function setSystemInfo(devices) {
-    let cpu = devices.all.cpu
+    let cpu = devices.all.cpu.name
     let allGPUs = Object.keys(devices.all).filter(d => d != 'cpu')
     let activeGPUs = Object.keys(devices.active)
-    const ID_TO_TEXT = d => `${devices.all[d]} (${d})`
+
+    function ID_TO_TEXT(d) {
+        let info = devices.all[d]
+        return `${info.name} (${d}) (${info.mem_free.toFixed(1)}Gb free / ${info.mem_total.toFixed(1)} Gb total)`
+    }
+
     allGPUs = allGPUs.map(ID_TO_TEXT)
     activeGPUs = activeGPUs.map(ID_TO_TEXT)
 
@@ -1422,7 +1427,7 @@ async function getDevices() {
 
             useGPUsField.innerHTML = ''
             allDeviceIds.forEach(device => {
-                let deviceName = res['all'][device]
+                let deviceName = res['all'][device]['name']
                 let deviceOption = `<option value="${device}">${deviceName} (${device})</option>`
                 useGPUsField.insertAdjacentHTML('beforeend', deviceOption)
             })
diff --git a/ui/sd_internal/task_manager.py b/ui/sd_internal/task_manager.py
index 650fa8e8..a1b70418 100644
--- a/ui/sd_internal/task_manager.py
+++ b/ui/sd_internal/task_manager.py
@@ -356,6 +356,20 @@ def get_devices():
         'active': {},
     }
 
+    def get_device_info(device):
+        if device == 'cpu':
+            return {'name': device_manager.get_processor_name()}
+            
+        mem_free, mem_total = torch.cuda.mem_get_info(device)
+        mem_free /= float(10**9)
+        mem_total /= float(10**9)
+
+        return {
+            'name': torch.cuda.get_device_name(device),
+            'mem_free': mem_free,
+            'mem_total': mem_total,
+        }
+
     # list the compatible devices
     gpu_count = torch.cuda.device_count()
     for device in range(gpu_count):
@@ -363,9 +377,9 @@ def get_devices():
         if not device_manager.is_device_compatible(device):
             continue
 
-        devices['all'].update({device: torch.cuda.get_device_name(device)})
+        devices['all'].update({device: get_device_info(device)})
 
-    devices['all'].update({'cpu': device_manager.get_processor_name()})
+    devices['all'].update({'cpu': get_device_info('cpu')})
 
     # list the activated devices
     if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('get_devices' + ERR_LOCK_FAILED)
@@ -376,7 +390,8 @@ def get_devices():
             weak_data = weak_thread_data.get(rthread)
             if not weak_data or not 'device' in weak_data or not 'device_name' in weak_data:
                 continue
-            devices['active'].update({weak_data['device']: weak_data['device_name']})
+            device = weak_data['device']
+            devices['active'].update({device: get_device_info(device)})
     finally:
         manager_lock.release()