merge process and interrogate, add exif handler

modules/interrogate.py

@@ -79,7 +79,7 @@ class InterrogateModels:
 
     def load_blip_model(self):
         self.create_fake_fairscale()
-        from repositories.blip import models
+        from repositories.blip import models # pylint: disable=unused-import
         from repositories.blip.models import blip
         import modules.modelloader as modelloader
         model_path = os.path.join(paths.models_path, "BLIP")
@@ -195,3 +195,162 @@ class InterrogateModels:
         self.unload()
         shared.state.end()
         return res
+
+# --------- interrrogate ui
+
+ci = None
+low_vram = False
+
+
+class BatchWriter:
+    def __init__(self, folder):
+        self.folder = folder
+        self.csv, self.file = None, None
+
+    def add(self, file, prompt):
+        txt_file = os.path.splitext(file)[0] + ".txt"
+        with open(os.path.join(self.folder, txt_file), 'w', encoding='utf-8') as f:
+            f.write(prompt)
+
+    def close(self):
+        if self.file is not None:
+            self.file.close()
+
+
+def get_clip_models():
+    import open_clip
+    return ['/'.join(x) for x in open_clip.list_pretrained()]
+
+
+def load_interrogator(model):
+    from clip_interrogator import Config, Interrogator
+    global ci # pylint: disable=global-statement
+    if ci is None:
+        config = Config(device=devices.get_optimal_device(), cache_path=os.path.join(paths.models_path, 'Interrogator'), clip_model_name=model, quiet=True)
+        if low_vram:
+            config.apply_low_vram_defaults()
+        shared.log.info(f'Interrogate load: config={config}')
+        ci = Interrogator(config)
+    elif model != ci.config.clip_model_name:
+        ci.config.clip_model_name = model
+        shared.log.info(f'Interrogate load: config={ci.config}')
+        ci.load_clip_model()
+
+
+def unload_clip_model():
+    if ci is not None:
+        shared.log.debug('Interrogate offload')
+        ci.caption_model = ci.caption_model.to(devices.cpu)
+        ci.clip_model = ci.clip_model.to(devices.cpu)
+        ci.caption_offloaded = True
+        ci.clip_offloaded = True
+        devices.torch_gc()
+
+
+def interrogate(image, mode, caption=None):
+    shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
+    if mode == 'best':
+        prompt = ci.interrogate(image, caption=caption)
+    elif mode == 'caption':
+        prompt = ci.generate_caption(image) if caption is None else caption
+    elif mode == 'classic':
+        prompt = ci.interrogate_classic(image, caption=caption)
+    elif mode == 'fast':
+        prompt = ci.interrogate_fast(image, caption=caption)
+    elif mode == 'negative':
+        prompt = ci.interrogate_negative(image)
+    else:
+        raise RuntimeError(f"Unknown mode {mode}")
+    return prompt
+
+
+def interrogate_image(image, model, mode):
+    shared.state.begin()
+    shared.state.job = 'interrogate'
+    try:
+        if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
+            lowvram.send_everything_to_cpu()
+            devices.torch_gc()
+        load_interrogator(model)
+        image = image.convert('RGB')
+        shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
+        prompt = interrogate(image, mode)
+    except Exception as e:
+        prompt = f"Exception {type(e)}"
+        shared.log.error(f'Interrogate: {e}')
+    shared.state.end()
+    return prompt
+
+
+def interrogate_batch(batch_files, batch_folder, batch_str, model, mode, write):
+    files = []
+    if batch_files is not None:
+        files += [f.name for f in batch_files]
+    if batch_folder is not None:
+        files += [f.name for f in batch_folder]
+    if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str):
+        files += [os.path.join(batch_str, f) for f in os.listdir(batch_str) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
+    if len(files) == 0:
+        shared.log.error('Interrogate batch no images')
+        return ''
+    shared.state.begin()
+    shared.state.job = 'batch interrogate'
+    prompts = []
+    try:
+        if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
+            lowvram.send_everything_to_cpu()
+            devices.torch_gc()
+        load_interrogator(model)
+        shared.log.info(f'Interrogate batch: images={len(files)} mode={mode} config={ci.config}')
+        captions = []
+        # first pass: generate captions
+        for file in files:
+            caption = ""
+            try:
+                if shared.state.interrupted:
+                    break
+                image = Image.open(file).convert('RGB')
+                caption = ci.generate_caption(image)
+            except Exception as e:
+                shared.log.error(f'Interrogate caption: {e}')
+            finally:
+                captions.append(caption)
+        # second pass: interrogate
+        if write:
+            writer = BatchWriter(os.path.dirname(files[0]))
+        for idx, file in enumerate(files):
+            try:
+                if shared.state.interrupted:
+                    break
+                image = Image.open(file).convert('RGB')
+                prompt = interrogate(image, mode, caption=captions[idx])
+                prompts.append(prompt)
+                if write:
+                    writer.add(file, prompt)
+            except OSError as e:
+                shared.log.error(f'Interrogate batch: {e}')
+        if write:
+            writer.close()
+        ci.config.quiet = False
+        unload_clip_model()
+    except Exception as e:
+        shared.log.error(f'Interrogate batch: {e}')
+    shared.state.end()
+    return '\n\n'.join(prompts)
+
+
+def analyze_image(image, model):
+    load_interrogator(model)
+    image = image.convert('RGB')
+    image_features = ci.image_to_features(image)
+    top_mediums = ci.mediums.rank(image_features, 5)
+    top_artists = ci.artists.rank(image_features, 5)
+    top_movements = ci.movements.rank(image_features, 5)
+    top_trendings = ci.trendings.rank(image_features, 5)
+    top_flavors = ci.flavors.rank(image_features, 5)
+    medium_ranks = dict(zip(top_mediums, ci.similarities(image_features, top_mediums)))
+    artist_ranks = dict(zip(top_artists, ci.similarities(image_features, top_artists)))
+    movement_ranks = dict(zip(top_movements, ci.similarities(image_features, top_movements)))
+    trending_ranks = dict(zip(top_trendings, ci.similarities(image_features, top_trendings)))
+    flavor_ranks = dict(zip(top_flavors, ci.similarities(image_features, top_flavors)))
+    return medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks