From 756b4599be10e5821fa52a56c4b1d4cdf54a79d4 Mon Sep 17 00:00:00 2001 From: vladmandic Date: Fri, 13 Mar 2026 12:22:46 +0100 Subject: [PATCH] merge: modules/caption Signed-off-by: vladmandic --- modules/caption/deepseek.py | 32 ++++++++++++++++++-------------- modules/caption/joycaption.py | 8 +++++--- modules/caption/joytag.py | 16 ++++++++++------ modules/caption/moondream3.py | 8 +++++--- modules/civitai/_video_helper.py | 19 +++++++++++++++++++ 5 files changed, 57 insertions(+), 26 deletions(-) create mode 100644 modules/civitai/_video_helper.py diff --git a/modules/caption/deepseek.py b/modules/caption/deepseek.py index ac661c924..4ddac224f 100644 --- a/modules/caption/deepseek.py +++ b/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 diff --git a/modules/caption/joycaption.py b/modules/caption/joycaption.py index 8a5fbf028..95d8413cd 100644 --- a/modules/caption/joycaption.py +++ b/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') diff --git a/modules/caption/joytag.py b/modules/caption/joytag.py index 8bea01a27..8338783ec 100644 --- a/modules/caption/joytag.py +++ b/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] diff --git a/modules/caption/moondream3.py b/modules/caption/moondream3.py index fae1be97c..d66e2b25c 100644 --- a/modules/caption/moondream3.py +++ b/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 diff --git a/modules/civitai/_video_helper.py b/modules/civitai/_video_helper.py new file mode 100644 index 000000000..8ffe159c6 --- /dev/null +++ b/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