merge: modules/caption

Signed-off-by: vladmandic <mandic00@live.com>

modules/caption/deepseek.py

@@ -13,6 +13,7 @@ import sys
 import importlib
 from transformers import AutoModelForCausalLM
 from modules import shared, devices, paths, sd_models
+from modules.sd_offload import register_aux, deregister_aux, move_aux_to_gpu, offload_aux
 from modules.logger import log
 
 
@@ -57,10 +58,11 @@ def load(repo: str):
         )
         vl_gpt.to(dtype=devices.dtype)
         vl_gpt.eval()  # required: trust_remote_code model
+        register_aux('deepseek', vl_gpt)
         loaded_repo = repo
         devices.torch_gc()
         log.info(f'Caption: type=vlm model="DeepSeek VL2" repo="{repo}"')
-    sd_models.move_model(vl_gpt, devices.device)
+    move_aux_to_gpu('deepseek')
     return True
 
 
@@ -69,6 +71,7 @@ def unload():
     global vl_gpt, vl_chat_processor, loaded_repo  # pylint: disable=global-statement
     if vl_gpt is not None:
         log.debug(f'DeepSeek unload: model="{loaded_repo}"')
+        deregister_aux('deepseek')
         sd_models.move_model(vl_gpt, devices.cpu, force=True)
         vl_gpt = None
         vl_chat_processor = None
@@ -102,18 +105,19 @@ def predict(question, image, repo):
     ).to(device=devices.device, dtype=devices.dtype)
     inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
     inputs_embeds = inputs_embeds.to(device=devices.device, dtype=devices.dtype)
-    with devices.inference_context():
-        outputs = vl_gpt.language.generate(
-            inputs_embeds=inputs_embeds,
-            attention_mask=prepare_inputs.attention_mask,
-            pad_token_id=vl_chat_processor.tokenizer.eos_token_id,
-            bos_token_id=vl_chat_processor.tokenizer.bos_token_id,
-            eos_token_id=vl_chat_processor.tokenizer.eos_token_id,
-            max_new_tokens=shared.opts.caption_vlm_max_length,
-            do_sample=False,
-            use_cache=True
-        )
-    if shared.opts.caption_offload:
-        sd_models.move_model(vl_gpt, devices.cpu, force=True)
+    try:
+        with devices.inference_context():
+            outputs = vl_gpt.language.generate(
+                inputs_embeds=inputs_embeds,
+                attention_mask=prepare_inputs.attention_mask,
+                pad_token_id=vl_chat_processor.tokenizer.eos_token_id,
+                bos_token_id=vl_chat_processor.tokenizer.bos_token_id,
+                eos_token_id=vl_chat_processor.tokenizer.eos_token_id,
+                max_new_tokens=shared.opts.caption_vlm_max_length,
+                do_sample=False,
+                use_cache=True
+            )
+    finally:
+        offload_aux('deepseek')
     answer = vl_chat_processor.tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
     return answer

modules/caption/joycaption.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass
 from transformers import AutoProcessor, LlavaForConditionalGeneration
 from modules import shared, devices, sd_models, model_quant
+from modules.sd_offload import register_aux, deregister_aux, move_aux_to_gpu, offload_aux
 from modules.logger import log
 
 
@@ -74,7 +75,8 @@ def load(repo: str = None):
             cache_dir=shared.opts.hfcache_dir,
             **quant_args,
         )
-    sd_models.move_model(llava_model, devices.device)
+        register_aux('joycaption', llava_model)
+    move_aux_to_gpu('joycaption')
 
 
 def unload():
@@ -82,6 +84,7 @@ def unload():
     global llava_model, processor  # pylint: disable=global-statement
     if llava_model is not None:
         log.debug(f'JoyCaption unload: model="{opts.repo}"')
+        deregister_aux('joycaption')
         sd_models.move_model(llava_model, devices.cpu, force=True)
         llava_model = None
         processor = None
@@ -124,5 +127,4 @@ def predict(question: str, image, vqa_model: str = None) -> str:
         caption = caption.replace('\n\n', '\n').strip()
         return caption
     finally:
-        if shared.opts.caption_offload:
-            sd_models.move_model(llava_model, devices.cpu, force=True)
+        offload_aux('joycaption')

modules/caption/joytag.py

@@ -17,6 +17,7 @@ import einops
 from einops.layers.torch import Rearrange
 import huggingface_hub
 from modules import shared, devices, sd_models
