diff --git a/tagger/interrogator.py b/tagger/interrogator.py index 138972a..974ffe8 100644 --- a/tagger/interrogator.py +++ b/tagger/interrogator.py @@ -100,93 +100,6 @@ class Interrogator: ]: raise NotImplementedError() - -class DeepDanbooruInterrogator(Interrogator): - def __init__(self, name: str, project_path: os.PathLike) -> None: - super().__init__(name) - self.project_path = project_path - - def load(self) -> None: - print(f'Loading {self.name} from {str(self.project_path)}') - - import tensorflow as tf - - # tensorflow maps nearly all vram by default, so we limit this - # https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth - # TODO: only run on the first run - for device in tf.config.experimental.list_physical_devices('GPU'): - tf.config.experimental.set_memory_growth(device, True) - - with tf.device(tf_device_name): - import deepdanbooru.project as ddp - - self.model = ddp.load_model_from_project( - project_path=self.project_path, - compile_model=False - ) - - print(f'Loaded {self.name} model from {str(self.project_path)}') - - self.tags = ddp.load_tags_from_project( - project_path=self.project_path - ) - - def unload(self) -> bool: - # unloaded = super().unload() - - # if unloaded: - # # tensorflow suck - # # https://github.com/keras-team/keras/issues/2102 - # import tensorflow as tf - # tf.keras.backend.clear_session() - # gc.collect() - - # return unloaded - - # There is a bug in Keras where it is not possible to release a model that has been loaded into memory. - # Downgrading to keras==2.1.6 may solve the issue, but it may cause compatibility issues with other packages. - # Using subprocess to create a new process may also solve the problem, but it can be too complex (like Automatic1111 did). - # It seems that for now, the best option is to keep the model in memory, as most users use the Waifu Diffusion model with onnx. - - return False - - def interrogate( - self, - image: Image - ) -> Tuple[ - Dict[str, float], # rating confidents - Dict[str, float] # tag confidents - ]: - # init model - if not hasattr(self, 'model') or self.model is None: - self.load() - - import deepdanbooru.data as ddd - - # convert an image to fit the model - image_bufs = BytesIO() - image.save(image_bufs, format='PNG') - image = ddd.load_image_for_evaluate( - image_bufs, - self.model.input_shape[2], - self.model.input_shape[1] - ) - - image = image.reshape((1, *image.shape[0:3])) - - # evaluate model - result = self.model.predict(image) - - confidents = result[0].tolist() - ratings = {} - tags = {} - - for i, tag in enumerate(self.tags): - tags[tag] = confidents[i] - - return ratings, tags - - class WaifuDiffusionInterrogator(Interrogator): def __init__( self,