diff --git a/converters/ConvertMasked.py b/converters/ConvertMasked.py index 9460a18..6a5b274 100644 --- a/converters/ConvertMasked.py +++ b/converters/ConvertMasked.py @@ -71,17 +71,17 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i if cfg.face_type == FaceType.FULL: FAN_dst_face_mask_a_0 = cv2.resize (dst_face_fanseg_mask, (output_size,output_size), cv2.INTER_CUBIC) - elif cfg.face_type == FaceType.HALF: - half_face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanseg_input_size, face_type=FaceType.HALF) + else: + face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanseg_input_size, face_type=cfg.face_type) fanseg_rect_corner_pts = np.array ( [ [0,0], [cfg.fanseg_input_size-1,0], [0,cfg.fanseg_input_size-1] ], dtype=np.float32 ) - a = LandmarksProcessor.transform_points (fanseg_rect_corner_pts, half_face_fanseg_mat, invert=True ) + a = LandmarksProcessor.transform_points (fanseg_rect_corner_pts, face_fanseg_mat, invert=True ) b = LandmarksProcessor.transform_points (a, full_face_fanseg_mat ) m = cv2.getAffineTransform(b, fanseg_rect_corner_pts) FAN_dst_face_mask_a_0 = cv2.warpAffine(dst_face_fanseg_mask, m, (cfg.fanseg_input_size,)*2, flags=cv2.INTER_CUBIC ) FAN_dst_face_mask_a_0 = cv2.resize (FAN_dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC) - else: - raise ValueError ("cfg.face_type unsupported") + #else: + # raise ValueError ("cfg.face_type unsupported") if cfg.mask_mode == 3: #FAN-prd prd_face_mask_a_0 = FAN_prd_face_mask_a_0 diff --git a/converters/ConverterConfig.py b/converters/ConverterConfig.py index 327a242..9f1e174 100644 --- a/converters/ConverterConfig.py +++ b/converters/ConverterConfig.py @@ -117,8 +117,8 @@ class ConverterConfigMasked(ConverterConfig): super().__init__(type=ConverterConfig.TYPE_MASKED) self.face_type = face_type - if self.face_type not in [FaceType.FULL, FaceType.HALF]: - raise ValueError("ConverterConfigMasked supports only full or half face masks.") + if self.face_type not in [FaceType.HALF, FaceType.MID_FULL, FaceType.FULL ]: + raise ValueError("ConverterConfigMasked does not support this type of face.") self.default_mode = default_mode self.clip_hborder_mask_per = clip_hborder_mask_per diff --git a/facelib/FaceType.py b/facelib/FaceType.py index a5e02d0..edd9b80 100644 --- a/facelib/FaceType.py +++ b/facelib/FaceType.py @@ -2,11 +2,12 @@ from enum import IntEnum class FaceType(IntEnum): #enumerating in order "next contains prev" - HALF = 0, - FULL = 1, - FULL_NO_ALIGN = 3, - HEAD = 4, - HEAD_NO_ALIGN = 5, + HALF = 0 + MID_FULL = 1 + FULL = 2 + FULL_NO_ALIGN = 3 + HEAD = 4 + HEAD_NO_ALIGN = 5 MARK_ONLY = 10, #no align at all, just embedded faceinfo @@ -22,6 +23,7 @@ class FaceType(IntEnum): return to_string_dict[face_type] from_string_dict = {'half_face': FaceType.HALF, + 'midfull_face': FaceType.MID_FULL, 'full_face': FaceType.FULL, 'head' : FaceType.HEAD, 'mark_only' : FaceType.MARK_ONLY, @@ -29,6 +31,7 @@ from_string_dict = {'half_face': FaceType.HALF, 'head_no_align' : FaceType.HEAD_NO_ALIGN, } to_string_dict = { FaceType.HALF : 'half_face', + FaceType.MID_FULL : 'midfull_face', FaceType.FULL : 'full_face', FaceType.HEAD : 'head', FaceType.MARK_ONLY :'mark_only', diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index 671129a..a87e6eb 100644 --- a/facelib/LandmarksProcessor.py +++ b/facelib/LandmarksProcessor.py @@ -271,6 +271,8 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0): if face_type == FaceType.HALF: padding = 0 + elif face_type == FaceType.MID_FULL: + padding = int(output_size * 0.06) elif face_type == FaceType.FULL: padding = (output_size / 64) * 12 elif face_type == FaceType.HEAD: @@ -435,9 +437,6 @@ def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0): ) ) - - #import code - #code.interact(local=dict(globals(), **locals())) eyes_fall_dist = w // 32 eyes_thickness = max( w // 64, 1 ) diff --git a/models/Model_SAEv2/Model.py b/models/Model_SAEv2/Model.py new file mode 100644 index 0000000..1fc5e78 --- /dev/null +++ b/models/Model_SAEv2/Model.py @@ -0,0 +1,658 @@ +from functools import partial + +import numpy as np + +import mathlib +from facelib import FaceType +from interact import interact as io +from models import ModelBase +from nnlib import nnlib +from samplelib import * + + +#SAE - Styled AutoEncoder +class SAEModel(ModelBase): + + #override + def onInitializeOptions(self, is_first_run, ask_override): + yn_str = {True:'y',False:'n'} + + default_resolution = 128 + default_archi = 'df' + default_face_type = 'f' + default_learn_mask = True + + if is_first_run: + resolution = io.input_int("Resolution ( 64-256 ?:help skip:128) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.") + resolution = np.clip (resolution, 64, 256) + while np.modf(resolution / 16)[0] != 0.0: + resolution -= 1 + self.options['resolution'] = resolution + + self.options['face_type'] = io.input_str ("Half, mid full, or full face? (h/mf/f, ?:help skip:f) : ", default_face_type, ['h','mf','f'], help_message="Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face.").lower() + self.options['learn_mask'] = io.input_bool ( f"Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]} ) : " , default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result.") + else: + self.options['resolution'] = self.options.get('resolution', default_resolution) + self.options['face_type'] = self.options.get('face_type', default_face_type) + self.options['learn_mask'] = self.options.get('learn_mask', default_learn_mask) + + if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend: + def_optimizer_mode = self.options.get('optimizer_mode', 1) + self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.") + else: + self.options['optimizer_mode'] = self.options.get('optimizer_mode', 1) + + if is_first_run: + self.options['archi'] = io.input_str ("AE architecture (df, liae ?:help skip:%s) : " % (default_archi) , default_archi, ['df','liae'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes.").lower() #-s version is slower, but has decreased change to collapse. + else: + self.options['archi'] = self.options.get('archi', default_archi) + + default_ae_dims = 256 if 'liae' in self.options['archi'] else 512 + default_e_ch_dims = 21 + default_d_ch_dims = default_e_ch_dims + def_ca_weights = False + + if is_first_run: + self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 ) + self.options['e_ch_dims'] = np.clip ( io.input_int("Encoder dims per channel (21-85 ?:help skip:%d) : " % (default_e_ch_dims) , default_e_ch_dims, help_message="More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 ) + default_d_ch_dims = self.options['e_ch_dims'] + self.options['d_ch_dims'] = np.clip ( io.input_int("Decoder dims per channel (10-85 ?:help skip:%d) : " % (default_d_ch_dims) , default_d_ch_dims, help_message="More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU." ), 10, 85 ) + self.options['ca_weights'] = io.input_bool (f"Use CA weights? (y/n, ?:help skip:{yn_str[def_ca_weights]} ) : ", def_ca_weights, help_message="Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes a time at first run.") + else: + self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims) + self.options['e_ch_dims'] = self.options.get('e_ch_dims', default_e_ch_dims) + self.options['d_ch_dims'] = self.options.get('d_ch_dims', default_d_ch_dims) + self.options['ca_weights'] = self.options.get('ca_weights', def_ca_weights) + + default_face_style_power = 0.0 + default_bg_style_power = 0.0 + if is_first_run or ask_override: + def_pixel_loss = self.options.get('pixel_loss', False) + self.options['pixel_loss'] = io.input_bool (f"Use pixel loss? (y/n, ?:help skip:{yn_str[def_pixel_loss]} ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time. Enabling this option too early increases the chance of model collapse.") + + default_face_style_power = default_face_style_power if is_first_run else self.options.get('face_style_power', default_face_style_power) + self.options['face_style_power'] = np.clip ( io.input_number("Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_face_style_power), default_face_style_power, + help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."), 0.0, 100.0 ) + + default_bg_style_power = default_bg_style_power if is_first_run else self.options.get('bg_style_power', default_bg_style_power) + self.options['bg_style_power'] = np.clip ( io.input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power, + help_message="Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."), 0.0, 100.0 ) + + default_apply_random_ct = False if is_first_run else self.options.get('apply_random_ct', False) + self.options['apply_random_ct'] = io.input_bool (f"Apply random color transfer to src faceset? (y/n, ?:help skip:{yn_str[default_apply_random_ct]}) : ", default_apply_random_ct, help_message="Increase variativity of src samples by apply LCT color transfer from random dst samples. It is like 'face_style' learning, but more precise color transfer and without risk of model collapse, also it does not require additional GPU resources, but the training time may be longer, due to the src faceset is becoming more diverse.") + + default_true_face_training = False if is_first_run else self.options.get('true_face_training', False) + self.options['true_face_training'] = io.input_bool (f"Enable 'true face' training? (y/n, ?:help skip:{yn_str[default_true_face_training]}) : ", default_true_face_training, help_message="Result face will be more like src. Enable it only after 100k iters, when face is sharp enough.") + + + if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301 + default_clipgrad = False if is_first_run else self.options.get('clipgrad', False) + self.options['clipgrad'] = io.input_bool (f"Enable gradient clipping? (y/n, ?:help skip:{yn_str[default_clipgrad]}) : ", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.") + else: + self.options['clipgrad'] = False + + else: + self.options['pixel_loss'] = self.options.get('pixel_loss', False) + self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power) + self.options['bg_style_power'] = self.options.get('bg_style_power', default_bg_style_power) + self.options['apply_random_ct'] = self.options.get('apply_random_ct', False) + self.options['clipgrad'] = self.options.get('clipgrad', False) + + if is_first_run: + self.options['pretrain'] = io.input_bool ("Pretrain the model? (y/n, ?:help skip:n) : ", False, help_message="Pretrain the model with large amount of various faces. This technique may help to train the fake with overly different face shapes and light conditions of src/dst data. Face will be look more like a morphed. To reduce the morph effect, some model files will be initialized but not be updated after pretrain: LIAE: inter_AB.h5 DF: encoder.h5. The longer you pretrain the model the more morphed face will look. After that, save and run the training again.") + else: + self.options['pretrain'] = False + + #override + def onInitialize(self): + exec(nnlib.import_all(), locals(), globals()) + self.set_vram_batch_requirements({1.5:4,4:8}) + + resolution = self.options['resolution'] + learn_mask = self.options['learn_mask'] + + ae_dims = self.options['ae_dims'] + e_ch_dims = self.options['e_ch_dims'] + d_ch_dims = self.options['d_ch_dims'] + self.pretrain = self.options['pretrain'] = self.options.get('pretrain', False) + if not self.pretrain: + self.options.pop('pretrain') + + bgr_shape = (resolution, resolution, 3) + mask_shape = (resolution, resolution, 1) + + apply_random_ct = self.options.get('apply_random_ct', False) + self.true_face_training = self.options.get('true_face_training', False) + masked_training = True + + class SAEDFModel(object): + def __init__(self, resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask): + super().__init__() + self.learn_mask = learn_mask + + output_nc = 3 + bgr_shape = (resolution, resolution, output_nc) + mask_shape = (resolution, resolution, 1) + lowest_dense_res = resolution // 16 + e_dims = output_nc*e_ch_dims + + def downscale (dim, kernel_size=5, dilation_rate=1): + def func(x): + return SubpixelDownscaler()(LeakyReLU(0.1)(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x))) + return func + + def upscale (dim): + def func(x): + return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same')(x))) + return func + + def enc_flow(e_dims, ae_dims, lowest_dense_res): + def func(inp): + x = downscale(e_dims , 3, 1 )(inp) + x = downscale(e_dims*2, 3, 1 )(x) + x = downscale(e_dims*4, 3, 1 )(x) + x0 = downscale(e_dims*8, 3, 1 )(x) + + x = downscale(e_dims , 5, 1 )(inp) + x = downscale(e_dims*2, 5, 1 )(x) + x = downscale(e_dims*4, 5, 1 )(x) + x1 = downscale(e_dims*8, 5, 1 )(x) + + x = downscale(e_dims , 5, 2 )(inp) + x = downscale(e_dims*2, 5, 2 )(x) + x = downscale(e_dims*4, 5, 2 )(x) + x2 = downscale(e_dims*8, 5, 2 )(x) + + x = downscale(e_dims , 7, 2 )(inp) + x = downscale(e_dims*2, 7, 2 )(x) + x = downscale(e_dims*4, 7, 2 )(x) + x3 = downscale(e_dims*8, 7, 2 )(x) + + x = Concatenate()([x0,x1,x2,x3]) + + x = Dense(ae_dims)(Flatten()(x)) + x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x) + x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x) + x = upscale(ae_dims)(x) + return x + return func + + def dec_flow(output_nc, d_ch_dims, add_residual_blocks=True): + dims = output_nc * d_ch_dims + def ResidualBlock(dim): + def func(inp): + x = Conv2D(dim, kernel_size=3, padding='same')(inp) + x = LeakyReLU(0.2)(x) + x = Conv2D(dim, kernel_size=3, padding='same')(x) + x = Add()([x, inp]) + x = LeakyReLU(0.2)(x) + return x + return func + + def func(x): + x = upscale(dims*8)(x) + + if add_residual_blocks: + x = ResidualBlock(dims*8)(x) + + x = upscale(dims*4)(x) + + if add_residual_blocks: + x = ResidualBlock(dims*4)(x) + + x = upscale(dims*2)(x) + + if add_residual_blocks: + x = ResidualBlock(dims*2)(x) + + return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x) + return func + + self.encoder = modelify(enc_flow(e_dims, ae_dims, lowest_dense_res)) ( Input(bgr_shape) ) + + sh = K.int_shape( self.encoder.outputs[0] )[1:] + self.decoder_src = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) ) + self.decoder_dst = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) ) + + if learn_mask: + self.decoder_srcm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) ) + self.decoder_dstm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) ) + + self.src_dst_trainable_weights = self.encoder.trainable_weights + self.decoder_src.trainable_weights + self.decoder_dst.trainable_weights + + if learn_mask: + self.src_dst_mask_trainable_weights = self.encoder.trainable_weights + self.decoder_srcm.trainable_weights + self.decoder_dstm.trainable_weights + + self.warped_src, self.warped_dst = Input(bgr_shape), Input(bgr_shape) + self.target_src, self.target_dst = Input(bgr_shape), Input(bgr_shape) + self.target_srcm, self.target_dstm = Input(mask_shape), Input(mask_shape) + self.src_code, self.dst_code = self.encoder(self.warped_src), self.encoder(self.warped_dst) + + self.pred_src_src = self.decoder_src(self.src_code) + self.pred_dst_dst = self.decoder_dst(self.dst_code) + self.pred_src_dst = self.decoder_src(self.dst_code) + + if learn_mask: + self.pred_src_srcm = self.decoder_srcm(self.src_code) + self.pred_dst_dstm = self.decoder_dstm(self.dst_code) + self.pred_src_dstm = self.decoder_srcm(self.dst_code) + + def get_model_filename_list(self, exclude_for_pretrain=False): + ar = [] + if not exclude_for_pretrain: + ar += [ [self.encoder, 'encoder.h5'] ] + ar += [ [self.decoder_src, 'decoder_src.h5'], + [self.decoder_dst, 'decoder_dst.h5'] ] + if self.learn_mask: + ar += [ [self.decoder_srcm, 'decoder_srcm.h5'], + [self.decoder_dstm, 'decoder_dstm.h5'] ] + return ar + + class SAELIAEModel(object): + def __init__(self, resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask): + super().__init__() + self.learn_mask = learn_mask + + output_nc = 3 + bgr_shape = (resolution, resolution, output_nc) + mask_shape = (resolution, resolution, 1) + + e_dims = output_nc*e_ch_dims + + lowest_dense_res = resolution // 16 + + def downscale (dim, kernel_size=5, dilation_rate=1): + def func(x): + return SubpixelDownscaler()(LeakyReLU(0.1)(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x))) + return func + + def upscale (dim): + def func(x): + return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same')(x))) + return func + + def enc_flow(e_dims): + def func(inp): + x = downscale(e_dims , 3, 1 )(inp) + x = downscale(e_dims*2, 3, 1 )(x) + x = downscale(e_dims*4, 3, 1 )(x) + x0 = downscale(e_dims*8, 3, 1 )(x) + + x = downscale(e_dims , 5, 1 )(inp) + x = downscale(e_dims*2, 5, 1 )(x) + x = downscale(e_dims*4, 5, 1 )(x) + x1 = downscale(e_dims*8, 5, 1 )(x) + + x = downscale(e_dims , 5, 2 )(inp) + x = downscale(e_dims*2, 5, 2 )(x) + x = downscale(e_dims*4, 5, 2 )(x) + x2 = downscale(e_dims*8, 5, 2 )(x) + + x = downscale(e_dims , 7, 2 )(inp) + x = downscale(e_dims*2, 7, 2 )(x) + x = downscale(e_dims*4, 7, 2 )(x) + x3 = downscale(e_dims*8, 7, 2 )(x) + + x = Concatenate()([x0,x1,x2,x3]) + + x = Flatten()(x) + return x + return func + + def inter_flow(lowest_dense_res, ae_dims): + def func(x): + x = Dense(ae_dims)(x) + x = Dense(lowest_dense_res * lowest_dense_res * ae_dims*2)(x) + x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims*2))(x) + x = upscale(ae_dims*2)(x) + return x + return func + + def dec_flow(output_nc, d_ch_dims, add_residual_blocks=True): + d_dims = output_nc*d_ch_dims + def ResidualBlock(dim): + def func(inp): + x = Conv2D(dim, kernel_size=3, padding='same')(inp) + x = LeakyReLU(0.2)(x) + x = Conv2D(dim, kernel_size=3, padding='same')(inp) + x = Add()([x, inp]) + x = LeakyReLU(0.2)(x) + return x + return func + + def func(x): + x = upscale(d_dims*8)(x) + + if add_residual_blocks: + x = ResidualBlock(d_dims*8)(x) + + x = upscale(d_dims*4)(x) + + if add_residual_blocks: + x = ResidualBlock(d_dims*4)(x) + + x = upscale(d_dims*2)(x) + + if add_residual_blocks: + x = ResidualBlock(d_dims*2)(x) + + return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x) + return func + + self.encoder = modelify(enc_flow(e_dims)) ( Input(bgr_shape) ) + + sh = K.int_shape( self.encoder.outputs[0] )[1:] + self.inter_B = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) ) + self.inter_AB = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) ) + + sh = np.array(K.int_shape( self.inter_B.outputs[0] )[1:])*(1,1,2) + self.decoder = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) ) + + if learn_mask: + self.decoderm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) ) + + self.src_dst_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights + + if learn_mask: + self.src_dst_mask_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoderm.trainable_weights + + self.warped_src, self.warped_dst = Input(bgr_shape), Input(bgr_shape) + self.target_src, self.target_dst = Input(bgr_shape), Input(bgr_shape) + self.target_srcm, self.target_dstm = Input(mask_shape), Input(mask_shape) + + warped_src_code = self.encoder (self.warped_src) + self.src_code = warped_src_inter_AB_code = self.inter_AB (warped_src_code) + src_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code]) + + warped_dst_code = self.encoder (self.warped_dst) + self.dst_code = warped_dst_inter_B_code = self.inter_B (warped_dst_code) + warped_dst_inter_AB_code = self.inter_AB (warped_dst_code) + dst_code = Concatenate()([warped_dst_inter_B_code,warped_dst_inter_AB_code]) + + src_dst_code = Concatenate()([warped_dst_inter_AB_code,warped_dst_inter_AB_code]) + + self.pred_src_src = self.decoder(src_code) + self.pred_dst_dst = self.decoder(dst_code) + self.pred_src_dst = self.decoder(src_dst_code) + + if learn_mask: + self.pred_src_srcm = self.decoderm(src_code) + self.pred_dst_dstm = self.decoderm(dst_code) + self.pred_src_dstm = self.decoderm(src_dst_code) + + def get_model_filename_list(self, exclude_for_pretrain=False): + ar = [ [self.encoder, 'encoder.h5'], + [self.inter_B, 'inter_B.h5'] ] + + if not exclude_for_pretrain: + ar += [ [self.inter_AB, 'inter_AB.h5'] ] + + ar += [ [self.decoder, 'decoder.h5'] ] + + if self.learn_mask: + ar += [ [self.decoderm, 'decoderm.h5'] ] + + return ar + + if 'df' in self.options['archi']: + self.model = SAEDFModel (resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask) + elif 'liae' in self.options['archi']: + self.model = SAELIAEModel (resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask) + + self.opt_dis_model = [] + + if self.true_face_training: + def dis_flow(ndf=256): + def func(x): + x, = x + + code_res = K.int_shape(x)[1] + + x = Conv2D( ndf, 4, strides=2, padding='valid')( ZeroPadding2D(1)(x) ) + x = LeakyReLU(0.1)(x) + + x = Conv2D( ndf*2, 3, strides=2, padding='valid')( ZeroPadding2D(1)(x) ) + x = LeakyReLU(0.1)(x) + + if code_res > 8: + x = Conv2D( ndf*4, 3, strides=2, padding='valid')( ZeroPadding2D(1)(x) ) + x = LeakyReLU(0.1)(x) + + if code_res > 16: + x = Conv2D( ndf*8, 3, strides=2, padding='valid')( ZeroPadding2D(1)(x) ) + x = LeakyReLU(0.1)(x) + + if code_res > 32: + x = Conv2D( ndf*8, 3, strides=2, padding='valid')( ZeroPadding2D(1)(x) ) + x = LeakyReLU(0.1)(x) + + return Conv2D( 1, 1, strides=1, padding='valid', activation='sigmoid')(x) + return func + + sh = [ Input( K.int_shape(self.model.src_code)[1:] ) ] + self.dis = modelify(dis_flow()) (sh) + + self.opt_dis_model = [ (self.dis, 'dis.h5') ] + + loaded, not_loaded = [], self.model.get_model_filename_list()+self.opt_dis_model + if not self.is_first_run(): + loaded, not_loaded = self.load_weights_safe(not_loaded) + + CA_models = [] + if self.options.get('ca_weights', False): + CA_models += [ model for model, _ in not_loaded ] + + CA_conv_weights_list = [] + for model in CA_models: + for layer in model.layers: + if type(layer) == keras.layers.Conv2D: + CA_conv_weights_list += [layer.weights[0]] #- is Conv2D kernel_weights + + if len(CA_conv_weights_list) != 0: + CAInitializerMP ( CA_conv_weights_list ) + + target_srcm = gaussian_blur( max(1, resolution // 32) )(self.model.target_srcm) + target_dstm = gaussian_blur( max(1, resolution // 32) )(self.model.target_dstm) + + target_src_masked = self.model.target_src*target_srcm + target_dst_masked = self.model.target_dst*target_dstm + target_dst_anti_masked = self.model.target_dst*(1.0 - target_dstm) + + target_src_masked_opt = target_src_masked if masked_training else self.model.target_src + target_dst_masked_opt = target_dst_masked if masked_training else self.model.target_dst + + pred_src_src_masked_opt = self.model.pred_src_src*target_srcm if masked_training else self.model.pred_src_src + pred_dst_dst_masked_opt = self.model.pred_dst_dst*target_dstm if masked_training else self.model.pred_dst_dst + + psd_target_dst_masked = self.model.pred_src_dst*target_dstm + psd_target_dst_anti_masked = self.model.pred_src_dst*(1.0 - target_dstm) + + if self.is_training_mode: + self.src_dst_opt = RMSprop(lr=5e-5, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1) + self.src_dst_mask_opt = RMSprop(lr=5e-5, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1) + self.D_opt = RMSprop(lr=1e-5, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1) + + if not self.options['pixel_loss']: + src_loss = K.mean ( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_opt, pred_src_src_masked_opt) ) + else: + src_loss = K.mean ( 50*K.square( target_src_masked_opt - pred_src_src_masked_opt ) ) + + face_style_power = self.options['face_style_power'] / 100.0 + if face_style_power != 0: + src_loss += style_loss(gaussian_blur_radius=resolution//16, loss_weight=face_style_power, wnd_size=0)( psd_target_dst_masked, target_dst_masked ) + + bg_style_power = self.options['bg_style_power'] / 100.0 + if bg_style_power != 0: + if not self.options['pixel_loss']: + src_loss += K.mean( (10*bg_style_power)*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( psd_target_dst_anti_masked, target_dst_anti_masked )) + else: + src_loss += K.mean( (50*bg_style_power)*K.square( psd_target_dst_anti_masked - target_dst_anti_masked )) + + if not self.options['pixel_loss']: + dst_loss = K.mean( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)(target_dst_masked_opt, pred_dst_dst_masked_opt) ) + else: + dst_loss = K.mean( 50*K.square( target_dst_masked_opt - pred_dst_dst_masked_opt ) ) + + opt_D_loss = [] + if self.true_face_training: + def DLoss(labels,logits): + return K.mean(K.binary_crossentropy(labels,logits)) + + src_code_d = self.dis( self.model.src_code ) + src_code_d_ones = K.ones_like(src_code_d) + src_code_d_zeros = K.zeros_like(src_code_d) + dst_code_d = self.dis( self.model.dst_code ) + dst_code_d_ones = K.ones_like(dst_code_d) + opt_D_loss = [ 0.2*DLoss(src_code_d_ones, src_code_d) ] + + loss_D = (DLoss(dst_code_d_ones , dst_code_d) + \ + DLoss(src_code_d_zeros, src_code_d) ) * 0.5 + + self.D_train = K.function ([self.model.warped_src, self.model.warped_dst],[loss_D], self.D_opt.get_updates(loss_D, self.dis.trainable_weights) ) + + self.src_dst_train = K.function ([self.model.warped_src, self.model.warped_dst, self.model.target_src, self.model.target_srcm, self.model.target_dst, self.model.target_dstm], + [src_loss,dst_loss], + self.src_dst_opt.get_updates( [src_loss+dst_loss]+opt_D_loss, self.model.src_dst_trainable_weights) + ) + + if self.options['learn_mask']: + src_mask_loss = K.mean(K.square(self.model.target_srcm-self.model.pred_src_srcm)) + dst_mask_loss = K.mean(K.square(self.model.target_dstm-self.model.pred_dst_dstm)) + self.src_dst_mask_train = K.function ([self.model.warped_src, self.model.warped_dst, self.model.target_srcm, self.model.target_dstm],[src_mask_loss, dst_mask_loss], self.src_dst_mask_opt.get_updates(src_mask_loss+dst_mask_loss, self.model.src_dst_mask_trainable_weights ) ) + + if self.options['learn_mask']: + self.AE_view = K.function ([self.model.warped_src, self.model.warped_dst], [self.model.pred_src_src, self.model.pred_dst_dst, self.model.pred_dst_dstm, self.model.pred_src_dst, self.model.pred_src_dstm]) + else: + self.AE_view = K.function ([self.model.warped_src, self.model.warped_dst], [self.model.pred_src_src, self.model.pred_dst_dst, self.model.pred_src_dst ]) + + else: + if self.options['learn_mask']: + self.AE_convert = K.function ([self.model.warped_dst],[ self.model.pred_src_dst, self.model.pred_dst_dstm, self.model.pred_src_dstm ]) + else: + self.AE_convert = K.function ([self.model.warped_dst],[ self.model.pred_src_dst ]) + + + if self.is_training_mode: + t = SampleProcessor.Types + + if self.options['face_type'] == 'h': + face_type = t.FACE_TYPE_HALF + elif self.options['face_type'] == 'mf': + face_type = t.FACE_TYPE_MID_FULL + elif self.options['face_type'] == 'f': + face_type = t.FACE_TYPE_FULL + + t_mode_bgr = t.MODE_BGR if not self.pretrain else t.MODE_BGR_SHUFFLE + + training_data_src_path = self.training_data_src_path + training_data_dst_path = self.training_data_dst_path + sort_by_yaw = self.sort_by_yaw + + if self.pretrain and self.pretraining_data_path is not None: + training_data_src_path = self.pretraining_data_path + training_data_dst_path = self.pretraining_data_path + sort_by_yaw = False + + self.set_training_data_generators ([ + SampleGeneratorFace(training_data_src_path, sort_by_yaw_target_samples_path=training_data_dst_path if sort_by_yaw else None, + random_ct_samples_path=training_data_dst_path if apply_random_ct else None, + debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), + output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution, 'apply_ct': apply_random_ct}, + {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution, 'apply_ct': apply_random_ct }, + {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution } ] + ), + + SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ), + output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution}, + {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution}, + {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution} ]) + ]) + + #override + def get_model_filename_list(self): + return self.model.get_model_filename_list ( exclude_for_pretrain=(self.pretrain and self.iter != 0) ) +self.opt_dis_model + + #override + def onSave(self): + self.save_weights_safe( self.get_model_filename_list()+self.opt_dis_model ) + + #override + def onTrainOneIter(self, generators_samples, generators_list): + warped_src, target_src, target_srcm = generators_samples[0] + warped_dst, target_dst, target_dstm = generators_samples[1] + + feed = [warped_src, warped_dst, target_src, target_srcm, target_dst, target_dstm] + + src_loss, dst_loss, = self.src_dst_train (feed) + + if self.true_face_training: + self.D_train([warped_src, warped_dst]) + + if self.options['learn_mask']: + feed = [ warped_src, warped_dst, target_srcm, target_dstm ] + src_mask_loss, dst_mask_loss, = self.src_dst_mask_train (feed) + + return ( ('src_loss', src_loss), ('dst_loss', dst_loss), ) + + #override + def onGetPreview(self, sample): + test_S = sample[0][1][0:4] #first 4 samples + test_S_m = sample[0][2][0:4] #first 4 samples + test_D = sample[1][1][0:4] + test_D_m = sample[1][2][0:4] + + if self.options['learn_mask']: + S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ] + DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ] + else: + S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ] + + result = [] + st = [] + for i in range(len(test_S)): + ar = S[i], SS[i], D[i], DD[i], SD[i] + + st.append ( np.concatenate ( ar, axis=1) ) + + result += [ ('SAE', np.concatenate (st, axis=0 )), ] + + if self.options['learn_mask']: + st_m = [] + for i in range(len(test_S)): + ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) + st_m.append ( np.concatenate ( ar, axis=1) ) + + result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ] + + return result + + def predictor_func (self, face=None, dummy_predict=False): + if dummy_predict: + self.AE_convert ([ np.zeros ( (1, self.options['resolution'], self.options['resolution'], 3), dtype=np.float32 ) ]) + else: + if self.options['learn_mask']: + bgr, mask_dst_dstm, mask_src_dstm = self.AE_convert ([face[np.newaxis,...]]) + mask = mask_dst_dstm[0] * mask_src_dstm[0] + return bgr[0], mask[...,0] + else: + bgr, = self.AE_convert ([face[np.newaxis,...]]) + return bgr[0] + + #override + def get_ConverterConfig(self): + if self.options['face_type'] == 'h': + face_type = FaceType.HALF + elif self.options['face_type'] == 'mf': + face_type = FaceType.MID_FULL + elif self.options['face_type'] == 'f': + face_type = FaceType.FULL + + import converters + return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type, + default_mode = 1 if self.options['apply_random_ct'] or self.options['face_style_power'] or self.options['bg_style_power'] else 4, + clip_hborder_mask_per=0.0625 if (face_type == FaceType.FULL) else 0, + ) + +Model = SAEModel diff --git a/models/Model_TrueFace/__init__.py b/models/Model_SAEv2/__init__.py similarity index 100% rename from models/Model_TrueFace/__init__.py rename to models/Model_SAEv2/__init__.py diff --git a/models/Model_TrueFace/Model.py b/models/Model_TrueFace/Model.py deleted file mode 100644 index 66678dc..0000000 --- a/models/Model_TrueFace/Model.py +++ /dev/null @@ -1,185 +0,0 @@ -import numpy as np - -from facelib import FaceType -from interact import interact as io -from models import ModelBase -from nnlib import nnlib, FUNIT -from samplelib import * - -class TrueFaceModel(ModelBase): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs, - ask_sort_by_yaw=False, - ask_random_flip=False, - ask_src_scale_mod=False) - - #override - def onInitializeOptions(self, is_first_run, ask_override): - default_resolution = 128 - default_face_type = 'f' - - if is_first_run: - resolution = self.options['resolution'] = io.input_int(f"Resolution ( 64-256 ?:help skip:{default_resolution}) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.") - resolution = np.clip (resolution, 64, 256) - while np.modf(resolution / 16)[0] != 0.0: - resolution -= 1 - else: - self.options['resolution'] = self.options.get('resolution', default_resolution) - - if is_first_run: - self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="").lower() - else: - self.options['face_type'] = self.options.get('face_type', default_face_type) - - if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend: - def_optimizer_mode = self.options.get('optimizer_mode', 3) - self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.") - else: - self.options['optimizer_mode'] = self.options.get('optimizer_mode', 1) - - #override - def onInitialize(self, batch_size=-1, **in_options): - exec(nnlib.code_import_all, locals(), globals()) - self.set_vram_batch_requirements({2:1,3:1,4:4,5:8,6:16}) - - resolution = self.options['resolution'] - face_type = self.face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF - - self.model = FUNIT( face_type_str=FaceType.toString(face_type), - batch_size=self.batch_size, - encoder_nf=64, - encoder_downs=2, - encoder_res_blk=2, - class_downs=4, - class_nf=64, - class_latent=64, - mlp_blks=2, - dis_nf=64, - dis_res_blks=10, - num_classes=2, - subpixel_decoder=True, - initialize_weights=self.is_first_run(), - is_training=self.is_training_mode, - tf_cpu_mode=self.options['optimizer_mode']-1 - ) - - if not self.is_first_run(): - self.load_weights_safe(self.model.get_model_filename_list()) - - t = SampleProcessor.Types - face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF - if self.is_training_mode: - - output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True}, - ] - - self.set_training_data_generators ([ - SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=True), - output_sample_types=output_sample_types ), - - SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=True), - output_sample_types=output_sample_types ) - ]) - else: - generator = SampleGeneratorFace(self.training_data_src_path, batch_size=1, - sample_process_options=SampleProcessor.Options(), - output_sample_types=[ {'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True} ] ) - - io.log_info("Calculating average src face style...") - codes = [] - for i in io.progress_bar_generator(range(generator.get_total_sample_count())): - codes += self.model.get_average_class_code( generator.generate_next() ) - - self.average_class_code = np.mean ( np.array(codes), axis=0 )[None,...] - - - #override - def get_model_filename_list(self): - return self.model.get_model_filename_list() - - #override - def onSave(self): - self.save_weights_safe(self.model.get_model_filename_list()) - - #override - def onTrainOneIter(self, generators_samples, generators_list): - bs = self.batch_size - lbs = bs // 2 - hbs = bs - lbs - - src, = generators_samples[0] - dst, = generators_samples[1] - - xa = np.concatenate ( [src[0:lbs], dst[0:lbs]], axis=0 ) - - la = np.concatenate ( [ np.array ([0]*lbs, np.int32), - np.array ([1]*lbs, np.int32) ] ) - - xb = np.concatenate ( [src[lbs:], dst[lbs:]], axis=0 ) - - lb = np.concatenate ( [ np.array ([0]*hbs, np.int32), - np.array ([1]*hbs, np.int32) ] ) - - rnd_list = np.arange(lbs*2) - np.random.shuffle(rnd_list) - xa = xa[rnd_list,...] - la = la[rnd_list,...] - la = la[...,None] - - rnd_list = np.arange(hbs*2) - np.random.shuffle(rnd_list) - xb = xb[rnd_list,...] - lb = lb[rnd_list,...] - lb = lb[...,None] - - G_loss, D_loss = self.model.train(xa,la,xb,lb) - - return ( ('G_loss', G_loss), ('D_loss', D_loss), ) - - #override - def onGetPreview(self, generators_samples): - xa = generators_samples[0][0] - xb = generators_samples[1][0] - - view_samples = min(4, xa.shape[0]) - - - s_xa_mean = self.model.get_average_class_code([xa])[0][None,...] - s_xb_mean = self.model.get_average_class_code([xb])[0][None,...] - - s_xab_mean = self.model.get_average_class_code([ np.concatenate( [xa,xb], axis=0) ])[0][None,...] - - lines = [] - - for i in range(view_samples): - xaxa, = self.model.convert ([ xa[i:i+1], s_xa_mean ] ) - xbxb, = self.model.convert ([ xb[i:i+1], s_xb_mean ] ) - xbxa, = self.model.convert ([ xb[i:i+1], s_xa_mean ] ) - - xa_i,xb_i,xaxa,xbxb,xbxa = [ np.clip(x/2+0.5, 0, 1) for x in [xa[i], xb[i], xaxa[0],xbxb[0],xbxa[0]] ] - - lines += [ np.concatenate( (xa_i, xaxa, xb_i, xbxb, xbxa), axis=1) ] - - r = np.concatenate ( lines, axis=0 ) - return [ ('TrueFace', r ) ] - - def predictor_func (self, face=None, dummy_predict=False): - if dummy_predict: - self.model.convert ([ np.zeros ( (1, self.options['resolution'], self.options['resolution'], 3), dtype=np.float32 ), self.average_class_code ]) - else: - bgr, = self.model.convert ([ face[np.newaxis,...]*2-1, self.average_class_code ]) - return bgr[0] / 2 + 0.5 - - #override - def get_ConverterConfig(self): - - import converters - return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=self.face_type, - default_mode = 1, - clip_hborder_mask_per=0.0625 if (self.face_type == FaceType.FULL) else 0, - ) - -Model = TrueFaceModel diff --git a/models/archived_models.zip b/models/archived_models.zip index 6e15f87..13e0198 100644 Binary files a/models/archived_models.zip and b/models/archived_models.zip differ diff --git a/nnlib/device.py b/nnlib/device.py index 144de43..82b2883 100644 --- a/nnlib/device.py +++ b/nnlib/device.py @@ -1,7 +1,8 @@ +import sys +import ctypes import os import json import numpy as np -from .pynvml import * #you can set DFL_TF_MIN_REQ_CAP manually for your build #the reason why we cannot check tensorflow.version is it requires import tensorflow @@ -88,13 +89,8 @@ class device: for i in range(plaidML_devices_count): yield i elif device.backend == "tensorflow": - for gpu_idx in range(nvmlDeviceGetCount()): - cap = device.getDeviceComputeCapability (gpu_idx) - if cap >= tf_min_req_cap: - yield gpu_idx - elif device.backend == "tensorflow-generic": - yield 0 - + for dev in cuda_devices: + yield dev['index'] @staticmethod def getValidDevicesWithAtLeastTotalMemoryGB(totalmemsize_gb): @@ -104,35 +100,20 @@ class device: if plaidML_devices[i]['globalMemSize'] >= totalmemsize_gb*1024*1024*1024: result.append (i) elif device.backend == "tensorflow": - for i in device.getValidDeviceIdxsEnumerator(): - handle = nvmlDeviceGetHandleByIndex(i) - memInfo = nvmlDeviceGetMemoryInfo( handle ) - if (memInfo.total) >= totalmemsize_gb*1024*1024*1024: + for dev in cuda_devices: + if dev['total_mem'] >= totalmemsize_gb*1024*1024*1024: result.append (i) - elif device.backend == "tensorflow-generic": - return [0] return result - @staticmethod - def getAllDevicesIdxsList(): - if device.backend == "plaidML": - return [ *range(plaidML_devices_count) ] - elif device.backend == "tensorflow": - return [ *range(nvmlDeviceGetCount() ) ] - elif device.backend == "tensorflow-generic": - return [0] - @staticmethod def getValidDevicesIdxsWithNamesList(): if device.backend == "plaidML": return [ (i, plaidML_devices[i]['description'] ) for i in device.getValidDeviceIdxsEnumerator() ] elif device.backend == "tensorflow": - return [ (i, nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)).decode() ) for i in device.getValidDeviceIdxsEnumerator() ] + return [ ( dev['index'], dev['name'] ) for dev in cuda_devices ] elif device.backend == "tensorflow-cpu": return [ (0, 'CPU') ] - elif device.backend == "tensorflow-generic": - return [ (0, device.getDeviceName(0) ) ] @staticmethod def getDeviceVRAMTotalGb (idx): @@ -140,13 +121,10 @@ class device: if idx < plaidML_devices_count: return plaidML_devices[idx]['globalMemSize'] / (1024*1024*1024) elif device.backend == "tensorflow": - if idx < nvmlDeviceGetCount(): - memInfo = nvmlDeviceGetMemoryInfo( nvmlDeviceGetHandleByIndex(idx) ) - return round ( memInfo.total / (1024*1024*1024) ) - + for dev in cuda_devices: + if idx == dev['index']: + return round ( dev['total_mem'] / (1024*1024*1024) ) return 0 - elif device.backend == "tensorflow-generic": - return 2 @staticmethod def getBestValidDeviceIdx(): @@ -163,15 +141,12 @@ class device: elif device.backend == "tensorflow": idx = -1 idx_mem = 0 - for i in device.getValidDeviceIdxsEnumerator(): - memInfo = nvmlDeviceGetMemoryInfo( nvmlDeviceGetHandleByIndex(i) ) - if memInfo.total > idx_mem: - idx = i - idx_mem = memInfo.total + for dev in cuda_devices: + if dev['total_mem'] > idx_mem: + idx = dev['index'] + idx_mem = dev['total_mem'] return idx - elif device.backend == "tensorflow-generic": - return 0 @staticmethod def getWorstValidDeviceIdx(): @@ -188,24 +163,22 @@ class device: elif device.backend == "tensorflow": idx = -1 idx_mem = sys.maxsize - for i in device.getValidDeviceIdxsEnumerator(): - memInfo = nvmlDeviceGetMemoryInfo( nvmlDeviceGetHandleByIndex(i) ) - if memInfo.total < idx_mem: - idx = i - idx_mem = memInfo.total + for dev in cuda_devices: + if dev['total_mem'] < idx_mem: + idx = dev['index'] + idx_mem = dev['total_mem'] return idx - elif device.backend == "tensorflow-generic": - return 0 @staticmethod def isValidDeviceIdx(idx): if device.backend == "plaidML": return idx in [*device.getValidDeviceIdxsEnumerator()] elif device.backend == "tensorflow": - return idx in [*device.getValidDeviceIdxsEnumerator()] - elif device.backend == "tensorflow-generic": - return (idx == 0) + for dev in cuda_devices: + if idx == dev['index']: + return True + return False @staticmethod def getDeviceIdxsEqualModel(idx): @@ -219,14 +192,13 @@ class device: return result elif device.backend == "tensorflow": result = [] - idx_name = nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode() - for i in device.getValidDeviceIdxsEnumerator(): - if nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)).decode() == idx_name: - result.append (i) + idx_name = device.getDeviceName(idx) + for dev in cuda_devices: + if dev['name'] == idx_name: + result.append ( dev['index'] ) + return result - elif device.backend == "tensorflow-generic": - return [0] if idx == 0 else [] @staticmethod def getDeviceName (idx): @@ -234,11 +206,9 @@ class device: if idx < plaidML_devices_count: return plaidML_devices[idx]['description'] elif device.backend == "tensorflow": - if idx < nvmlDeviceGetCount(): - return nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode() - elif device.backend == "tensorflow-generic": - if idx == 0: - return "Generic GeForce GPU" + for dev in cuda_devices: + if dev['index'] == idx: + return dev['name'] return None @@ -252,35 +222,22 @@ class device: @staticmethod def getDeviceComputeCapability(idx): - result = 0 if device.backend == "plaidML": return 99 elif device.backend == "tensorflow": - if idx < nvmlDeviceGetCount(): - result = nvmlDeviceGetCudaComputeCapability(nvmlDeviceGetHandleByIndex(idx)) - elif device.backend == "tensorflow-generic": - return 99 if idx == 0 else 0 - - return result[0] * 10 + result[1] - - -force_plaidML = os.environ.get("DFL_FORCE_PLAIDML", "0") == "1" #for OpenCL build , forcing using plaidML even if NVIDIA found -force_tf_cpu = os.environ.get("DFL_FORCE_TF_CPU", "0") == "1" #for OpenCL build , forcing using tf-cpu if plaidML failed -has_nvml = False -has_nvml_cap = False - -#use DFL_FORCE_HAS_NVIDIA_DEVICE=1 if -#- your NVIDIA cannot be seen by OpenCL -#- CUDA build of DFL -has_nvidia_device = os.environ.get("DFL_FORCE_HAS_NVIDIA_DEVICE", "0") == "1" + for dev in cuda_devices: + if dev['index'] == idx: + return dev['cc'] + return 0 +plaidML_build = os.environ.get("DFL_PLAIDML_BUILD", "0") == "1" plaidML_devices = None -def get_plaidML_devices(): - global plaidML_devices - global has_nvidia_device +cuda_devices = None + +if plaidML_build: if plaidML_devices is None: plaidML_devices = [] - # Using plaidML OpenCL backend to determine system devices and has_nvidia_device + # Using plaidML OpenCL backend to determine system devices try: os.environ['PLAIDML_EXPERIMENTAL'] = 'false' #this enables work plaidML without run 'plaidml-setup' import plaidml @@ -289,8 +246,6 @@ def get_plaidML_devices(): details = json.loads(d.details) if details['type'] == 'CPU': #skipping opencl-CPU continue - if 'nvidia' in details['vendor'].lower(): - has_nvidia_device = True plaidML_devices += [ {'id':d.id, 'globalMemSize' : int(details['globalMemSize']), 'description' : d.description.decode() @@ -298,60 +253,58 @@ def get_plaidML_devices(): ctx.shutdown() except: pass - return plaidML_devices - -if not has_nvidia_device: - get_plaidML_devices() - -#choosing backend - -if device.backend is None and not force_tf_cpu: - #first trying to load NVSMI and detect CUDA devices for tensorflow backend, - #even force_plaidML is choosed, because if plaidML will fail, we can choose tensorflow - try: - nvmlInit() - has_nvml = True - device.backend = "tensorflow" #set tensorflow backend in order to use device.*device() functions - - gpu_idxs = device.getAllDevicesIdxsList() - gpu_caps = np.array ( [ device.getDeviceComputeCapability(gpu_idx) for gpu_idx in gpu_idxs ] ) - - if len ( np.ndarray.flatten ( np.argwhere (gpu_caps >= tf_min_req_cap) ) ) == 0: - if not force_plaidML: - print ("No CUDA devices found with minimum required compute capability: %d.%d. Falling back to OpenCL mode." % (tf_min_req_cap // 10, tf_min_req_cap % 10) ) - device.backend = None - nvmlShutdown() - else: - has_nvml_cap = True - except: - #if no NVSMI installed exception will occur - device.backend = None - has_nvml = False - -if force_plaidML or (device.backend is None and not has_nvidia_device): - #tensorflow backend was failed without has_nvidia_device , or forcing plaidML, trying to use plaidML backend - if len(get_plaidML_devices()) == 0: - #print ("plaidML: No capable OpenCL devices found. Falling back to tensorflow backend.") - device.backend = None - else: + + if len(plaidML_devices) != 0: device.backend = "plaidML" - plaidML_devices_count = len(get_plaidML_devices()) +else: + if cuda_devices is None: + cuda_devices = [] + libnames = ('libcuda.so', 'libcuda.dylib', 'nvcuda.dll') + cuda = None + for libname in libnames: + try: + cuda = ctypes.CDLL(libname) + except: + continue + else: + break + + if cuda is not None: + nGpus = ctypes.c_int() + name = b' ' * 200 + cc_major = ctypes.c_int() + cc_minor = ctypes.c_int() + freeMem = ctypes.c_size_t() + totalMem = ctypes.c_size_t() + + result = ctypes.c_int() + device_t = ctypes.c_int() + context = ctypes.c_void_p() + error_str = ctypes.c_char_p() + + if cuda.cuInit(0) == 0 and \ + cuda.cuDeviceGetCount(ctypes.byref(nGpus)) == 0: + for i in range(nGpus.value): + if cuda.cuDeviceGet(ctypes.byref(device_t), i) != 0 or \ + cuda.cuDeviceGetName(ctypes.c_char_p(name), len(name), device_t) != 0 or \ + cuda.cuDeviceComputeCapability(ctypes.byref(cc_major), ctypes.byref(cc_minor), device_t) != 0: + continue + + if cuda.cuCtxCreate_v2(ctypes.byref(context), 0, device_t) == 0: + if cuda.cuMemGetInfo_v2(ctypes.byref(freeMem), ctypes.byref(totalMem)) == 0: + cc = cc_major.value * 10 + cc_minor.value + if cc >= tf_min_req_cap: + cuda_devices.append ( {'index':i, + 'name':name.split(b'\0', 1)[0].decode(), + 'total_mem':totalMem.value, + 'free_mem':freeMem.value, + 'cc':cc + } + ) + cuda.cuCtxDetach(context) + + if len(cuda_devices) != 0: + device.backend = "tensorflow" if device.backend is None: - if force_tf_cpu: - device.backend = "tensorflow-cpu" - elif not has_nvml: - if has_nvidia_device: - #some notebook systems have NVIDIA card without NVSMI in official drivers - #in that case considering we have system with one capable GPU and let tensorflow to choose best GPU - device.backend = "tensorflow-generic" - else: - #no NVSMI and no NVIDIA cards, also plaidML was failed, then CPU only - device.backend = "tensorflow-cpu" - else: - if has_nvml_cap: - #has NVSMI and capable CUDA-devices, but force_plaidML was failed, then we choosing tensorflow - device.backend = "tensorflow" - else: - #has NVSMI, no capable CUDA-devices, also plaidML was failed, then CPU only - device.backend = "tensorflow-cpu" + device.backend = "tensorflow-cpu" diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index f2bd9ab..30f4894 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -96,6 +96,7 @@ dssim = nnlib.dssim PixelShuffler = nnlib.PixelShuffler SubpixelUpscaler = nnlib.SubpixelUpscaler +SubpixelDownscaler = nnlib.SubpixelDownscaler Scale = nnlib.Scale BlurPool = nnlib.BlurPool FUNITAdain = nnlib.FUNITAdain @@ -156,13 +157,13 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator else: config = tf.ConfigProto() - if device_config.force_gpu_idx != -1 and device_config.backend != "tensorflow-generic": - #tensorflow-generic is system with NVIDIA card, but w/o NVSMI - #so dont hide devices and let tensorflow to choose best card - visible_device_list = '' - for idx in device_config.gpu_idxs: - visible_device_list += str(idx) + ',' - config.gpu_options.visible_device_list=visible_device_list[:-1] + #if device_config.force_gpu_idx != -1 and device_config.backend != "tensorflow-generic": + # #tensorflow-generic is system with NVIDIA card, but w/o NVSMI + # #so dont hide devices and let tensorflow to choose best card + # visible_device_list = '' + # for idx in device_config.gpu_idxs: + # visible_device_list += str(idx) + ',' + # config.gpu_options.visible_device_list=visible_device_list[:-1] config.gpu_options.force_gpu_compatible = True config.gpu_options.allow_growth = device_config.allow_growth @@ -472,7 +473,50 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator nnlib.PixelShuffler = PixelShuffler nnlib.SubpixelUpscaler = PixelShuffler + + class SubpixelDownscaler(KL.Layer): + def __init__(self, size=(2, 2), data_format='channels_last', **kwargs): + super(SubpixelDownscaler, self).__init__(**kwargs) + self.data_format = data_format + self.size = size + def call(self, inputs): + + input_shape = K.shape(inputs) + if K.int_shape(input_shape)[0] != 4: + raise ValueError('Inputs should have rank 4; Received input shape:', str(K.int_shape(inputs))) + + batch_size, h, w, c = input_shape[0], input_shape[1], input_shape[2], K.int_shape(inputs)[-1] + rh, rw = self.size + oh, ow = h // rh, w // rw + oc = c * (rh * rw) + + out = K.reshape(inputs, (batch_size, oh, rh, ow, rw, c)) + out = K.permute_dimensions(out, (0, 1, 3, 2, 4, 5)) + out = K.reshape(out, (batch_size, oh, ow, oc)) + return out + + def compute_output_shape(self, input_shape): + if len(input_shape) != 4: + raise ValueError('Inputs should have rank ' + + str(4) + + '; Received input shape:', str(input_shape)) + + height = input_shape[1] // self.size[0] if input_shape[1] is not None else None + width = input_shape[2] // self.size[1] if input_shape[2] is not None else None + channels = input_shape[3] * self.size[0] * self.size[1] + + return (input_shape[0], height, width, channels) + + def get_config(self): + config = {'size': self.size, + 'data_format': self.data_format} + base_config = super(SubpixelDownscaler, self).get_config() + + return dict(list(base_config.items()) + list(config.items())) + + nnlib.SubpixelDownscaler = SubpixelDownscaler + class BlurPool(KL.Layer): """ https://arxiv.org/abs/1904.11486 https://github.com/adobe/antialiased-cnns diff --git a/nnlib/pynvml.py b/nnlib/pynvml.py deleted file mode 100644 index 5cc5a50..0000000 --- a/nnlib/pynvml.py +++ /dev/null @@ -1,1727 +0,0 @@ -##### -# Copyright (c) 2011-2015, NVIDIA Corporation. All rights reserved. -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# * Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# * Redistributions in binary form must reproduce the above copyright -# notice, this list of conditions and the following disclaimer in the -# documentation and/or other materials provided with the distribution. -# * Neither the name of the NVIDIA Corporation nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF -# THE POSSIBILITY OF SUCH DAMAGE. -##### - -## -# Python bindings for the NVML library -## -from ctypes import * -from ctypes.util import find_library -import sys -import os -import threading -import string - -## C Type mappings ## -## Enums -_nvmlEnableState_t = c_uint -NVML_FEATURE_DISABLED = 0 -NVML_FEATURE_ENABLED = 1 - -_nvmlBrandType_t = c_uint -NVML_BRAND_UNKNOWN = 0 -NVML_BRAND_QUADRO = 1 -NVML_BRAND_TESLA = 2 -NVML_BRAND_NVS = 3 -NVML_BRAND_GRID = 4 -NVML_BRAND_GEFORCE = 5 -NVML_BRAND_COUNT = 6 - -_nvmlTemperatureThresholds_t = c_uint -NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 -NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 -NVML_TEMPERATURE_THRESHOLD_COUNT = 1 - -_nvmlTemperatureSensors_t = c_uint -NVML_TEMPERATURE_GPU = 0 -NVML_TEMPERATURE_COUNT = 1 - -_nvmlComputeMode_t = c_uint -NVML_COMPUTEMODE_DEFAULT = 0 -NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 -NVML_COMPUTEMODE_PROHIBITED = 2 -NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3 -NVML_COMPUTEMODE_COUNT = 4 - -_nvmlMemoryLocation_t = c_uint -NVML_MEMORY_LOCATION_L1_CACHE = 0 -NVML_MEMORY_LOCATION_L2_CACHE = 1 -NVML_MEMORY_LOCATION_DEVICE_MEMORY = 2 -NVML_MEMORY_LOCATION_REGISTER_FILE = 3 -NVML_MEMORY_LOCATION_TEXTURE_MEMORY = 4 -NVML_MEMORY_LOCATION_COUNT = 5 - -# These are deprecated, instead use _nvmlMemoryErrorType_t -_nvmlEccBitType_t = c_uint -NVML_SINGLE_BIT_ECC = 0 -NVML_DOUBLE_BIT_ECC = 1 -NVML_ECC_ERROR_TYPE_COUNT = 2 - -_nvmlEccCounterType_t = c_uint -NVML_VOLATILE_ECC = 0 -NVML_AGGREGATE_ECC = 1 -NVML_ECC_COUNTER_TYPE_COUNT = 2 - -_nvmlMemoryErrorType_t = c_uint -NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 -NVML_MEMORY_ERROR_TYPE_UNCORRECTED = 1 -NVML_MEMORY_ERROR_TYPE_COUNT = 2 - -_nvmlClockType_t = c_uint -NVML_CLOCK_GRAPHICS = 0 -NVML_CLOCK_SM = 1 -NVML_CLOCK_MEM = 2 -NVML_CLOCK_COUNT = 3 - -_nvmlDriverModel_t = c_uint -NVML_DRIVER_WDDM = 0 -NVML_DRIVER_WDM = 1 - -_nvmlPstates_t = c_uint -NVML_PSTATE_0 = 0 -NVML_PSTATE_1 = 1 -NVML_PSTATE_2 = 2 -NVML_PSTATE_3 = 3 -NVML_PSTATE_4 = 4 -NVML_PSTATE_5 = 5 -NVML_PSTATE_6 = 6 -NVML_PSTATE_7 = 7 -NVML_PSTATE_8 = 8 -NVML_PSTATE_9 = 9 -NVML_PSTATE_10 = 10 -NVML_PSTATE_11 = 11 -NVML_PSTATE_12 = 12 -NVML_PSTATE_13 = 13 -NVML_PSTATE_14 = 14 -NVML_PSTATE_15 = 15 -NVML_PSTATE_UNKNOWN = 32 - -_nvmlInforomObject_t = c_uint -NVML_INFOROM_OEM = 0 -NVML_INFOROM_ECC = 1 -NVML_INFOROM_POWER = 2 -NVML_INFOROM_COUNT = 3 - -_nvmlReturn_t = c_uint -NVML_SUCCESS = 0 -NVML_ERROR_UNINITIALIZED = 1 -NVML_ERROR_INVALID_ARGUMENT = 2 -NVML_ERROR_NOT_SUPPORTED = 3 -NVML_ERROR_NO_PERMISSION = 4 -NVML_ERROR_ALREADY_INITIALIZED = 5 -NVML_ERROR_NOT_FOUND = 6 -NVML_ERROR_INSUFFICIENT_SIZE = 7 -NVML_ERROR_INSUFFICIENT_POWER = 8 -NVML_ERROR_DRIVER_NOT_LOADED = 9 -NVML_ERROR_TIMEOUT = 10 -NVML_ERROR_IRQ_ISSUE = 11 -NVML_ERROR_LIBRARY_NOT_FOUND = 12 -NVML_ERROR_FUNCTION_NOT_FOUND = 13 -NVML_ERROR_CORRUPTED_INFOROM = 14 -NVML_ERROR_GPU_IS_LOST = 15 -NVML_ERROR_RESET_REQUIRED = 16 -NVML_ERROR_OPERATING_SYSTEM = 17 -NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 -NVML_ERROR_UNKNOWN = 999 - -_nvmlFanState_t = c_uint -NVML_FAN_NORMAL = 0 -NVML_FAN_FAILED = 1 - -_nvmlLedColor_t = c_uint -NVML_LED_COLOR_GREEN = 0 -NVML_LED_COLOR_AMBER = 1 - -_nvmlGpuOperationMode_t = c_uint -NVML_GOM_ALL_ON = 0 -NVML_GOM_COMPUTE = 1 -NVML_GOM_LOW_DP = 2 - -_nvmlPageRetirementCause_t = c_uint -NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 0 -NVML_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 1 -NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 - -_nvmlRestrictedAPI_t = c_uint -NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 -NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 -NVML_RESTRICTED_API_COUNT = 2 - -_nvmlBridgeChipType_t = c_uint -NVML_BRIDGE_CHIP_PLX = 0 -NVML_BRIDGE_CHIP_BRO4 = 1 -NVML_MAX_PHYSICAL_BRIDGE = 128 - -_nvmlValueType_t = c_uint -NVML_VALUE_TYPE_DOUBLE = 0 -NVML_VALUE_TYPE_UNSIGNED_INT = 1 -NVML_VALUE_TYPE_UNSIGNED_LONG = 2 -NVML_VALUE_TYPE_UNSIGNED_LONG_LONG = 3 -NVML_VALUE_TYPE_COUNT = 4 - -_nvmlPerfPolicyType_t = c_uint -NVML_PERF_POLICY_POWER = 0 -NVML_PERF_POLICY_THERMAL = 1 -NVML_PERF_POLICY_COUNT = 2 - -_nvmlSamplingType_t = c_uint -NVML_TOTAL_POWER_SAMPLES = 0 -NVML_GPU_UTILIZATION_SAMPLES = 1 -NVML_MEMORY_UTILIZATION_SAMPLES = 2 -NVML_ENC_UTILIZATION_SAMPLES = 3 -NVML_DEC_UTILIZATION_SAMPLES = 4 -NVML_PROCESSOR_CLK_SAMPLES = 5 -NVML_MEMORY_CLK_SAMPLES = 6 -NVML_SAMPLINGTYPE_COUNT = 7 - -_nvmlPcieUtilCounter_t = c_uint -NVML_PCIE_UTIL_TX_BYTES = 0 -NVML_PCIE_UTIL_RX_BYTES = 1 -NVML_PCIE_UTIL_COUNT = 2 - -_nvmlGpuTopologyLevel_t = c_uint -NVML_TOPOLOGY_INTERNAL = 0 -NVML_TOPOLOGY_SINGLE = 10 -NVML_TOPOLOGY_MULTIPLE = 20 -NVML_TOPOLOGY_HOSTBRIDGE = 30 -NVML_TOPOLOGY_CPU = 40 -NVML_TOPOLOGY_SYSTEM = 50 - -# C preprocessor defined values -nvmlFlagDefault = 0 -nvmlFlagForce = 1 - -# buffer size -NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 -NVML_DEVICE_UUID_BUFFER_SIZE = 80 -NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 81 -NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 -NVML_DEVICE_NAME_BUFFER_SIZE = 64 -NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 -NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 -NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 16 - -NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1) -NVML_VALUE_NOT_AVAILABLE_uint = c_uint(-1) - -## Lib loading ## -nvmlLib = None -libLoadLock = threading.Lock() -_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown - -## Error Checking ## -class NVMLError(Exception): - _valClassMapping = dict() - # List of currently known error codes - _errcode_to_string = { - NVML_ERROR_UNINITIALIZED: "Uninitialized", - NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", - NVML_ERROR_NOT_SUPPORTED: "Not Supported", - NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", - NVML_ERROR_ALREADY_INITIALIZED: "Already Initialized", - NVML_ERROR_NOT_FOUND: "Not Found", - NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", - NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", - NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", - NVML_ERROR_TIMEOUT: "Timeout", - NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", - NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", - NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", - NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", - NVML_ERROR_GPU_IS_LOST: "GPU is lost", - NVML_ERROR_RESET_REQUIRED: "GPU requires restart", - NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", - NVML_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an NVML/RM version mismatch.", - NVML_ERROR_UNKNOWN: "Unknown Error", - } - def __new__(typ, value): - ''' - Maps value to a proper subclass of NVMLError. - See _extractNVMLErrorsAsClasses function for more details - ''' - if typ == NVMLError: - typ = NVMLError._valClassMapping.get(value, typ) - obj = Exception.__new__(typ) - obj.value = value - return obj - def __str__(self): - try: - if self.value not in NVMLError._errcode_to_string: - NVMLError._errcode_to_string[self.value] = str(nvmlErrorString(self.value)) - return NVMLError._errcode_to_string[self.value] - except NVMLError_Uninitialized: - return "NVML Error with code %d" % self.value - def __eq__(self, other): - return self.value == other.value - -def _extractNVMLErrorsAsClasses(): - ''' - Generates a hierarchy of classes on top of NVMLError class. - - Each NVML Error gets a new NVMLError subclass. This way try,except blocks can filter appropriate - exceptions more easily. - - NVMLError is a parent class. Each NVML_ERROR_* gets it's own subclass. - e.g. NVML_ERROR_ALREADY_INITIALIZED will be turned into NVMLError_AlreadyInitialized - ''' - this_module = sys.modules[__name__] - nvmlErrorsNames = filter(lambda x: x.startswith("NVML_ERROR_"), dir(this_module)) - for err_name in nvmlErrorsNames: - # e.g. Turn NVML_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized - class_name = "NVMLError_" + string.capwords(err_name.replace("NVML_ERROR_", ""), "_").replace("_", "") - err_val = getattr(this_module, err_name) - def gen_new(val): - def new(typ): - obj = NVMLError.__new__(typ, val) - return obj - return new - new_error_class = type(class_name, (NVMLError,), {'__new__': gen_new(err_val)}) - new_error_class.__module__ = __name__ - setattr(this_module, class_name, new_error_class) - NVMLError._valClassMapping[err_val] = new_error_class -_extractNVMLErrorsAsClasses() - -def _nvmlCheckReturn(ret): - if (ret != NVML_SUCCESS): - raise NVMLError(ret) - return ret - -## Function access ## -_nvmlGetFunctionPointer_cache = dict() # function pointers are cached to prevent unnecessary libLoadLock locking -def _nvmlGetFunctionPointer(name): - global nvmlLib - - if name in _nvmlGetFunctionPointer_cache: - return _nvmlGetFunctionPointer_cache[name] - - libLoadLock.acquire() - try: - # ensure library was loaded - if (nvmlLib == None): - raise NVMLError(NVML_ERROR_UNINITIALIZED) - try: - _nvmlGetFunctionPointer_cache[name] = getattr(nvmlLib, name) - return _nvmlGetFunctionPointer_cache[name] - except AttributeError: - raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) - finally: - # lock is always freed - libLoadLock.release() - -## Alternative object -# Allows the object to be printed -# Allows mismatched types to be assigned -# - like None when the Structure variant requires c_uint -class nvmlFriendlyObject(object): - def __init__(self, dictionary): - for x in dictionary: - setattr(self, x, dictionary[x]) - def __str__(self): - return self.__dict__.__str__() - -def nvmlStructToFriendlyObject(struct): - d = {} - for x in struct._fields_: - key = x[0] - value = getattr(struct, key) - d[key] = value - obj = nvmlFriendlyObject(d) - return obj - -# pack the object so it can be passed to the NVML library -def nvmlFriendlyObjectToStruct(obj, model): - for x in model._fields_: - key = x[0] - value = obj.__dict__[key] - setattr(model, key, value) - return model - -## Unit structures -class struct_c_nvmlUnit_t(Structure): - pass # opaque handle -c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t) - -class _PrintableStructure(Structure): - """ - Abstract class that produces nicer __str__ output than ctypes.Structure. - e.g. instead of: - >>> print str(obj) - - this class will print - class_name(field_name: formatted_value, field_name: formatted_value) - - _fmt_ dictionary of -> - e.g. class that has _field_ 'hex_value', c_uint could be formatted with - _fmt_ = {"hex_value" : "%08X"} - to produce nicer output. - Default fomratting string for all fields can be set with key "" like: - _fmt_ = {"" : "%d MHz"} # e.g all values are numbers in MHz. - If not set it's assumed to be just "%s" - - Exact format of returned str from this class is subject to change in the future. - """ - _fmt_ = {} - def __str__(self): - result = [] - for x in self._fields_: - key = x[0] - value = getattr(self, key) - fmt = "%s" - if key in self._fmt_: - fmt = self._fmt_[key] - elif "" in self._fmt_: - fmt = self._fmt_[""] - result.append(("%s: " + fmt) % (key, value)) - return self.__class__.__name__ + "(" + string.join(result, ", ") + ")" - -class c_nvmlUnitInfo_t(_PrintableStructure): - _fields_ = [ - ('name', c_char * 96), - ('id', c_char * 96), - ('serial', c_char * 96), - ('firmwareVersion', c_char * 96), - ] - -class c_nvmlLedState_t(_PrintableStructure): - _fields_ = [ - ('cause', c_char * 256), - ('color', _nvmlLedColor_t), - ] - -class c_nvmlPSUInfo_t(_PrintableStructure): - _fields_ = [ - ('state', c_char * 256), - ('current', c_uint), - ('voltage', c_uint), - ('power', c_uint), - ] - -class c_nvmlUnitFanInfo_t(_PrintableStructure): - _fields_ = [ - ('speed', c_uint), - ('state', _nvmlFanState_t), - ] - -class c_nvmlUnitFanSpeeds_t(_PrintableStructure): - _fields_ = [ - ('fans', c_nvmlUnitFanInfo_t * 24), - ('count', c_uint) - ] - -## Device structures -class struct_c_nvmlDevice_t(Structure): - pass # opaque handle -c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t) - -class nvmlPciInfo_t(_PrintableStructure): - _fields_ = [ - ('busId', c_char * 16), - ('domain', c_uint), - ('bus', c_uint), - ('device', c_uint), - ('pciDeviceId', c_uint), - - # Added in 2.285 - ('pciSubSystemId', c_uint), - ('reserved0', c_uint), - ('reserved1', c_uint), - ('reserved2', c_uint), - ('reserved3', c_uint), - ] - _fmt_ = { - 'domain' : "0x%04X", - 'bus' : "0x%02X", - 'device' : "0x%02X", - 'pciDeviceId' : "0x%08X", - 'pciSubSystemId' : "0x%08X", - } - -class c_nvmlMemory_t(_PrintableStructure): - _fields_ = [ - ('total', c_ulonglong), - ('free', c_ulonglong), - ('used', c_ulonglong), - ] - _fmt_ = {'': "%d B"} - -class c_nvmlBAR1Memory_t(_PrintableStructure): - _fields_ = [ - ('bar1Total', c_ulonglong), - ('bar1Free', c_ulonglong), - ('bar1Used', c_ulonglong), - ] - _fmt_ = {'': "%d B"} - -# On Windows with the WDDM driver, usedGpuMemory is reported as None -# Code that processes this structure should check for None, I.E. -# -# if (info.usedGpuMemory == None): -# # TODO handle the error -# pass -# else: -# print("Using %d MiB of memory" % (info.usedGpuMemory / 1024 / 1024)) -# -# See NVML documentation for more information -class c_nvmlProcessInfo_t(_PrintableStructure): - _fields_ = [ - ('pid', c_uint), - ('usedGpuMemory', c_ulonglong), - ] - _fmt_ = {'usedGpuMemory': "%d B"} - -class c_nvmlBridgeChipInfo_t(_PrintableStructure): - _fields_ = [ - ('type', _nvmlBridgeChipType_t), - ('fwVersion', c_uint), - ] - -class c_nvmlBridgeChipHierarchy_t(_PrintableStructure): - _fields_ = [ - ('bridgeCount', c_uint), - ('bridgeChipInfo', c_nvmlBridgeChipInfo_t * 128), - ] - -class c_nvmlEccErrorCounts_t(_PrintableStructure): - _fields_ = [ - ('l1Cache', c_ulonglong), - ('l2Cache', c_ulonglong), - ('deviceMemory', c_ulonglong), - ('registerFile', c_ulonglong), - ] - -class c_nvmlUtilization_t(_PrintableStructure): - _fields_ = [ - ('gpu', c_uint), - ('memory', c_uint), - ] - _fmt_ = {'': "%d %%"} - -# Added in 2.285 -class c_nvmlHwbcEntry_t(_PrintableStructure): - _fields_ = [ - ('hwbcId', c_uint), - ('firmwareVersion', c_char * 32), - ] - -class c_nvmlValue_t(Union): - _fields_ = [ - ('dVal', c_double), - ('uiVal', c_uint), - ('ulVal', c_ulong), - ('ullVal', c_ulonglong), - ] - -class c_nvmlSample_t(_PrintableStructure): - _fields_ = [ - ('timeStamp', c_ulonglong), - ('sampleValue', c_nvmlValue_t), - ] - -class c_nvmlViolationTime_t(_PrintableStructure): - _fields_ = [ - ('referenceTime', c_ulonglong), - ('violationTime', c_ulonglong), - ] - -## Event structures -class struct_c_nvmlEventSet_t(Structure): - pass # opaque handle -c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t) - -nvmlEventTypeSingleBitEccError = 0x0000000000000001 -nvmlEventTypeDoubleBitEccError = 0x0000000000000002 -nvmlEventTypePState = 0x0000000000000004 -nvmlEventTypeXidCriticalError = 0x0000000000000008 -nvmlEventTypeClock = 0x0000000000000010 -nvmlEventTypeNone = 0x0000000000000000 -nvmlEventTypeAll = ( - nvmlEventTypeNone | - nvmlEventTypeSingleBitEccError | - nvmlEventTypeDoubleBitEccError | - nvmlEventTypePState | - nvmlEventTypeClock | - nvmlEventTypeXidCriticalError - ) - -## Clock Throttle Reasons defines -nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 -nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002 -nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting -nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 -nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 -nvmlClocksThrottleReasonUnknown = 0x8000000000000000 -nvmlClocksThrottleReasonNone = 0x0000000000000000 -nvmlClocksThrottleReasonAll = ( - nvmlClocksThrottleReasonNone | - nvmlClocksThrottleReasonGpuIdle | - nvmlClocksThrottleReasonApplicationsClocksSetting | - nvmlClocksThrottleReasonSwPowerCap | - nvmlClocksThrottleReasonHwSlowdown | - nvmlClocksThrottleReasonUnknown - ) - -class c_nvmlEventData_t(_PrintableStructure): - _fields_ = [ - ('device', c_nvmlDevice_t), - ('eventType', c_ulonglong), - ('eventData', c_ulonglong) - ] - _fmt_ = {'eventType': "0x%08X"} - -class c_nvmlAccountingStats_t(_PrintableStructure): - _fields_ = [ - ('gpuUtilization', c_uint), - ('memoryUtilization', c_uint), - ('maxMemoryUsage', c_ulonglong), - ('time', c_ulonglong), - ('startTime', c_ulonglong), - ('isRunning', c_uint), - ('reserved', c_uint * 5) - ] - -## C function wrappers ## -def nvmlInit(): - _LoadNvmlLibrary() - - # - # Initialize the library - # - fn = _nvmlGetFunctionPointer("nvmlInit_v2") - ret = fn() - _nvmlCheckReturn(ret) - - # Atomically update refcount - global _nvmlLib_refcount - libLoadLock.acquire() - _nvmlLib_refcount += 1 - libLoadLock.release() - return None - -def _LoadNvmlLibrary(): - ''' - Load the library if it isn't loaded already - ''' - global nvmlLib - - if (nvmlLib == None): - # lock to ensure only one caller loads the library - libLoadLock.acquire() - - try: - # ensure the library still isn't loaded - if (nvmlLib == None): - try: - if (sys.platform[:3] == "win"): - searchPaths = [ - os.path.join(os.getenv("ProgramFiles", r"C:\Program Files"), r"NVIDIA Corporation\NVSMI\nvml.dll"), - os.path.join(os.getenv("WinDir", r"C:\Windows"), r"System32\nvml.dll"), - ] - nvmlPath = next((x for x in searchPaths if os.path.isfile(x)), None) - if (nvmlPath == None): - _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) - else: - # cdecl calling convention - nvmlLib = CDLL(nvmlPath) - else: - # assume linux - nvmlLib = CDLL("libnvidia-ml.so.1") - except OSError as ose: - _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) - if (nvmlLib == None): - _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) - finally: - # lock is always freed - libLoadLock.release() - -def nvmlShutdown(): - # - # Leave the library loaded, but shutdown the interface - # - fn = _nvmlGetFunctionPointer("nvmlShutdown") - ret = fn() - _nvmlCheckReturn(ret) - - # Atomically update refcount - global _nvmlLib_refcount - libLoadLock.acquire() - if (0 < _nvmlLib_refcount): - _nvmlLib_refcount -= 1 - libLoadLock.release() - return None - -# Added in 2.285 -def nvmlErrorString(result): - fn = _nvmlGetFunctionPointer("nvmlErrorString") - fn.restype = c_char_p # otherwise return is an int - ret = fn(result) - return ret - -# Added in 2.285 -def nvmlSystemGetNVMLVersion(): - c_version = create_string_buffer(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlSystemGetNVMLVersion") - ret = fn(c_version, c_uint(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE)) - _nvmlCheckReturn(ret) - return c_version.value - -# Added in 2.285 -def nvmlSystemGetProcessName(pid): - c_name = create_string_buffer(1024) - fn = _nvmlGetFunctionPointer("nvmlSystemGetProcessName") - ret = fn(c_uint(pid), c_name, c_uint(1024)) - _nvmlCheckReturn(ret) - return c_name.value - -def nvmlSystemGetDriverVersion(): - c_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverVersion") - ret = fn(c_version, c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) - _nvmlCheckReturn(ret) - return c_version.value - -# Added in 2.285 -def nvmlSystemGetHicVersion(): - c_count = c_uint(0) - hics = None - fn = _nvmlGetFunctionPointer("nvmlSystemGetHicVersion") - - # get the count - ret = fn(byref(c_count), None) - - # this should only fail with insufficient size - if ((ret != NVML_SUCCESS) and - (ret != NVML_ERROR_INSUFFICIENT_SIZE)): - raise NVMLError(ret) - - # if there are no hics - if (c_count.value == 0): - return [] - - hic_array = c_nvmlHwbcEntry_t * c_count.value - hics = hic_array() - ret = fn(byref(c_count), hics) - _nvmlCheckReturn(ret) - return hics - -## Unit get functions -def nvmlUnitGetCount(): - c_count = c_uint() - fn = _nvmlGetFunctionPointer("nvmlUnitGetCount") - ret = fn(byref(c_count)) - _nvmlCheckReturn(ret) - return c_count.value - -def nvmlUnitGetHandleByIndex(index): - c_index = c_uint(index) - unit = c_nvmlUnit_t() - fn = _nvmlGetFunctionPointer("nvmlUnitGetHandleByIndex") - ret = fn(c_index, byref(unit)) - _nvmlCheckReturn(ret) - return unit - -def nvmlUnitGetUnitInfo(unit): - c_info = c_nvmlUnitInfo_t() - fn = _nvmlGetFunctionPointer("nvmlUnitGetUnitInfo") - ret = fn(unit, byref(c_info)) - _nvmlCheckReturn(ret) - return c_info - -def nvmlUnitGetLedState(unit): - c_state = c_nvmlLedState_t() - fn = _nvmlGetFunctionPointer("nvmlUnitGetLedState") - ret = fn(unit, byref(c_state)) - _nvmlCheckReturn(ret) - return c_state - -def nvmlUnitGetPsuInfo(unit): - c_info = c_nvmlPSUInfo_t() - fn = _nvmlGetFunctionPointer("nvmlUnitGetPsuInfo") - ret = fn(unit, byref(c_info)) - _nvmlCheckReturn(ret) - return c_info - -def nvmlUnitGetTemperature(unit, type): - c_temp = c_uint() - fn = _nvmlGetFunctionPointer("nvmlUnitGetTemperature") - ret = fn(unit, c_uint(type), byref(c_temp)) - _nvmlCheckReturn(ret) - return c_temp.value - -def nvmlUnitGetFanSpeedInfo(unit): - c_speeds = c_nvmlUnitFanSpeeds_t() - fn = _nvmlGetFunctionPointer("nvmlUnitGetFanSpeedInfo") - ret = fn(unit, byref(c_speeds)) - _nvmlCheckReturn(ret) - return c_speeds - -# added to API -def nvmlUnitGetDeviceCount(unit): - c_count = c_uint(0) - # query the unit to determine device count - fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") - ret = fn(unit, byref(c_count), None) - if (ret == NVML_ERROR_INSUFFICIENT_SIZE): - ret = NVML_SUCCESS - _nvmlCheckReturn(ret) - return c_count.value - -def nvmlUnitGetDevices(unit): - c_count = c_uint(nvmlUnitGetDeviceCount(unit)) - device_array = c_nvmlDevice_t * c_count.value - c_devices = device_array() - fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") - ret = fn(unit, byref(c_count), c_devices) - _nvmlCheckReturn(ret) - return c_devices - -## Device get functions -def nvmlDeviceGetCount(): - c_count = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetCount_v2") - ret = fn(byref(c_count)) - _nvmlCheckReturn(ret) - return c_count.value - -def nvmlDeviceGetHandleByIndex(index): - c_index = c_uint(index) - device = c_nvmlDevice_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByIndex_v2") - ret = fn(c_index, byref(device)) - _nvmlCheckReturn(ret) - return device - -def nvmlDeviceGetHandleBySerial(serial): - c_serial = c_char_p(serial) - device = c_nvmlDevice_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleBySerial") - ret = fn(c_serial, byref(device)) - _nvmlCheckReturn(ret) - return device - -def nvmlDeviceGetHandleByUUID(uuid): - c_uuid = c_char_p(uuid) - device = c_nvmlDevice_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByUUID") - ret = fn(c_uuid, byref(device)) - _nvmlCheckReturn(ret) - return device - -def nvmlDeviceGetHandleByPciBusId(pciBusId): - c_busId = c_char_p(pciBusId) - device = c_nvmlDevice_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByPciBusId_v2") - ret = fn(c_busId, byref(device)) - _nvmlCheckReturn(ret) - return device - -def nvmlDeviceGetName(handle): - c_name = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetName") - ret = fn(handle, c_name, c_uint(NVML_DEVICE_NAME_BUFFER_SIZE)) - _nvmlCheckReturn(ret) - return c_name.value - -def nvmlDeviceGetBoardId(handle): - c_id = c_uint(); - fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardId") - ret = fn(handle, byref(c_id)) - _nvmlCheckReturn(ret) - return c_id.value - -def nvmlDeviceGetMultiGpuBoard(handle): - c_multiGpu = c_uint(); - fn = _nvmlGetFunctionPointer("nvmlDeviceGetMultiGpuBoard") - ret = fn(handle, byref(c_multiGpu)) - _nvmlCheckReturn(ret) - return c_multiGpu.value - -def nvmlDeviceGetBrand(handle): - c_type = _nvmlBrandType_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetBrand") - ret = fn(handle, byref(c_type)) - _nvmlCheckReturn(ret) - return c_type.value - -def nvmlDeviceGetSerial(handle): - c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetSerial") - ret = fn(handle, c_serial, c_uint(NVML_DEVICE_SERIAL_BUFFER_SIZE)) - _nvmlCheckReturn(ret) - return c_serial.value - -def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): - affinity_array = c_ulonglong * cpuSetSize - c_affinity = affinity_array() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinity") - ret = fn(handle, cpuSetSize, byref(c_affinity)) - _nvmlCheckReturn(ret) - return c_affinity - -def nvmlDeviceSetCpuAffinity(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetCpuAffinity") - ret = fn(handle) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceClearCpuAffinity(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceClearCpuAffinity") - ret = fn(handle) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceGetMinorNumber(handle): - c_minor_number = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinorNumber") - ret = fn(handle, byref(c_minor_number)) - _nvmlCheckReturn(ret) - return c_minor_number.value - -def nvmlDeviceGetUUID(handle): - c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetUUID") - ret = fn(handle, c_uuid, c_uint(NVML_DEVICE_UUID_BUFFER_SIZE)) - _nvmlCheckReturn(ret) - return c_uuid.value - -def nvmlDeviceGetInforomVersion(handle, infoRomObject): - c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion") - ret = fn(handle, _nvmlInforomObject_t(infoRomObject), - c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) - _nvmlCheckReturn(ret) - return c_version.value - -# Added in 4.304 -def nvmlDeviceGetInforomImageVersion(handle): - c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomImageVersion") - ret = fn(handle, c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) - _nvmlCheckReturn(ret) - return c_version.value - -# Added in 4.304 -def nvmlDeviceGetInforomConfigurationChecksum(handle): - c_checksum = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomConfigurationChecksum") - ret = fn(handle, byref(c_checksum)) - _nvmlCheckReturn(ret) - return c_checksum.value - -# Added in 4.304 -def nvmlDeviceValidateInforom(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom") - ret = fn(handle) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceGetDisplayMode(handle): - c_mode = _nvmlEnableState_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayMode") - ret = fn(handle, byref(c_mode)) - _nvmlCheckReturn(ret) - return c_mode.value - -def nvmlDeviceGetDisplayActive(handle): - c_mode = _nvmlEnableState_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive") - ret = fn(handle, byref(c_mode)) - _nvmlCheckReturn(ret) - return c_mode.value - - -def nvmlDeviceGetPersistenceMode(handle): - c_state = _nvmlEnableState_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPersistenceMode") - ret = fn(handle, byref(c_state)) - _nvmlCheckReturn(ret) - return c_state.value - -def nvmlDeviceGetPciInfo(handle): - c_info = nvmlPciInfo_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v2") - ret = fn(handle, byref(c_info)) - _nvmlCheckReturn(ret) - return c_info - -def nvmlDeviceGetClockInfo(handle, type): - c_clock = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo") - ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) - _nvmlCheckReturn(ret) - return c_clock.value - -# Added in 2.285 -def nvmlDeviceGetMaxClockInfo(handle, type): - c_clock = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxClockInfo") - ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) - _nvmlCheckReturn(ret) - return c_clock.value - -# Added in 4.304 -def nvmlDeviceGetApplicationsClock(handle, type): - c_clock = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetApplicationsClock") - ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) - _nvmlCheckReturn(ret) - return c_clock.value - -# Added in 5.319 -def nvmlDeviceGetDefaultApplicationsClock(handle, type): - c_clock = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultApplicationsClock") - ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) - _nvmlCheckReturn(ret) - return c_clock.value - -# Added in 4.304 -def nvmlDeviceGetSupportedMemoryClocks(handle): - # first call to get the size - c_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks") - ret = fn(handle, byref(c_count), None) - - if (ret == NVML_SUCCESS): - # special case, no clocks - return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): - # typical case - clocks_array = c_uint * c_count.value - c_clocks = clocks_array() - - # make the call again - ret = fn(handle, byref(c_count), c_clocks) - _nvmlCheckReturn(ret) - - procs = [] - for i in range(c_count.value): - procs.append(c_clocks[i]) - - return procs - else: - # error case - raise NVMLError(ret) - -# Added in 4.304 -def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): - # first call to get the size - c_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks") - ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None) - - if (ret == NVML_SUCCESS): - # special case, no clocks - return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): - # typical case - clocks_array = c_uint * c_count.value - c_clocks = clocks_array() - - # make the call again - ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), c_clocks) - _nvmlCheckReturn(ret) - - procs = [] - for i in range(c_count.value): - procs.append(c_clocks[i]) - - return procs - else: - # error case - raise NVMLError(ret) - -def nvmlDeviceGetFanSpeed(handle): - c_speed = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed") - ret = fn(handle, byref(c_speed)) - _nvmlCheckReturn(ret) - return c_speed.value - -def nvmlDeviceGetTemperature(handle, sensor): - c_temp = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperature") - ret = fn(handle, _nvmlTemperatureSensors_t(sensor), byref(c_temp)) - _nvmlCheckReturn(ret) - return c_temp.value - -def nvmlDeviceGetTemperatureThreshold(handle, threshold): - c_temp = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureThreshold") - ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) - _nvmlCheckReturn(ret) - return c_temp.value - -# DEPRECATED use nvmlDeviceGetPerformanceState -def nvmlDeviceGetPowerState(handle): - c_pstate = _nvmlPstates_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerState") - ret = fn(handle, byref(c_pstate)) - _nvmlCheckReturn(ret) - return c_pstate.value - -def nvmlDeviceGetPerformanceState(handle): - c_pstate = _nvmlPstates_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState") - ret = fn(handle, byref(c_pstate)) - _nvmlCheckReturn(ret) - return c_pstate.value - -def nvmlDeviceGetPowerManagementMode(handle): - c_pcapMode = _nvmlEnableState_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementMode") - ret = fn(handle, byref(c_pcapMode)) - _nvmlCheckReturn(ret) - return c_pcapMode.value - -def nvmlDeviceGetPowerManagementLimit(handle): - c_limit = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit") - ret = fn(handle, byref(c_limit)) - _nvmlCheckReturn(ret) - return c_limit.value - -# Added in 4.304 -def nvmlDeviceGetPowerManagementLimitConstraints(handle): - c_minLimit = c_uint() - c_maxLimit = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimitConstraints") - ret = fn(handle, byref(c_minLimit), byref(c_maxLimit)) - _nvmlCheckReturn(ret) - return [c_minLimit.value, c_maxLimit.value] - -# Added in 4.304 -def nvmlDeviceGetPowerManagementDefaultLimit(handle): - c_limit = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementDefaultLimit") - ret = fn(handle, byref(c_limit)) - _nvmlCheckReturn(ret) - return c_limit.value - - -# Added in 331 -def nvmlDeviceGetEnforcedPowerLimit(handle): - c_limit = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetEnforcedPowerLimit") - ret = fn(handle, byref(c_limit)) - _nvmlCheckReturn(ret) - return c_limit.value - -def nvmlDeviceGetPowerUsage(handle): - c_watts = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerUsage") - ret = fn(handle, byref(c_watts)) - _nvmlCheckReturn(ret) - return c_watts.value - -# Added in 4.304 -def nvmlDeviceGetGpuOperationMode(handle): - c_currState = _nvmlGpuOperationMode_t() - c_pendingState = _nvmlGpuOperationMode_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuOperationMode") - ret = fn(handle, byref(c_currState), byref(c_pendingState)) - _nvmlCheckReturn(ret) - return [c_currState.value, c_pendingState.value] - -# Added in 4.304 -def nvmlDeviceGetCurrentGpuOperationMode(handle): - return nvmlDeviceGetGpuOperationMode(handle)[0] - -# Added in 4.304 -def nvmlDeviceGetPendingGpuOperationMode(handle): - return nvmlDeviceGetGpuOperationMode(handle)[1] - -def nvmlDeviceGetMemoryInfo(handle): - c_memory = c_nvmlMemory_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo") - ret = fn(handle, byref(c_memory)) - _nvmlCheckReturn(ret) - return c_memory - -def nvmlDeviceGetBAR1MemoryInfo(handle): - c_bar1_memory = c_nvmlBAR1Memory_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetBAR1MemoryInfo") - ret = fn(handle, byref(c_bar1_memory)) - _nvmlCheckReturn(ret) - return c_bar1_memory - -def nvmlDeviceGetComputeMode(handle): - c_mode = _nvmlComputeMode_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode") - ret = fn(handle, byref(c_mode)) - _nvmlCheckReturn(ret) - return c_mode.value - -def nvmlDeviceGetEccMode(handle): - c_currState = _nvmlEnableState_t() - c_pendingState = _nvmlEnableState_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetEccMode") - ret = fn(handle, byref(c_currState), byref(c_pendingState)) - _nvmlCheckReturn(ret) - return [c_currState.value, c_pendingState.value] - -# added to API -def nvmlDeviceGetCurrentEccMode(handle): - return nvmlDeviceGetEccMode(handle)[0] - -# added to API -def nvmlDeviceGetPendingEccMode(handle): - return nvmlDeviceGetEccMode(handle)[1] - -def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType): - c_count = c_ulonglong() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEccErrors") - ret = fn(handle, _nvmlMemoryErrorType_t(errorType), - _nvmlEccCounterType_t(counterType), byref(c_count)) - _nvmlCheckReturn(ret) - return c_count.value - -# This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter -def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType): - c_counts = c_nvmlEccErrorCounts_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetDetailedEccErrors") - ret = fn(handle, _nvmlMemoryErrorType_t(errorType), - _nvmlEccCounterType_t(counterType), byref(c_counts)) - _nvmlCheckReturn(ret) - return c_counts - -# Added in 4.304 -def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType): - c_count = c_ulonglong() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryErrorCounter") - ret = fn(handle, - _nvmlMemoryErrorType_t(errorType), - _nvmlEccCounterType_t(counterType), - _nvmlMemoryLocation_t(locationType), - byref(c_count)) - _nvmlCheckReturn(ret) - return c_count.value - -def nvmlDeviceGetUtilizationRates(handle): - c_util = c_nvmlUtilization_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates") - ret = fn(handle, byref(c_util)) - _nvmlCheckReturn(ret) - return c_util - -def nvmlDeviceGetEncoderUtilization(handle): - c_util = c_uint() - c_samplingPeriod = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderUtilization") - ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) - _nvmlCheckReturn(ret) - return [c_util.value, c_samplingPeriod.value] - -def nvmlDeviceGetDecoderUtilization(handle): - c_util = c_uint() - c_samplingPeriod = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetDecoderUtilization") - ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) - _nvmlCheckReturn(ret) - return [c_util.value, c_samplingPeriod.value] - -def nvmlDeviceGetPcieReplayCounter(handle): - c_replay = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieReplayCounter") - ret = fn(handle, byref(c_replay)) - _nvmlCheckReturn(ret) - return c_replay.value - -def nvmlDeviceGetDriverModel(handle): - c_currModel = _nvmlDriverModel_t() - c_pendingModel = _nvmlDriverModel_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetDriverModel") - ret = fn(handle, byref(c_currModel), byref(c_pendingModel)) - _nvmlCheckReturn(ret) - return [c_currModel.value, c_pendingModel.value] - -# added to API -def nvmlDeviceGetCurrentDriverModel(handle): - return nvmlDeviceGetDriverModel(handle)[0] - -# added to API -def nvmlDeviceGetPendingDriverModel(handle): - return nvmlDeviceGetDriverModel(handle)[1] - -# Added in 2.285 -def nvmlDeviceGetVbiosVersion(handle): - c_version = create_string_buffer(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVbiosVersion") - ret = fn(handle, c_version, c_uint(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE)) - _nvmlCheckReturn(ret) - return c_version.value - -# Added in 2.285 -def nvmlDeviceGetComputeRunningProcesses(handle): - # first call to get the size - c_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses") - ret = fn(handle, byref(c_count), None) - - if (ret == NVML_SUCCESS): - # special case, no running processes - return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): - # typical case - # oversize the array incase more processes are created - c_count.value = c_count.value * 2 + 5 - proc_array = c_nvmlProcessInfo_t * c_count.value - c_procs = proc_array() - - # make the call again - ret = fn(handle, byref(c_count), c_procs) - _nvmlCheckReturn(ret) - - procs = [] - for i in range(c_count.value): - # use an alternative struct for this object - obj = nvmlStructToFriendlyObject(c_procs[i]) - if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): - # special case for WDDM on Windows, see comment above - obj.usedGpuMemory = None - procs.append(obj) - - return procs - else: - # error case - raise NVMLError(ret) - -def nvmlDeviceGetGraphicsRunningProcesses(handle): - # first call to get the size - c_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses") - ret = fn(handle, byref(c_count), None) - - if (ret == NVML_SUCCESS): - # special case, no running processes - return [] - elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): - # typical case - # oversize the array incase more processes are created - c_count.value = c_count.value * 2 + 5 - proc_array = c_nvmlProcessInfo_t * c_count.value - c_procs = proc_array() - - # make the call again - ret = fn(handle, byref(c_count), c_procs) - _nvmlCheckReturn(ret) - - procs = [] - for i in range(c_count.value): - # use an alternative struct for this object - obj = nvmlStructToFriendlyObject(c_procs[i]) - if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): - # special case for WDDM on Windows, see comment above - obj.usedGpuMemory = None - procs.append(obj) - - return procs - else: - # error case - raise NVMLError(ret) - -def nvmlDeviceGetAutoBoostedClocksEnabled(handle): - c_isEnabled = _nvmlEnableState_t() - c_defaultIsEnabled = _nvmlEnableState_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetAutoBoostedClocksEnabled") - ret = fn(handle, byref(c_isEnabled), byref(c_defaultIsEnabled)) - _nvmlCheckReturn(ret) - return [c_isEnabled.value, c_defaultIsEnabled.value] - #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks - -## Set functions -def nvmlUnitSetLedState(unit, color): - fn = _nvmlGetFunctionPointer("nvmlUnitSetLedState") - ret = fn(unit, _nvmlLedColor_t(color)) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceSetPersistenceMode(handle, mode): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode") - ret = fn(handle, _nvmlEnableState_t(mode)) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceSetComputeMode(handle, mode): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode") - ret = fn(handle, _nvmlComputeMode_t(mode)) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceSetEccMode(handle, mode): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode") - ret = fn(handle, _nvmlEnableState_t(mode)) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceClearEccErrorCounts(handle, counterType): - fn = _nvmlGetFunctionPointer("nvmlDeviceClearEccErrorCounts") - ret = fn(handle, _nvmlEccCounterType_t(counterType)) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceSetDriverModel(handle, model): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetDriverModel") - ret = fn(handle, _nvmlDriverModel_t(model)) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled") - ret = fn(handle, _nvmlEnableState_t(enabled)) - _nvmlCheckReturn(ret) - return None - #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks - -def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled") - ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags)) - _nvmlCheckReturn(ret) - return None - #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks - -# Added in 4.304 -def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetApplicationsClocks") - ret = fn(handle, c_uint(maxMemClockMHz), c_uint(maxGraphicsClockMHz)) - _nvmlCheckReturn(ret) - return None - -# Added in 4.304 -def nvmlDeviceResetApplicationsClocks(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks") - ret = fn(handle) - _nvmlCheckReturn(ret) - return None - -# Added in 4.304 -def nvmlDeviceSetPowerManagementLimit(handle, limit): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit") - ret = fn(handle, c_uint(limit)) - _nvmlCheckReturn(ret) - return None - -# Added in 4.304 -def nvmlDeviceSetGpuOperationMode(handle, mode): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode") - ret = fn(handle, _nvmlGpuOperationMode_t(mode)) - _nvmlCheckReturn(ret) - return None - -# Added in 2.285 -def nvmlEventSetCreate(): - fn = _nvmlGetFunctionPointer("nvmlEventSetCreate") - eventSet = c_nvmlEventSet_t() - ret = fn(byref(eventSet)) - _nvmlCheckReturn(ret) - return eventSet - -# Added in 2.285 -def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): - fn = _nvmlGetFunctionPointer("nvmlDeviceRegisterEvents") - ret = fn(handle, c_ulonglong(eventTypes), eventSet) - _nvmlCheckReturn(ret) - return None - -# Added in 2.285 -def nvmlDeviceGetSupportedEventTypes(handle): - c_eventTypes = c_ulonglong() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedEventTypes") - ret = fn(handle, byref(c_eventTypes)) - _nvmlCheckReturn(ret) - return c_eventTypes.value - -# Added in 2.285 -# raises NVML_ERROR_TIMEOUT exception on timeout -def nvmlEventSetWait(eventSet, timeoutms): - fn = _nvmlGetFunctionPointer("nvmlEventSetWait") - data = c_nvmlEventData_t() - ret = fn(eventSet, byref(data), c_uint(timeoutms)) - _nvmlCheckReturn(ret) - return data - -# Added in 2.285 -def nvmlEventSetFree(eventSet): - fn = _nvmlGetFunctionPointer("nvmlEventSetFree") - ret = fn(eventSet) - _nvmlCheckReturn(ret) - return None - -# Added in 3.295 -def nvmlDeviceOnSameBoard(handle1, handle2): - fn = _nvmlGetFunctionPointer("nvmlDeviceOnSameBoard") - onSameBoard = c_int() - ret = fn(handle1, handle2, byref(onSameBoard)) - _nvmlCheckReturn(ret) - return (onSameBoard.value != 0) - -# Added in 3.295 -def nvmlDeviceGetCurrPcieLinkGeneration(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkGeneration") - gen = c_uint() - ret = fn(handle, byref(gen)) - _nvmlCheckReturn(ret) - return gen.value - -# Added in 3.295 -def nvmlDeviceGetMaxPcieLinkGeneration(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkGeneration") - gen = c_uint() - ret = fn(handle, byref(gen)) - _nvmlCheckReturn(ret) - return gen.value - -# Added in 3.295 -def nvmlDeviceGetCurrPcieLinkWidth(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkWidth") - width = c_uint() - ret = fn(handle, byref(width)) - _nvmlCheckReturn(ret) - return width.value - -# Added in 3.295 -def nvmlDeviceGetMaxPcieLinkWidth(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkWidth") - width = c_uint() - ret = fn(handle, byref(width)) - _nvmlCheckReturn(ret) - return width.value - -# Added in 4.304 -def nvmlDeviceGetSupportedClocksThrottleReasons(handle): - c_reasons= c_ulonglong() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksThrottleReasons") - ret = fn(handle, byref(c_reasons)) - _nvmlCheckReturn(ret) - return c_reasons.value - -# Added in 4.304 -def nvmlDeviceGetCurrentClocksThrottleReasons(handle): - c_reasons= c_ulonglong() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksThrottleReasons") - ret = fn(handle, byref(c_reasons)) - _nvmlCheckReturn(ret) - return c_reasons.value - -# Added in 5.319 -def nvmlDeviceGetIndex(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetIndex") - c_index = c_uint() - ret = fn(handle, byref(c_index)) - _nvmlCheckReturn(ret) - return c_index.value - -# Added in 5.319 -def nvmlDeviceGetAccountingMode(handle): - c_mode = _nvmlEnableState_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingMode") - ret = fn(handle, byref(c_mode)) - _nvmlCheckReturn(ret) - return c_mode.value - -def nvmlDeviceSetAccountingMode(handle, mode): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode") - ret = fn(handle, _nvmlEnableState_t(mode)) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceClearAccountingPids(handle): - fn = _nvmlGetFunctionPointer("nvmlDeviceClearAccountingPids") - ret = fn(handle) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceGetAccountingStats(handle, pid): - stats = c_nvmlAccountingStats_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingStats") - ret = fn(handle, c_uint(pid), byref(stats)) - _nvmlCheckReturn(ret) - if (stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): - # special case for WDDM on Windows, see comment above - stats.maxMemoryUsage = None - return stats - -def nvmlDeviceGetAccountingPids(handle): - count = c_uint(nvmlDeviceGetAccountingBufferSize(handle)) - pids = (c_uint * count.value)() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids") - ret = fn(handle, byref(count), pids) - _nvmlCheckReturn(ret) - return map(int, pids[0:count.value]) - -def nvmlDeviceGetAccountingBufferSize(handle): - bufferSize = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingBufferSize") - ret = fn(handle, byref(bufferSize)) - _nvmlCheckReturn(ret) - return int(bufferSize.value) - -def nvmlDeviceGetRetiredPages(device, sourceFilter): - c_source = _nvmlPageRetirementCause_t(sourceFilter) - c_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages") - - # First call will get the size - ret = fn(device, c_source, byref(c_count), None) - - # this should only fail with insufficient size - if ((ret != NVML_SUCCESS) and - (ret != NVML_ERROR_INSUFFICIENT_SIZE)): - raise NVMLError(ret) - - # call again with a buffer - # oversize the array for the rare cases where additional pages - # are retired between NVML calls - c_count.value = c_count.value * 2 + 5 - page_array = c_ulonglong * c_count.value - c_pages = page_array() - ret = fn(device, c_source, byref(c_count), c_pages) - _nvmlCheckReturn(ret) - return map(int, c_pages[0:c_count.value]) - -def nvmlDeviceGetRetiredPagesPendingStatus(device): - c_pending = _nvmlEnableState_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPagesPendingStatus") - ret = fn(device, byref(c_pending)) - _nvmlCheckReturn(ret) - return int(c_pending.value) - -def nvmlDeviceGetAPIRestriction(device, apiType): - c_permission = _nvmlEnableState_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetAPIRestriction") - ret = fn(device, _nvmlRestrictedAPI_t(apiType), byref(c_permission)) - _nvmlCheckReturn(ret) - return int(c_permission.value) - -def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetAPIRestriction") - ret = fn(handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted)) - _nvmlCheckReturn(ret) - return None - -def nvmlDeviceGetBridgeChipInfo(handle): - bridgeHierarchy = c_nvmlBridgeChipHierarchy_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetBridgeChipInfo") - ret = fn(handle, byref(bridgeHierarchy)) - _nvmlCheckReturn(ret) - return bridgeHierarchy - -def nvmlDeviceGetSamples(device, sampling_type, timeStamp): - c_sampling_type = _nvmlSamplingType_t(sampling_type) - c_time_stamp = c_ulonglong(timeStamp) - c_sample_count = c_uint(0) - c_sample_value_type = _nvmlValueType_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetSamples") - - ## First Call gets the size - ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), None) - - # Stop if this fails - if (ret != NVML_SUCCESS): - raise NVMLError(ret) - - sampleArray = c_sample_count.value * c_nvmlSample_t - c_samples = sampleArray() - ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), c_samples) - _nvmlCheckReturn(ret) - return (c_sample_value_type.value, c_samples[0:c_sample_count.value]) - -def nvmlDeviceGetViolationStatus(device, perfPolicyType): - c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType) - c_violTime = c_nvmlViolationTime_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetViolationStatus") - - ## Invoke the method to get violation time - ret = fn(device, c_perfPolicy_type, byref(c_violTime)) - _nvmlCheckReturn(ret) - return c_violTime - -def nvmlDeviceGetPcieThroughput(device, counter): - c_util = c_uint() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput") - ret = fn(device, _nvmlPcieUtilCounter_t(counter), byref(c_util)) - _nvmlCheckReturn(ret) - return c_util.value - -def nvmlSystemGetTopologyGpuSet(cpuNumber): - c_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlSystemGetTopologyGpuSet") - - # First call will get the size - ret = fn(cpuNumber, byref(c_count), None) - - if ret != NVML_SUCCESS: - raise NVMLError(ret) - print(c_count.value) - # call again with a buffer - device_array = c_nvmlDevice_t * c_count.value - c_devices = device_array() - ret = fn(cpuNumber, byref(c_count), c_devices) - _nvmlCheckReturn(ret) - return map(None, c_devices[0:c_count.value]) - -def nvmlDeviceGetTopologyNearestGpus(device, level): - c_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyNearestGpus") - - # First call will get the size - ret = fn(device, level, byref(c_count), None) - - if ret != NVML_SUCCESS: - raise NVMLError(ret) - - # call again with a buffer - device_array = c_nvmlDevice_t * c_count.value - c_devices = device_array() - ret = fn(device, level, byref(c_count), c_devices) - _nvmlCheckReturn(ret) - return map(None, c_devices[0:c_count.value]) - -def nvmlDeviceGetTopologyCommonAncestor(device1, device2): - c_level = _nvmlGpuTopologyLevel_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyCommonAncestor") - ret = fn(device1, device2, byref(c_level)) - _nvmlCheckReturn(ret) - return c_level.value - -#DeepFaceLab additions -def nvmlDeviceGetCudaComputeCapability(device): - c_major = c_int() - c_minor = c_int() - - try: - fn = _nvmlGetFunctionPointer("nvmlDeviceGetCudaComputeCapability") - except: - return 9, 9 - - # get the count - ret = fn(device, byref(c_major), byref(c_minor)) - - # this should only fail with insufficient size - if (ret != NVML_SUCCESS): - raise NVMLError(ret) - - return c_major.value, c_minor.value \ No newline at end of file diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 25a6644..eb72c31 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -58,11 +58,12 @@ class SampleProcessor(object): FACE_TYPE_BEGIN = 10 FACE_TYPE_HALF = 10 - FACE_TYPE_FULL = 11 - FACE_TYPE_HEAD = 12 #currently unused - FACE_TYPE_AVATAR = 13 #currently unused - FACE_TYPE_FULL_NO_ALIGN = 14 - FACE_TYPE_HEAD_NO_ALIGN = 15 + FACE_TYPE_MID_FULL = 11 + FACE_TYPE_FULL = 12 + FACE_TYPE_HEAD = 13 #currently unused + FACE_TYPE_AVATAR = 14 #currently unused + FACE_TYPE_FULL_NO_ALIGN = 15 + FACE_TYPE_HEAD_NO_ALIGN = 16 FACE_TYPE_END = 20 MODE_BEGIN = 40 @@ -82,6 +83,7 @@ class SampleProcessor(object): self.ty_range = ty_range SPTF_FACETYPE_TO_FACETYPE = { Types.FACE_TYPE_HALF : FaceType.HALF, + Types.FACE_TYPE_MID_FULL : FaceType.MID_FULL, Types.FACE_TYPE_FULL : FaceType.FULL, Types.FACE_TYPE_HEAD : FaceType.HEAD, Types.FACE_TYPE_FULL_NO_ALIGN : FaceType.FULL_NO_ALIGN,