+from modules.sd_offload import register_aux, deregister_aux, move_aux_to_gpu, offload_aux
 from modules.logger import log
 from modules.image import convert
 
@@ -1050,8 +1051,9 @@ def load():
         model.eval()  # required: custom loader, not from_pretrained
         with open(os.path.join(folder, 'top_tags.txt'), encoding='utf8') as f:
             tags = [line.strip() for line in f.readlines() if line.strip()]
+        register_aux('joytag', model)
         log.info(f'Caption: type=vlm model="JoyTag" repo="{MODEL_REPO}" tags={len(tags)}')
-    sd_models.move_model(model, devices.device)
+    move_aux_to_gpu('joytag')
 
 
 def unload():
@@ -1059,6 +1061,7 @@ def unload():
     global model, tags  # pylint: disable=global-statement
     if model is not None:
         log.debug('JoyTag unload')
+        deregister_aux('joytag')
         sd_models.move_model(model, devices.cpu, force=True)
         model = None
         tags = None
@@ -1070,11 +1073,12 @@ def unload():
 def predict(image: Image.Image):
     load()
     image_tensor = prepare_image(image, model.image_size).unsqueeze(0).to(device=devices.device, dtype=devices.dtype)
-    with devices.inference_context():
-        preds = model({'image': image_tensor})
-        tag_preds = preds['tags'].sigmoid().cpu()
-    if shared.opts.caption_offload:
-        sd_models.move_model(model, devices.cpu, force=True)
+    try:
+        with devices.inference_context():
+            preds = model({'image': image_tensor})
+            tag_preds = preds['tags'].sigmoid().cpu()
+    finally:
+        offload_aux('joytag')
     scores = {tags[i]: tag_preds[0][i] for i in range(len(tags))}
     if shared.opts.tagger_show_scores:
         predicted_tags = [f'{tag}:{score:.2f}' for tag, score in scores.items() if score > THRESHOLD]

modules/caption/moondream3.py

@@ -9,6 +9,7 @@ import collections
 import transformers
 from PIL import Image
 from modules import shared, devices, sd_models
+from modules.sd_offload import register_aux, deregister_aux, move_aux_to_gpu, offload_aux
 from modules.logger import log
 from modules.caption import vqa_detection
 
@@ -69,11 +70,12 @@ def load_model(repo: str):
         if hasattr(moondream3_model, 'model') and hasattr(moondream3_model.model, 'use_flex_decoding'):
             moondream3_model.model.use_flex_decoding = False
 
+        register_aux('moondream3', moondream3_model)
         loaded = repo
         devices.torch_gc()
 
     # Move model to active device
-    sd_models.move_model(moondream3_model, devices.device)
+    move_aux_to_gpu('moondream3')
     return moondream3_model
 
 
@@ -402,8 +404,7 @@ def predict(question: str, image: Image.Image, repo: str, model_name: str = None
         errors.display(e, 'Moondream3')
         return f"Error: {str(e)}"
     finally:
-        if shared.opts.caption_offload and moondream3_model is not None:
-            sd_models.move_model(moondream3_model, devices.cpu, force=True)
+        offload_aux('moondream3')
 
 
 def clear_cache():
@@ -419,6 +420,7 @@ def unload():
     global moondream3_model, loaded  # pylint: disable=global-statement
     if moondream3_model is not None:
         log.debug(f'Moondream3 unload: model="{loaded}"')
+        deregister_aux('moondream3')
         sd_models.move_model(moondream3_model, devices.cpu, force=True)
         moondream3_model = None
         loaded = None

modules/civitai/_video_helper.py

@@ -0,0 +1,19 @@
+import os
+from modules.logger import log
+
+
+def save_video_frame(filepath: str):
+    from modules import video
+    try:
+        frames, fps, duration, w, h, codec, frame = video.get_video_params(filepath, capture=True)
+    except Exception as e:
+        log.error(f'Video: file={filepath} {e}')
+        return None
+    if frame is not None:
+        basename = os.path.splitext(filepath)
+        thumb = f'{basename[0]}.thumb.jpg'
+        log.debug(f'Video: file={filepath} frames={frames} fps={fps} size={w}x{h} codec={codec} duration={duration} thumb={thumb}')
+        frame.save(thumb)
+    else:
+        log.error(f'Video: file={filepath} no frames found')
+    return frame