Delete unused code.
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user