usage

app.py

@@ -6,7 +6,7 @@ import math
 from pathlib import Path
 import sys
 sys.path.insert(1, './taming-transformers')
-# from IPython import display
+from IPython import display
 from base64 import b64encode
 from omegaconf import OmegaConf
 from PIL import Image
@@ -24,8 +24,14 @@ import imageio
 from PIL import ImageFile, Image
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 import gradio as gr
-torch.hub.download_url_to_file('https://i.imgur.com/WEHmKef.jpg', 'gpu.jpg')
+import nvidia_smi
+
+nvidia_smi.nvmlInit()
+handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
+# card id 0 hardcoded here, there is also a call to get all available card ids, so we could iterate
 
+
+torch.hub.download_url_to_file('https://i.imgur.com/WEHmKef.jpg', 'gpu.jpg')
 def sinc(x):
     return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))
 def lanczos(x, a):
@@ -290,8 +296,10 @@ def inference(text, seed, step_size, max_iterations, width, height):
         losses_str = ', '.join(f'{loss.item():g}' for loss in losses)
         tqdm.write(f'i: {i}, loss: {sum(losses).item():g}, losses: {losses_str}')
         out = synth(z)
-        # TF.to_pil_image(out[0].cpu()).save('progress.png')
-        # display.display(display.Image('progress.png'))
+        TF.to_pil_image(out[0].cpu()).save('progress.png')
+        display.display(display.Image('progress.png'))
+        res = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)
+        print(f'gpu: {res.gpu}%, gpu-mem: {res.memory}%')
     def ascend_txt():
         # global i
         out = synth(z)
@@ -338,11 +346,6 @@ def load_image( infilename ) :
     data = np.asarray( img, dtype="int32" )
     return data
 def throttled_inference(text, seed, step_size, max_iterations, width, height):
-    t = torch.cuda.get_device_properties(0).total_memory
-    r = torch.cuda.memory_reserved(0)
-    a = torch.cuda.memory_allocated(0)
-    f = r-a  # free inside reserved
-    print(f)
     global inferences_running
     current = inferences_running
     if current >= 2: