mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-09 04:36:50 -04:00
e76ab16bfe
* Add preview functionality to effmpeg. (#435) * Add preview functionality to effmpeg. effmpeg tool: Preview for actions that have a video output now available. Preview does not work when muxing audio. * Model json unicode fix1 (#443) * fixed Windows 10 path error while loading weights * - fixed TypeError: the JSON object must be str, not 'bytes' with OriginalHighRes Model * MTCNN Extractor and Extraction refactor (#453) * implement mtcnn extractor * mtcnn refactor and vram management changes * cli arguments update for mtcnn/dlib split * Add mtcnn models to gitignore * Change multiprocessing on extract * GUI changes to handle nargs defaults * Early exit bugfix (#455) * Fix extract early termination bug * Fix extract early exit bug * Multi face detection bugfix (#456) * Multi face extraction fix * Original high res cleanup 1 (#457) * slight model re-factoring - removed excess threading code - added random kernel initialization to dense layer * Slight OriginalHighRes re-factoring an code cleanup
107 lines
3.1 KiB
Python
107 lines
3.1 KiB
Python
#!/usr/bin python3
|
|
""" Information on available Nvidia GPUs """
|
|
|
|
import pynvml
|
|
|
|
|
|
class GPUStats(object):
|
|
""" Holds information about system GPU(s) """
|
|
def __init__(self):
|
|
self.verbose = False
|
|
|
|
self.initialized = False
|
|
self.device_count = 0
|
|
self.handles = None
|
|
self.driver = None
|
|
self.devices = None
|
|
self.vram = None
|
|
|
|
self.initialize()
|
|
|
|
if self.device_count == 0:
|
|
return
|
|
|
|
self.driver = self.get_driver()
|
|
self.devices = self.get_devices()
|
|
self.vram = self.get_vram()
|
|
|
|
self.shutdown()
|
|
|
|
def initialize(self):
|
|
""" Initialize pynvml """
|
|
if not self.initialized:
|
|
try:
|
|
pynvml.nvmlInit()
|
|
except pynvml.NVMLError_LibraryNotFound:
|
|
self.initialized = True
|
|
return
|
|
self.initialized = True
|
|
self.get_device_count()
|
|
self.get_handles()
|
|
|
|
def shutdown(self):
|
|
""" Shutdown pynvml """
|
|
if self.initialized:
|
|
self.handles = None
|
|
pynvml.nvmlShutdown()
|
|
self.initialized = False
|
|
|
|
def get_device_count(self):
|
|
""" Return count of Nvidia devices """
|
|
try:
|
|
self.device_count = pynvml.nvmlDeviceGetCount()
|
|
except pynvml.NVMLError:
|
|
self.device_count = 0
|
|
|
|
def get_handles(self):
|
|
""" Return all listed Nvidia handles """
|
|
self.handles = [pynvml.nvmlDeviceGetHandleByIndex(i)
|
|
for i in range(self.device_count)]
|
|
|
|
@staticmethod
|
|
def get_driver():
|
|
""" Get the driver version """
|
|
try:
|
|
driver = pynvml.nvmlSystemGetDriverVersion().decode("utf-8")
|
|
except pynvml.NVMLError:
|
|
driver = "No Nvidia driver found"
|
|
return driver
|
|
|
|
def get_devices(self):
|
|
""" Return total vram in megabytes per device """
|
|
vram = [pynvml.nvmlDeviceGetName(handle).decode("utf-8")
|
|
for handle in self.handles]
|
|
return vram
|
|
|
|
def get_vram(self):
|
|
""" Return total vram in megabytes per device """
|
|
vram = [pynvml.nvmlDeviceGetMemoryInfo(handle).total / (1024 * 1024)
|
|
for handle in self.handles]
|
|
return vram
|
|
|
|
def get_used(self):
|
|
""" Return the vram in use """
|
|
self.initialize()
|
|
vram = [pynvml.nvmlDeviceGetMemoryInfo(handle).used / (1024 * 1024)
|
|
for handle in self.handles]
|
|
self.shutdown()
|
|
|
|
if self.verbose:
|
|
print("GPU VRAM used: {}".format(vram))
|
|
|
|
return vram
|
|
|
|
def get_free(self):
|
|
""" Return the vram available """
|
|
self.initialize()
|
|
vram = [pynvml.nvmlDeviceGetMemoryInfo(handle).free / (1024 * 1024)
|
|
for handle in self.handles]
|
|
self.shutdown()
|
|
return vram
|
|
|
|
def print_info(self):
|
|
""" Output GPU info in verbose mode """
|
|
print("GPU Driver: {}".format(self.driver))
|
|
print("GPU Device count: {}".format(self.device_count))
|
|
print("GPU Devices: {}".format(self.devices))
|
|
print("GPU VRAM: {}".format(self.vram))
|