mirror of
https://github.com/vladmandic/automatic
synced 2026-09-20 01:31:13 +02:00
refactor: rename interrogate module to caption
Move all caption-related modules from modules/interrogate/ to modules/caption/ for better naming consistency: - Rename deepbooru, deepseek, joycaption, joytag, moondream3, openclip, tagger, vqa, vqa_detection, waifudiffusion modules - Add new caption.py dispatcher module - Remove old interrogate.py (functionality moved to caption.py)
This commit is contained in:
CalamitousFelicitousness
@@ -0,0 +1,48 @@
|
||||
import time
|
||||
from PIL import Image
|
||||
from modules import shared
|
||||
|
||||
|
||||
def caption(image):
|
||||
if isinstance(image, list):
|
||||
image = image[0] if len(image) > 0 else None
|
||||
if isinstance(image, dict) and 'name' in image:
|
||||
image = Image.open(image['name'])
|
||||
if image is None:
|
||||
shared.log.error('Caption: no image provided')
|
||||
return ''
|
||||
t0 = time.time()
|
||||
if shared.opts.caption_default_type == 'OpenCLiP':
|
||||
shared.log.info(f'Caption: type={shared.opts.caption_default_type} clip="{shared.opts.caption_openclip_model}" blip="{shared.opts.caption_openclip_blip_model}" mode="{shared.opts.caption_openclip_mode}"')
|
||||
from modules.caption import openclip
|
||||
openclip.load_captioner(clip_model=shared.opts.caption_openclip_model, blip_model=shared.opts.caption_openclip_blip_model)
|
||||
openclip.update_caption_params()
|
||||
prompt = openclip.caption(image, mode=shared.opts.caption_openclip_mode)
|
||||
shared.log.debug(f'Caption: time={time.time()-t0:.2f} answer="{prompt}"')
|
||||
return prompt
|
||||
elif shared.opts.caption_default_type == 'Tagger':
|
||||
shared.log.info(f'Caption: type={shared.opts.caption_default_type} model="{shared.opts.waifudiffusion_model}"')
|
||||
from modules.caption import tagger
|
||||
prompt = tagger.tag(
|
||||
image=image,
|
||||
model_name=shared.opts.waifudiffusion_model,
|
||||
general_threshold=shared.opts.tagger_threshold,
|
||||
character_threshold=shared.opts.waifudiffusion_character_threshold,
|
||||
include_rating=shared.opts.tagger_include_rating,
|
||||
exclude_tags=shared.opts.tagger_exclude_tags,
|
||||
max_tags=shared.opts.tagger_max_tags,
|
||||
sort_alpha=shared.opts.tagger_sort_alpha,
|
||||
use_spaces=shared.opts.tagger_use_spaces,
|
||||
escape_brackets=shared.opts.tagger_escape_brackets,
|
||||
)
|
||||
shared.log.debug(f'Caption: time={time.time()-t0:.2f} answer="{prompt}"')
|
||||
return prompt
|
||||
elif shared.opts.caption_default_type == 'VLM':
|
||||
shared.log.info(f'Caption: type={shared.opts.caption_default_type} vlm="{shared.opts.caption_vlm_model}" prompt="{shared.opts.caption_vlm_prompt}"')
|
||||
from modules.caption import vqa
|
||||
prompt = vqa.caption(image=image, model_name=shared.opts.caption_vlm_model, question=shared.opts.caption_vlm_prompt, prompt=None, system_prompt=shared.opts.caption_vlm_system)
|
||||
shared.log.debug(f'Caption: time={time.time()-t0:.2f} answer="{prompt}"')
|
||||
return prompt
|
||||
else:
|
||||
shared.log.error(f'Caption: type="{shared.opts.caption_default_type}" unknown')
|
||||
return ''
|
||||
@@ -0,0 +1,328 @@
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from modules import modelloader, devices, shared
|
||||
|
||||
re_special = re.compile(r'([\\()])')
|
||||
load_lock = threading.Lock()
|
||||
|
||||
|
||||
class DeepDanbooru:
|
||||
def __init__(self):
|
||||
self.model = None
|
||||
|
||||
def load(self):
|
||||
with load_lock:
|
||||
if self.model is not None:
|
||||
return
|
||||
model_path = os.path.join(shared.opts.clip_models_path, "DeepDanbooru")
|
||||
shared.log.debug(f'Caption load: module=DeepDanbooru folder="{model_path}"')
|
||||
files = modelloader.load_models(
|
||||
model_path=model_path,
|
||||
model_url='https://github.com/AUTOMATIC1111/TorchDeepDanbooru/releases/download/v1/model-resnet_custom_v3.pt',
|
||||
ext_filter=[".pt"],
|
||||
download_name='model-resnet_custom_v3.pt',
|
||||
)
|
||||
|
||||
from modules.caption.deepbooru_model import DeepDanbooruModel
|
||||
self.model = DeepDanbooruModel()
|
||||
self.model.load_state_dict(torch.load(files[0], map_location="cpu"))
|
||||
self.model.eval()
|
||||
self.model.to(devices.cpu, devices.dtype)
|
||||
|
||||
def start(self):
|
||||
self.load()
|
||||
self.model.to(devices.device)
|
||||
|
||||
def stop(self):
|
||||
if shared.opts.caption_offload:
|
||||
self.model.to(devices.cpu)
|
||||
devices.torch_gc()
|
||||
|
||||
def tag(self, pil_image, **kwargs):
|
||||
self.start()
|
||||
res = self.tag_multi(pil_image, **kwargs)
|
||||
self.stop()
|
||||
|
||||
return res
|
||||
|
||||
def tag_multi(
|
||||
self,
|
||||
pil_image,
|
||||
general_threshold: float = None,
|
||||
include_rating: bool = None,
|
||||
exclude_tags: str = None,
|
||||
max_tags: int = None,
|
||||
sort_alpha: bool = None,
|
||||
use_spaces: bool = None,
|
||||
escape_brackets: bool = None,
|
||||
):
|
||||
"""Run inference and return formatted tag string.
|
||||
|
||||
Args:
|
||||
pil_image: PIL Image to tag
|
||||
general_threshold: Threshold for tag scores (0-1)
|
||||
include_rating: Whether to include rating tags
|
||||
exclude_tags: Comma-separated tags to exclude
|
||||
max_tags: Maximum number of tags to return
|
||||
sort_alpha: Sort tags alphabetically vs by confidence
|
||||
use_spaces: Use spaces instead of underscores
|
||||
escape_brackets: Escape parentheses/brackets in tags
|
||||
|
||||
Returns:
|
||||
Formatted tag string
|
||||
"""
|
||||
# Use settings defaults if not specified
|
||||
general_threshold = general_threshold or shared.opts.tagger_threshold
|
||||
include_rating = include_rating if include_rating is not None else shared.opts.tagger_include_rating
|
||||
exclude_tags = exclude_tags or shared.opts.tagger_exclude_tags
|
||||
max_tags = max_tags or shared.opts.tagger_max_tags
|
||||
sort_alpha = sort_alpha if sort_alpha is not None else shared.opts.tagger_sort_alpha
|
||||
use_spaces = use_spaces if use_spaces is not None else shared.opts.tagger_use_spaces
|
||||
escape_brackets = escape_brackets if escape_brackets is not None else shared.opts.tagger_escape_brackets
|
||||
|
||||
if isinstance(pil_image, list):
|
||||
pil_image = pil_image[0] if len(pil_image) > 0 else None
|
||||
if isinstance(pil_image, dict) and 'name' in pil_image:
|
||||
pil_image = Image.open(pil_image['name'])
|
||||
if pil_image is None:
|
||||
return ''
|
||||
pic = pil_image.resize((512, 512), resample=Image.Resampling.LANCZOS).convert("RGB")
|
||||
a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
|
||||
with devices.inference_context():
|
||||
x = torch.from_numpy(a).to(device=devices.device, dtype=devices.dtype)
|
||||
y = self.model(x)[0].detach().float().cpu().numpy()
|
||||
probability_dict = {}
|
||||
for current, probability in zip(self.model.tags, y):
|
||||
if probability < general_threshold:
|
||||
continue
|
||||
if current.startswith("rating:") and not include_rating:
|
||||
continue
|
||||
probability_dict[current] = probability
|
||||
if sort_alpha:
|
||||
tags = sorted(probability_dict)
|
||||
else:
|
||||
tags = [tag for tag, _ in sorted(probability_dict.items(), key=lambda x: -x[1])]
|
||||
res = []
|
||||
filtertags = {x.strip().replace(' ', '_') for x in exclude_tags.split(",")}
|
||||
for filtertag in [x for x in tags if x not in filtertags]:
|
||||
probability = probability_dict[filtertag]
|
||||
tag_outformat = filtertag
|
||||
if use_spaces:
|
||||
tag_outformat = tag_outformat.replace('_', ' ')
|
||||
if escape_brackets:
|
||||
tag_outformat = re.sub(re_special, r'\\\1', tag_outformat)
|
||||
if shared.opts.tagger_show_scores:
|
||||
tag_outformat = f"({tag_outformat}:{probability:.2f})"
|
||||
res.append(tag_outformat)
|
||||
if max_tags > 0 and len(res) > max_tags:
|
||||
res = res[:max_tags]
|
||||
return ", ".join(res)
|
||||
|
||||
|
||||
model = DeepDanbooru()
|
||||
|
||||
|
||||
def _save_tags_to_file(img_path, tags_str: str, save_append: bool) -> bool:
|
||||
"""Save tags to a text file with error handling.
|
||||
|
||||
Args:
|
||||
img_path: Path to the image file
|
||||
tags_str: Tags string to save
|
||||
save_append: If True, append to existing file; otherwise overwrite
|
||||
|
||||
Returns:
|
||||
True if save succeeded, False otherwise
|
||||
"""
|
||||
try:
|
||||
txt_path = img_path.with_suffix('.txt')
|
||||
if save_append and txt_path.exists():
|
||||
with open(txt_path, 'a', encoding='utf-8') as f:
|
||||
f.write(f', {tags_str}')
|
||||
else:
|
||||
with open(txt_path, 'w', encoding='utf-8') as f:
|
||||
f.write(tags_str)
|
||||
return True
|
||||
except Exception as e:
|
||||
shared.log.error(f'DeepBooru batch: failed to save file="{img_path}" error={e}')
|
||||
return False
|
||||
|
||||
|
||||
def get_models() -> list:
|
||||
"""Return list of available DeepBooru models (just one)."""
|
||||
return ["DeepBooru"]
|
||||
|
||||
|
||||
def load_model(model_name: str = None) -> bool: # pylint: disable=unused-argument
|
||||
"""Load the DeepBooru model."""
|
||||
try:
|
||||
model.load()
|
||||
return model.model is not None
|
||||
except Exception as e:
|
||||
shared.log.error(f'DeepBooru load: {e}')
|
||||
return False
|
||||
|
||||
|
||||
def unload_model():
|
||||
"""Unload the DeepBooru model and free memory."""
|
||||
if model.model is not None:
|
||||
shared.log.debug('DeepBooru unload')
|
||||
model.model = None
|
||||
devices.torch_gc(force=True)
|
||||
|
||||
|
||||
def tag(image, **kwargs) -> str:
|
||||
"""Tag an image using DeepBooru.
|
||||
|
||||
Args:
|
||||
image: PIL Image to tag
|
||||
**kwargs: Tagger parameters (general_threshold, include_rating, exclude_tags,
|
||||
max_tags, sort_alpha, use_spaces, escape_brackets)
|
||||
|
||||
Returns:
|
||||
Formatted tag string
|
||||
"""
|
||||
import time
|
||||
t0 = time.time()
|
||||
jobid = shared.state.begin('DeepBooru Tag')
|
||||
shared.log.info(f'DeepBooru: image_size={image.size if image else None}')
|
||||
|
||||
try:
|
||||
result = model.tag(image, **kwargs)
|
||||
shared.log.debug(f'DeepBooru: complete time={time.time()-t0:.2f} tags={len(result.split(", ")) if result else 0}')
|
||||
except Exception as e:
|
||||
result = f"Exception {type(e)}"
|
||||
shared.log.error(f'DeepBooru: {e}')
|
||||
|
||||
shared.state.end(jobid)
|
||||
return result
|
||||
|
||||
|
||||
def batch(
|
||||
model_name: str, # pylint: disable=unused-argument
|
||||
batch_files: list,
|
||||
batch_folder: str,
|
||||
batch_str: str,
|
||||
save_output: bool = True,
|
||||
save_append: bool = False,
|
||||
recursive: bool = False,
|
||||
**kwargs
|
||||
) -> str:
|
||||
"""Process multiple images in batch mode.
|
||||
|
||||
Args:
|
||||
model_name: Model name (ignored, only DeepBooru available)
|
||||
batch_files: List of file paths
|
||||
batch_folder: Folder path from file picker
|
||||
batch_str: Folder path as string
|
||||
save_output: Save caption to .txt files
|
||||
save_append: Append to existing caption files
|
||||
recursive: Recursively process subfolders
|
||||
**kwargs: Additional arguments (for interface compatibility)
|
||||
|
||||
Returns:
|
||||
Combined tag results
|
||||
"""
|
||||
import time
|
||||
from pathlib import Path
|
||||
import rich.progress as rp
|
||||
|
||||
# Load model
|
||||
model.load()
|
||||
|
||||
# Collect image files
|
||||
image_files = []
|
||||
image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
|
||||
|
||||
# From file picker
|
||||
if batch_files:
|
||||
for f in batch_files:
|
||||
if isinstance(f, dict):
|
||||
image_files.append(Path(f['name']))
|
||||
elif hasattr(f, 'name'):
|
||||
image_files.append(Path(f.name))
|
||||
else:
|
||||
image_files.append(Path(f))
|
||||
|
||||
# From folder picker
|
||||
if batch_folder:
|
||||
folder_path = None
|
||||
if isinstance(batch_folder, list) and len(batch_folder) > 0:
|
||||
f = batch_folder[0]
|
||||
if isinstance(f, dict):
|
||||
folder_path = Path(f['name']).parent
|
||||
elif hasattr(f, 'name'):
|
||||
folder_path = Path(f.name).parent
|
||||
if folder_path and folder_path.is_dir():
|
||||
if recursive:
|
||||
for ext in image_extensions:
|
||||
image_files.extend(folder_path.rglob(f'*{ext}'))
|
||||
else:
|
||||
for ext in image_extensions:
|
||||
image_files.extend(folder_path.glob(f'*{ext}'))
|
||||
|
||||
# From string path
|
||||
if batch_str and batch_str.strip():
|
||||
folder_path = Path(batch_str.strip())
|
||||
if folder_path.is_dir():
|
||||
if recursive:
|
||||
for ext in image_extensions:
|
||||
image_files.extend(folder_path.rglob(f'*{ext}'))
|
||||
else:
|
||||
for ext in image_extensions:
|
||||
image_files.extend(folder_path.glob(f'*{ext}'))
|
||||
|
||||
# Remove duplicates while preserving order
|
||||
seen = set()
|
||||
unique_files = []
|
||||
for f in image_files:
|
||||
f_resolved = f.resolve()
|
||||
if f_resolved not in seen:
|
||||
seen.add(f_resolved)
|
||||
unique_files.append(f)
|
||||
image_files = unique_files
|
||||
|
||||
if not image_files:
|
||||
shared.log.warning('DeepBooru batch: no images found')
|
||||
return ''
|
||||
|
||||
t0 = time.time()
|
||||
jobid = shared.state.begin('DeepBooru Batch')
|
||||
shared.log.info(f'DeepBooru batch: images={len(image_files)} write={save_output} append={save_append} recursive={recursive}')
|
||||
|
||||
results = []
|
||||
model.start()
|
||||
|
||||
# Progress bar
|
||||
pbar = rp.Progress(rp.TextColumn('[cyan]DeepBooru:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
|
||||
|
||||
with pbar:
|
||||
task = pbar.add_task(total=len(image_files), description='starting...')
|
||||
for img_path in image_files:
|
||||
pbar.update(task, advance=1, description=str(img_path.name))
|
||||
try:
|
||||
if shared.state.interrupted:
|
||||
shared.log.info('DeepBooru batch: interrupted')
|
||||
break
|
||||
|
||||
image = Image.open(img_path)
|
||||
tags_str = model.tag_multi(image, **kwargs)
|
||||
|
||||
if save_output:
|
||||
_save_tags_to_file(img_path, tags_str, save_append)
|
||||
|
||||
results.append(f'{img_path.name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{img_path.name}: {tags_str}')
|
||||
|
||||
except Exception as e:
|
||||
shared.log.error(f'DeepBooru batch: file="{img_path}" error={e}')
|
||||
results.append(f'{img_path.name}: ERROR - {e}')
|
||||
|
||||
model.stop()
|
||||
elapsed = time.time() - t0
|
||||
shared.log.info(f'DeepBooru batch: complete images={len(results)} time={elapsed:.1f}s')
|
||||
shared.state.end(jobid)
|
||||
|
||||
return '\n'.join(results)
|
||||
@@ -0,0 +1,674 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modules import devices
|
||||
|
||||
# see https://github.com/AUTOMATIC1111/TorchDeepDanbooru for more
|
||||
|
||||
|
||||
class DeepDanbooruModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.tags = []
|
||||
self.n_Conv_0 = nn.Conv2d(kernel_size=(7, 7), in_channels=3, out_channels=64, stride=(2, 2))
|
||||
self.n_MaxPool_0 = nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2))
|
||||
self.n_Conv_1 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256)
|
||||
self.n_Conv_2 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=64)
|
||||
self.n_Conv_3 = nn.Conv2d(kernel_size=(3, 3), in_channels=64, out_channels=64)
|
||||
self.n_Conv_4 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256)
|
||||
self.n_Conv_5 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=64)
|
||||
self.n_Conv_6 = nn.Conv2d(kernel_size=(3, 3), in_channels=64, out_channels=64)
|
||||
self.n_Conv_7 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256)
|
||||
self.n_Conv_8 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=64)
|
||||
self.n_Conv_9 = nn.Conv2d(kernel_size=(3, 3), in_channels=64, out_channels=64)
|
||||
self.n_Conv_10 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256)
|
||||
self.n_Conv_11 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=512, stride=(2, 2))
|
||||
self.n_Conv_12 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=128)
|
||||
self.n_Conv_13 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128, stride=(2, 2))
|
||||
self.n_Conv_14 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512)
|
||||
self.n_Conv_15 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128)
|
||||
self.n_Conv_16 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128)
|
||||
self.n_Conv_17 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512)
|
||||
self.n_Conv_18 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128)
|
||||
self.n_Conv_19 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128)
|
||||
self.n_Conv_20 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512)
|
||||
self.n_Conv_21 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128)
|
||||
self.n_Conv_22 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128)
|
||||
self.n_Conv_23 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512)
|
||||
self.n_Conv_24 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128)
|
||||
self.n_Conv_25 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128)
|
||||
self.n_Conv_26 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512)
|
||||
self.n_Conv_27 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128)
|
||||
self.n_Conv_28 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128)
|
||||
self.n_Conv_29 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512)
|
||||
self.n_Conv_30 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128)
|
||||
self.n_Conv_31 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128)
|
||||
self.n_Conv_32 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512)
|
||||
self.n_Conv_33 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128)
|
||||
self.n_Conv_34 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128)
|
||||
self.n_Conv_35 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512)
|
||||
self.n_Conv_36 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=1024, stride=(2, 2))
|
||||
self.n_Conv_37 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=256)
|
||||
self.n_Conv_38 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256, stride=(2, 2))
|
||||
self.n_Conv_39 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_40 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_41 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_42 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_43 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_44 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_45 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_46 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_47 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_48 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_49 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_50 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_51 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_52 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_53 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_54 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_55 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_56 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_57 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_58 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_59 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_60 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_61 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_62 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_63 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_64 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_65 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_66 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_67 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_68 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_69 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_70 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_71 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_72 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_73 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_74 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_75 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_76 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_77 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_78 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_79 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_80 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_81 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_82 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_83 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_84 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_85 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_86 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_87 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_88 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_89 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_90 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_91 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_92 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_93 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_94 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_95 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_96 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_97 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_98 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256, stride=(2, 2))
|
||||
self.n_Conv_99 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_100 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=1024, stride=(2, 2))
|
||||
self.n_Conv_101 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_102 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_103 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_104 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_105 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_106 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_107 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_108 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_109 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_110 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_111 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_112 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_113 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_114 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_115 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_116 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_117 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_118 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_119 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_120 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_121 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_122 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_123 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_124 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_125 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_126 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_127 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_128 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_129 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_130 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_131 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_132 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_133 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_134 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_135 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_136 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_137 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_138 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_139 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_140 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_141 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_142 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_143 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_144 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_145 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_146 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_147 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_148 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_149 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_150 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_151 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_152 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_153 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_154 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_155 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256)
|
||||
self.n_Conv_156 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256)
|
||||
self.n_Conv_157 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024)
|
||||
self.n_Conv_158 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=2048, stride=(2, 2))
|
||||
self.n_Conv_159 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=512)
|
||||
self.n_Conv_160 = nn.Conv2d(kernel_size=(3, 3), in_channels=512, out_channels=512, stride=(2, 2))
|
||||
self.n_Conv_161 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=2048)
|
||||
self.n_Conv_162 = nn.Conv2d(kernel_size=(1, 1), in_channels=2048, out_channels=512)
|
||||
self.n_Conv_163 = nn.Conv2d(kernel_size=(3, 3), in_channels=512, out_channels=512)
|
||||
self.n_Conv_164 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=2048)
|
||||
self.n_Conv_165 = nn.Conv2d(kernel_size=(1, 1), in_channels=2048, out_channels=512)
|
||||
self.n_Conv_166 = nn.Conv2d(kernel_size=(3, 3), in_channels=512, out_channels=512)
|
||||
self.n_Conv_167 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=2048)
|
||||
self.n_Conv_168 = nn.Conv2d(kernel_size=(1, 1), in_channels=2048, out_channels=4096, stride=(2, 2))
|
||||
self.n_Conv_169 = nn.Conv2d(kernel_size=(1, 1), in_channels=2048, out_channels=1024)
|
||||
self.n_Conv_170 = nn.Conv2d(kernel_size=(3, 3), in_channels=1024, out_channels=1024, stride=(2, 2))
|
||||
self.n_Conv_171 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=4096)
|
||||
self.n_Conv_172 = nn.Conv2d(kernel_size=(1, 1), in_channels=4096, out_channels=1024)
|
||||
self.n_Conv_173 = nn.Conv2d(kernel_size=(3, 3), in_channels=1024, out_channels=1024)
|
||||
self.n_Conv_174 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=4096)
|
||||
self.n_Conv_175 = nn.Conv2d(kernel_size=(1, 1), in_channels=4096, out_channels=1024)
|
||||
self.n_Conv_176 = nn.Conv2d(kernel_size=(3, 3), in_channels=1024, out_channels=1024)
|
||||
self.n_Conv_177 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=4096)
|
||||
self.n_Conv_178 = nn.Conv2d(kernel_size=(1, 1), in_channels=4096, out_channels=9176, bias=False)
|
||||
|
||||
def forward(self, *inputs):
|
||||
t_358, = inputs
|
||||
t_359 = t_358.permute(*[0, 3, 1, 2])
|
||||
t_359_padded = F.pad(t_359, [2, 3, 2, 3], value=0)
|
||||
t_360 = self.n_Conv_0(t_359_padded.to(self.n_Conv_0.bias.dtype) if devices.unet_needs_upcast else t_359_padded)
|
||||
t_361 = F.relu(t_360)
|
||||
t_361 = F.pad(t_361, [0, 1, 0, 1], value=float('-inf'))
|
||||
t_362 = self.n_MaxPool_0(t_361)
|
||||
t_363 = self.n_Conv_1(t_362)
|
||||
t_364 = self.n_Conv_2(t_362)
|
||||
t_365 = F.relu(t_364)
|
||||
t_365_padded = F.pad(t_365, [1, 1, 1, 1], value=0)
|
||||
t_366 = self.n_Conv_3(t_365_padded)
|
||||
t_367 = F.relu(t_366)
|
||||
t_368 = self.n_Conv_4(t_367)
|
||||
t_369 = torch.add(t_368, t_363)
|
||||
t_370 = F.relu(t_369)
|
||||
t_371 = self.n_Conv_5(t_370)
|
||||
t_372 = F.relu(t_371)
|
||||
t_372_padded = F.pad(t_372, [1, 1, 1, 1], value=0)
|
||||
t_373 = self.n_Conv_6(t_372_padded)
|
||||
t_374 = F.relu(t_373)
|
||||
t_375 = self.n_Conv_7(t_374)
|
||||
t_376 = torch.add(t_375, t_370)
|
||||
t_377 = F.relu(t_376)
|
||||
t_378 = self.n_Conv_8(t_377)
|
||||
t_379 = F.relu(t_378)
|
||||
t_379_padded = F.pad(t_379, [1, 1, 1, 1], value=0)
|
||||
t_380 = self.n_Conv_9(t_379_padded)
|
||||
t_381 = F.relu(t_380)
|
||||
t_382 = self.n_Conv_10(t_381)
|
||||
t_383 = torch.add(t_382, t_377)
|
||||
t_384 = F.relu(t_383)
|
||||
t_385 = self.n_Conv_11(t_384)
|
||||
t_386 = self.n_Conv_12(t_384)
|
||||
t_387 = F.relu(t_386)
|
||||
t_387_padded = F.pad(t_387, [0, 1, 0, 1], value=0)
|
||||
t_388 = self.n_Conv_13(t_387_padded)
|
||||
t_389 = F.relu(t_388)
|
||||
t_390 = self.n_Conv_14(t_389)
|
||||
t_391 = torch.add(t_390, t_385)
|
||||
t_392 = F.relu(t_391)
|
||||
t_393 = self.n_Conv_15(t_392)
|
||||
t_394 = F.relu(t_393)
|
||||
t_394_padded = F.pad(t_394, [1, 1, 1, 1], value=0)
|
||||
t_395 = self.n_Conv_16(t_394_padded)
|
||||
t_396 = F.relu(t_395)
|
||||
t_397 = self.n_Conv_17(t_396)
|
||||
t_398 = torch.add(t_397, t_392)
|
||||
t_399 = F.relu(t_398)
|
||||
t_400 = self.n_Conv_18(t_399)
|
||||
t_401 = F.relu(t_400)
|
||||
t_401_padded = F.pad(t_401, [1, 1, 1, 1], value=0)
|
||||
t_402 = self.n_Conv_19(t_401_padded)
|
||||
t_403 = F.relu(t_402)
|
||||
t_404 = self.n_Conv_20(t_403)
|
||||
t_405 = torch.add(t_404, t_399)
|
||||
t_406 = F.relu(t_405)
|
||||
t_407 = self.n_Conv_21(t_406)
|
||||
t_408 = F.relu(t_407)
|
||||
t_408_padded = F.pad(t_408, [1, 1, 1, 1], value=0)
|
||||
t_409 = self.n_Conv_22(t_408_padded)
|
||||
t_410 = F.relu(t_409)
|
||||
t_411 = self.n_Conv_23(t_410)
|
||||
t_412 = torch.add(t_411, t_406)
|
||||
t_413 = F.relu(t_412)
|
||||
t_414 = self.n_Conv_24(t_413)
|
||||
t_415 = F.relu(t_414)
|
||||
t_415_padded = F.pad(t_415, [1, 1, 1, 1], value=0)
|
||||
t_416 = self.n_Conv_25(t_415_padded)
|
||||
t_417 = F.relu(t_416)
|
||||
t_418 = self.n_Conv_26(t_417)
|
||||
t_419 = torch.add(t_418, t_413)
|
||||
t_420 = F.relu(t_419)
|
||||
t_421 = self.n_Conv_27(t_420)
|
||||
t_422 = F.relu(t_421)
|
||||
t_422_padded = F.pad(t_422, [1, 1, 1, 1], value=0)
|
||||
t_423 = self.n_Conv_28(t_422_padded)
|
||||
t_424 = F.relu(t_423)
|
||||
t_425 = self.n_Conv_29(t_424)
|
||||
t_426 = torch.add(t_425, t_420)
|
||||
t_427 = F.relu(t_426)
|
||||
t_428 = self.n_Conv_30(t_427)
|
||||
t_429 = F.relu(t_428)
|
||||
t_429_padded = F.pad(t_429, [1, 1, 1, 1], value=0)
|
||||
t_430 = self.n_Conv_31(t_429_padded)
|
||||
t_431 = F.relu(t_430)
|
||||
t_432 = self.n_Conv_32(t_431)
|
||||
t_433 = torch.add(t_432, t_427)
|
||||
t_434 = F.relu(t_433)
|
||||
t_435 = self.n_Conv_33(t_434)
|
||||
t_436 = F.relu(t_435)
|
||||
t_436_padded = F.pad(t_436, [1, 1, 1, 1], value=0)
|
||||
t_437 = self.n_Conv_34(t_436_padded)
|
||||
t_438 = F.relu(t_437)
|
||||
t_439 = self.n_Conv_35(t_438)
|
||||
t_440 = torch.add(t_439, t_434)
|
||||
t_441 = F.relu(t_440)
|
||||
t_442 = self.n_Conv_36(t_441)
|
||||
t_443 = self.n_Conv_37(t_441)
|
||||
t_444 = F.relu(t_443)
|
||||
t_444_padded = F.pad(t_444, [0, 1, 0, 1], value=0)
|
||||
t_445 = self.n_Conv_38(t_444_padded)
|
||||
t_446 = F.relu(t_445)
|
||||
t_447 = self.n_Conv_39(t_446)
|
||||
t_448 = torch.add(t_447, t_442)
|
||||
t_449 = F.relu(t_448)
|
||||
t_450 = self.n_Conv_40(t_449)
|
||||
t_451 = F.relu(t_450)
|
||||
t_451_padded = F.pad(t_451, [1, 1, 1, 1], value=0)
|
||||
t_452 = self.n_Conv_41(t_451_padded)
|
||||
t_453 = F.relu(t_452)
|
||||
t_454 = self.n_Conv_42(t_453)
|
||||
t_455 = torch.add(t_454, t_449)
|
||||
t_456 = F.relu(t_455)
|
||||
t_457 = self.n_Conv_43(t_456)
|
||||
t_458 = F.relu(t_457)
|
||||
t_458_padded = F.pad(t_458, [1, 1, 1, 1], value=0)
|
||||
t_459 = self.n_Conv_44(t_458_padded)
|
||||
t_460 = F.relu(t_459)
|
||||
t_461 = self.n_Conv_45(t_460)
|
||||
t_462 = torch.add(t_461, t_456)
|
||||
t_463 = F.relu(t_462)
|
||||
t_464 = self.n_Conv_46(t_463)
|
||||
t_465 = F.relu(t_464)
|
||||
t_465_padded = F.pad(t_465, [1, 1, 1, 1], value=0)
|
||||
t_466 = self.n_Conv_47(t_465_padded)
|
||||
t_467 = F.relu(t_466)
|
||||
t_468 = self.n_Conv_48(t_467)
|
||||
t_469 = torch.add(t_468, t_463)
|
||||
t_470 = F.relu(t_469)
|
||||
t_471 = self.n_Conv_49(t_470)
|
||||
t_472 = F.relu(t_471)
|
||||
t_472_padded = F.pad(t_472, [1, 1, 1, 1], value=0)
|
||||
t_473 = self.n_Conv_50(t_472_padded)
|
||||
t_474 = F.relu(t_473)
|
||||
t_475 = self.n_Conv_51(t_474)
|
||||
t_476 = torch.add(t_475, t_470)
|
||||
t_477 = F.relu(t_476)
|
||||
t_478 = self.n_Conv_52(t_477)
|
||||
t_479 = F.relu(t_478)
|
||||
t_479_padded = F.pad(t_479, [1, 1, 1, 1], value=0)
|
||||
t_480 = self.n_Conv_53(t_479_padded)
|
||||
t_481 = F.relu(t_480)
|
||||
t_482 = self.n_Conv_54(t_481)
|
||||
t_483 = torch.add(t_482, t_477)
|
||||
t_484 = F.relu(t_483)
|
||||
t_485 = self.n_Conv_55(t_484)
|
||||
t_486 = F.relu(t_485)
|
||||
t_486_padded = F.pad(t_486, [1, 1, 1, 1], value=0)
|
||||
t_487 = self.n_Conv_56(t_486_padded)
|
||||
t_488 = F.relu(t_487)
|
||||
t_489 = self.n_Conv_57(t_488)
|
||||
t_490 = torch.add(t_489, t_484)
|
||||
t_491 = F.relu(t_490)
|
||||
t_492 = self.n_Conv_58(t_491)
|
||||
t_493 = F.relu(t_492)
|
||||
t_493_padded = F.pad(t_493, [1, 1, 1, 1], value=0)
|
||||
t_494 = self.n_Conv_59(t_493_padded)
|
||||
t_495 = F.relu(t_494)
|
||||
t_496 = self.n_Conv_60(t_495)
|
||||
t_497 = torch.add(t_496, t_491)
|
||||
t_498 = F.relu(t_497)
|
||||
t_499 = self.n_Conv_61(t_498)
|
||||
t_500 = F.relu(t_499)
|
||||
t_500_padded = F.pad(t_500, [1, 1, 1, 1], value=0)
|
||||
t_501 = self.n_Conv_62(t_500_padded)
|
||||
t_502 = F.relu(t_501)
|
||||
t_503 = self.n_Conv_63(t_502)
|
||||
t_504 = torch.add(t_503, t_498)
|
||||
t_505 = F.relu(t_504)
|
||||
t_506 = self.n_Conv_64(t_505)
|
||||
t_507 = F.relu(t_506)
|
||||
t_507_padded = F.pad(t_507, [1, 1, 1, 1], value=0)
|
||||
t_508 = self.n_Conv_65(t_507_padded)
|
||||
t_509 = F.relu(t_508)
|
||||
t_510 = self.n_Conv_66(t_509)
|
||||
t_511 = torch.add(t_510, t_505)
|
||||
t_512 = F.relu(t_511)
|
||||
t_513 = self.n_Conv_67(t_512)
|
||||
t_514 = F.relu(t_513)
|
||||
t_514_padded = F.pad(t_514, [1, 1, 1, 1], value=0)
|
||||
t_515 = self.n_Conv_68(t_514_padded)
|
||||
t_516 = F.relu(t_515)
|
||||
t_517 = self.n_Conv_69(t_516)
|
||||
t_518 = torch.add(t_517, t_512)
|
||||
t_519 = F.relu(t_518)
|
||||
t_520 = self.n_Conv_70(t_519)
|
||||
t_521 = F.relu(t_520)
|
||||
t_521_padded = F.pad(t_521, [1, 1, 1, 1], value=0)
|
||||
t_522 = self.n_Conv_71(t_521_padded)
|
||||
t_523 = F.relu(t_522)
|
||||
t_524 = self.n_Conv_72(t_523)
|
||||
t_525 = torch.add(t_524, t_519)
|
||||
t_526 = F.relu(t_525)
|
||||
t_527 = self.n_Conv_73(t_526)
|
||||
t_528 = F.relu(t_527)
|
||||
t_528_padded = F.pad(t_528, [1, 1, 1, 1], value=0)
|
||||
t_529 = self.n_Conv_74(t_528_padded)
|
||||
t_530 = F.relu(t_529)
|
||||
t_531 = self.n_Conv_75(t_530)
|
||||
t_532 = torch.add(t_531, t_526)
|
||||
t_533 = F.relu(t_532)
|
||||
t_534 = self.n_Conv_76(t_533)
|
||||
t_535 = F.relu(t_534)
|
||||
t_535_padded = F.pad(t_535, [1, 1, 1, 1], value=0)
|
||||
t_536 = self.n_Conv_77(t_535_padded)
|
||||
t_537 = F.relu(t_536)
|
||||
t_538 = self.n_Conv_78(t_537)
|
||||
t_539 = torch.add(t_538, t_533)
|
||||
t_540 = F.relu(t_539)
|
||||
t_541 = self.n_Conv_79(t_540)
|
||||
t_542 = F.relu(t_541)
|
||||
t_542_padded = F.pad(t_542, [1, 1, 1, 1], value=0)
|
||||
t_543 = self.n_Conv_80(t_542_padded)
|
||||
t_544 = F.relu(t_543)
|
||||
t_545 = self.n_Conv_81(t_544)
|
||||
t_546 = torch.add(t_545, t_540)
|
||||
t_547 = F.relu(t_546)
|
||||
t_548 = self.n_Conv_82(t_547)
|
||||
t_549 = F.relu(t_548)
|
||||
t_549_padded = F.pad(t_549, [1, 1, 1, 1], value=0)
|
||||
t_550 = self.n_Conv_83(t_549_padded)
|
||||
t_551 = F.relu(t_550)
|
||||
t_552 = self.n_Conv_84(t_551)
|
||||
t_553 = torch.add(t_552, t_547)
|
||||
t_554 = F.relu(t_553)
|
||||
t_555 = self.n_Conv_85(t_554)
|
||||
t_556 = F.relu(t_555)
|
||||
t_556_padded = F.pad(t_556, [1, 1, 1, 1], value=0)
|
||||
t_557 = self.n_Conv_86(t_556_padded)
|
||||
t_558 = F.relu(t_557)
|
||||
t_559 = self.n_Conv_87(t_558)
|
||||
t_560 = torch.add(t_559, t_554)
|
||||
t_561 = F.relu(t_560)
|
||||
t_562 = self.n_Conv_88(t_561)
|
||||
t_563 = F.relu(t_562)
|
||||
t_563_padded = F.pad(t_563, [1, 1, 1, 1], value=0)
|
||||
t_564 = self.n_Conv_89(t_563_padded)
|
||||
t_565 = F.relu(t_564)
|
||||
t_566 = self.n_Conv_90(t_565)
|
||||
t_567 = torch.add(t_566, t_561)
|
||||
t_568 = F.relu(t_567)
|
||||
t_569 = self.n_Conv_91(t_568)
|
||||
t_570 = F.relu(t_569)
|
||||
t_570_padded = F.pad(t_570, [1, 1, 1, 1], value=0)
|
||||
t_571 = self.n_Conv_92(t_570_padded)
|
||||
t_572 = F.relu(t_571)
|
||||
t_573 = self.n_Conv_93(t_572)
|
||||
t_574 = torch.add(t_573, t_568)
|
||||
t_575 = F.relu(t_574)
|
||||
t_576 = self.n_Conv_94(t_575)
|
||||
t_577 = F.relu(t_576)
|
||||
t_577_padded = F.pad(t_577, [1, 1, 1, 1], value=0)
|
||||
t_578 = self.n_Conv_95(t_577_padded)
|
||||
t_579 = F.relu(t_578)
|
||||
t_580 = self.n_Conv_96(t_579)
|
||||
t_581 = torch.add(t_580, t_575)
|
||||
t_582 = F.relu(t_581)
|
||||
t_583 = self.n_Conv_97(t_582)
|
||||
t_584 = F.relu(t_583)
|
||||
t_584_padded = F.pad(t_584, [0, 1, 0, 1], value=0)
|
||||
t_585 = self.n_Conv_98(t_584_padded)
|
||||
t_586 = F.relu(t_585)
|
||||
t_587 = self.n_Conv_99(t_586)
|
||||
t_588 = self.n_Conv_100(t_582)
|
||||
t_589 = torch.add(t_587, t_588)
|
||||
t_590 = F.relu(t_589)
|
||||
t_591 = self.n_Conv_101(t_590)
|
||||
t_592 = F.relu(t_591)
|
||||
t_592_padded = F.pad(t_592, [1, 1, 1, 1], value=0)
|
||||
t_593 = self.n_Conv_102(t_592_padded)
|
||||
t_594 = F.relu(t_593)
|
||||
t_595 = self.n_Conv_103(t_594)
|
||||
t_596 = torch.add(t_595, t_590)
|
||||
t_597 = F.relu(t_596)
|
||||
t_598 = self.n_Conv_104(t_597)
|
||||
t_599 = F.relu(t_598)
|
||||
t_599_padded = F.pad(t_599, [1, 1, 1, 1], value=0)
|
||||
t_600 = self.n_Conv_105(t_599_padded)
|
||||
t_601 = F.relu(t_600)
|
||||
t_602 = self.n_Conv_106(t_601)
|
||||
t_603 = torch.add(t_602, t_597)
|
||||
t_604 = F.relu(t_603)
|
||||
t_605 = self.n_Conv_107(t_604)
|
||||
t_606 = F.relu(t_605)
|
||||
t_606_padded = F.pad(t_606, [1, 1, 1, 1], value=0)
|
||||
t_607 = self.n_Conv_108(t_606_padded)
|
||||
t_608 = F.relu(t_607)
|
||||
t_609 = self.n_Conv_109(t_608)
|
||||
t_610 = torch.add(t_609, t_604)
|
||||
t_611 = F.relu(t_610)
|
||||
t_612 = self.n_Conv_110(t_611)
|
||||
t_613 = F.relu(t_612)
|
||||
t_613_padded = F.pad(t_613, [1, 1, 1, 1], value=0)
|
||||
t_614 = self.n_Conv_111(t_613_padded)
|
||||
t_615 = F.relu(t_614)
|
||||
t_616 = self.n_Conv_112(t_615)
|
||||
t_617 = torch.add(t_616, t_611)
|
||||
t_618 = F.relu(t_617)
|
||||
t_619 = self.n_Conv_113(t_618)
|
||||
t_620 = F.relu(t_619)
|
||||
t_620_padded = F.pad(t_620, [1, 1, 1, 1], value=0)
|
||||
t_621 = self.n_Conv_114(t_620_padded)
|
||||
t_622 = F.relu(t_621)
|
||||
t_623 = self.n_Conv_115(t_622)
|
||||
t_624 = torch.add(t_623, t_618)
|
||||
t_625 = F.relu(t_624)
|
||||
t_626 = self.n_Conv_116(t_625)
|
||||
t_627 = F.relu(t_626)
|
||||
t_627_padded = F.pad(t_627, [1, 1, 1, 1], value=0)
|
||||
t_628 = self.n_Conv_117(t_627_padded)
|
||||
t_629 = F.relu(t_628)
|
||||
t_630 = self.n_Conv_118(t_629)
|
||||
t_631 = torch.add(t_630, t_625)
|
||||
t_632 = F.relu(t_631)
|
||||
t_633 = self.n_Conv_119(t_632)
|
||||
t_634 = F.relu(t_633)
|
||||
t_634_padded = F.pad(t_634, [1, 1, 1, 1], value=0)
|
||||
t_635 = self.n_Conv_120(t_634_padded)
|
||||
t_636 = F.relu(t_635)
|
||||
t_637 = self.n_Conv_121(t_636)
|
||||
t_638 = torch.add(t_637, t_632)
|
||||
t_639 = F.relu(t_638)
|
||||
t_640 = self.n_Conv_122(t_639)
|
||||
t_641 = F.relu(t_640)
|
||||
t_641_padded = F.pad(t_641, [1, 1, 1, 1], value=0)
|
||||
t_642 = self.n_Conv_123(t_641_padded)
|
||||
t_643 = F.relu(t_642)
|
||||
t_644 = self.n_Conv_124(t_643)
|
||||
t_645 = torch.add(t_644, t_639)
|
||||
t_646 = F.relu(t_645)
|
||||
t_647 = self.n_Conv_125(t_646)
|
||||
t_648 = F.relu(t_647)
|
||||
t_648_padded = F.pad(t_648, [1, 1, 1, 1], value=0)
|
||||
t_649 = self.n_Conv_126(t_648_padded)
|
||||
t_650 = F.relu(t_649)
|
||||
t_651 = self.n_Conv_127(t_650)
|
||||
t_652 = torch.add(t_651, t_646)
|
||||
t_653 = F.relu(t_652)
|
||||
t_654 = self.n_Conv_128(t_653)
|
||||
t_655 = F.relu(t_654)
|
||||
t_655_padded = F.pad(t_655, [1, 1, 1, 1], value=0)
|
||||
t_656 = self.n_Conv_129(t_655_padded)
|
||||
t_657 = F.relu(t_656)
|
||||
t_658 = self.n_Conv_130(t_657)
|
||||
t_659 = torch.add(t_658, t_653)
|
||||
t_660 = F.relu(t_659)
|
||||
t_661 = self.n_Conv_131(t_660)
|
||||
t_662 = F.relu(t_661)
|
||||
t_662_padded = F.pad(t_662, [1, 1, 1, 1], value=0)
|
||||
t_663 = self.n_Conv_132(t_662_padded)
|
||||
t_664 = F.relu(t_663)
|
||||
t_665 = self.n_Conv_133(t_664)
|
||||
t_666 = torch.add(t_665, t_660)
|
||||
t_667 = F.relu(t_666)
|
||||
t_668 = self.n_Conv_134(t_667)
|
||||
t_669 = F.relu(t_668)
|
||||
t_669_padded = F.pad(t_669, [1, 1, 1, 1], value=0)
|
||||
t_670 = self.n_Conv_135(t_669_padded)
|
||||
t_671 = F.relu(t_670)
|
||||
t_672 = self.n_Conv_136(t_671)
|
||||
t_673 = torch.add(t_672, t_667)
|
||||
t_674 = F.relu(t_673)
|
||||
t_675 = self.n_Conv_137(t_674)
|
||||
t_676 = F.relu(t_675)
|
||||
t_676_padded = F.pad(t_676, [1, 1, 1, 1], value=0)
|
||||
t_677 = self.n_Conv_138(t_676_padded)
|
||||
t_678 = F.relu(t_677)
|
||||
t_679 = self.n_Conv_139(t_678)
|
||||
t_680 = torch.add(t_679, t_674)
|
||||
t_681 = F.relu(t_680)
|
||||
t_682 = self.n_Conv_140(t_681)
|
||||
t_683 = F.relu(t_682)
|
||||
t_683_padded = F.pad(t_683, [1, 1, 1, 1], value=0)
|
||||
t_684 = self.n_Conv_141(t_683_padded)
|
||||
t_685 = F.relu(t_684)
|
||||
t_686 = self.n_Conv_142(t_685)
|
||||
t_687 = torch.add(t_686, t_681)
|
||||
t_688 = F.relu(t_687)
|
||||
t_689 = self.n_Conv_143(t_688)
|
||||
t_690 = F.relu(t_689)
|
||||
t_690_padded = F.pad(t_690, [1, 1, 1, 1], value=0)
|
||||
t_691 = self.n_Conv_144(t_690_padded)
|
||||
t_692 = F.relu(t_691)
|
||||
t_693 = self.n_Conv_145(t_692)
|
||||
t_694 = torch.add(t_693, t_688)
|
||||
t_695 = F.relu(t_694)
|
||||
t_696 = self.n_Conv_146(t_695)
|
||||
t_697 = F.relu(t_696)
|
||||
t_697_padded = F.pad(t_697, [1, 1, 1, 1], value=0)
|
||||
t_698 = self.n_Conv_147(t_697_padded)
|
||||
t_699 = F.relu(t_698)
|
||||
t_700 = self.n_Conv_148(t_699)
|
||||
t_701 = torch.add(t_700, t_695)
|
||||
t_702 = F.relu(t_701)
|
||||
t_703 = self.n_Conv_149(t_702)
|
||||
t_704 = F.relu(t_703)
|
||||
t_704_padded = F.pad(t_704, [1, 1, 1, 1], value=0)
|
||||
t_705 = self.n_Conv_150(t_704_padded)
|
||||
t_706 = F.relu(t_705)
|
||||
t_707 = self.n_Conv_151(t_706)
|
||||
t_708 = torch.add(t_707, t_702)
|
||||
t_709 = F.relu(t_708)
|
||||
t_710 = self.n_Conv_152(t_709)
|
||||
t_711 = F.relu(t_710)
|
||||
t_711_padded = F.pad(t_711, [1, 1, 1, 1], value=0)
|
||||
t_712 = self.n_Conv_153(t_711_padded)
|
||||
t_713 = F.relu(t_712)
|
||||
t_714 = self.n_Conv_154(t_713)
|
||||
t_715 = torch.add(t_714, t_709)
|
||||
t_716 = F.relu(t_715)
|
||||
t_717 = self.n_Conv_155(t_716)
|
||||
t_718 = F.relu(t_717)
|
||||
t_718_padded = F.pad(t_718, [1, 1, 1, 1], value=0)
|
||||
t_719 = self.n_Conv_156(t_718_padded)
|
||||
t_720 = F.relu(t_719)
|
||||
t_721 = self.n_Conv_157(t_720)
|
||||
t_722 = torch.add(t_721, t_716)
|
||||
t_723 = F.relu(t_722)
|
||||
t_724 = self.n_Conv_158(t_723)
|
||||
t_725 = self.n_Conv_159(t_723)
|
||||
t_726 = F.relu(t_725)
|
||||
t_726_padded = F.pad(t_726, [0, 1, 0, 1], value=0)
|
||||
t_727 = self.n_Conv_160(t_726_padded)
|
||||
t_728 = F.relu(t_727)
|
||||
t_729 = self.n_Conv_161(t_728)
|
||||
t_730 = torch.add(t_729, t_724)
|
||||
t_731 = F.relu(t_730)
|
||||
t_732 = self.n_Conv_162(t_731)
|
||||
t_733 = F.relu(t_732)
|
||||
t_733_padded = F.pad(t_733, [1, 1, 1, 1], value=0)
|
||||
t_734 = self.n_Conv_163(t_733_padded)
|
||||
t_735 = F.relu(t_734)
|
||||
t_736 = self.n_Conv_164(t_735)
|
||||
t_737 = torch.add(t_736, t_731)
|
||||
t_738 = F.relu(t_737)
|
||||
t_739 = self.n_Conv_165(t_738)
|
||||
t_740 = F.relu(t_739)
|
||||
t_740_padded = F.pad(t_740, [1, 1, 1, 1], value=0)
|
||||
t_741 = self.n_Conv_166(t_740_padded)
|
||||
t_742 = F.relu(t_741)
|
||||
t_743 = self.n_Conv_167(t_742)
|
||||
t_744 = torch.add(t_743, t_738)
|
||||
t_745 = F.relu(t_744)
|
||||
t_746 = self.n_Conv_168(t_745)
|
||||
t_747 = self.n_Conv_169(t_745)
|
||||
t_748 = F.relu(t_747)
|
||||
t_748_padded = F.pad(t_748, [0, 1, 0, 1], value=0)
|
||||
t_749 = self.n_Conv_170(t_748_padded)
|
||||
t_750 = F.relu(t_749)
|
||||
t_751 = self.n_Conv_171(t_750)
|
||||
t_752 = torch.add(t_751, t_746)
|
||||
t_753 = F.relu(t_752)
|
||||
t_754 = self.n_Conv_172(t_753)
|
||||
t_755 = F.relu(t_754)
|
||||
t_755_padded = F.pad(t_755, [1, 1, 1, 1], value=0)
|
||||
t_756 = self.n_Conv_173(t_755_padded)
|
||||
t_757 = F.relu(t_756)
|
||||
t_758 = self.n_Conv_174(t_757)
|
||||
t_759 = torch.add(t_758, t_753)
|
||||
t_760 = F.relu(t_759)
|
||||
t_761 = self.n_Conv_175(t_760)
|
||||
t_762 = F.relu(t_761)
|
||||
t_762_padded = F.pad(t_762, [1, 1, 1, 1], value=0)
|
||||
t_763 = self.n_Conv_176(t_762_padded)
|
||||
t_764 = F.relu(t_763)
|
||||
t_765 = self.n_Conv_177(t_764)
|
||||
t_766 = torch.add(t_765, t_760)
|
||||
t_767 = F.relu(t_766)
|
||||
t_768 = self.n_Conv_178(t_767)
|
||||
t_769 = F.avg_pool2d(t_768, kernel_size=t_768.shape[-2:])
|
||||
t_770 = torch.squeeze(t_769, 3)
|
||||
t_770 = torch.squeeze(t_770, 2)
|
||||
t_771 = torch.sigmoid(t_770)
|
||||
return t_771
|
||||
|
||||
def load_state_dict(self, state_dict, **kwargs): # pylint: disable=arguments-differ,unused-argument
|
||||
self.tags = state_dict.get('tags', [])
|
||||
super(DeepDanbooruModel, self).load_state_dict({k: v for k, v in state_dict.items() if k != 'tags'}) # pylint: disable=R1725
|
||||
@@ -0,0 +1,116 @@
|
||||
# source: <https://huggingface.co/deepseek-ai/deepseek-vl2-tiny>
|
||||
# implementation: <https://github.com/deepseek-ai/DeepSeek-VL2/tree/main/deepseek_vl2/serve>
|
||||
"""
|
||||
- run `git clone https://github.com/deepseek-ai/DeepSeek-VL2 repositories/deepseek-vl2 --depth 1`
|
||||
- remove hardcoded `python==3.9` requirement due to obsolete attrdict package dependency
|
||||
- patch transformers due to internal changes as deepseek requires obsolete `transformers==4.38.2`
|
||||
- deepseek requires `xformers`
|
||||
- broken flash_attention
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import importlib
|
||||
from transformers import AutoModelForCausalLM
|
||||
from modules import shared, devices, paths, sd_models
|
||||
|
||||
|
||||
# model_path = "deepseek-ai/deepseek-vl2-small"
|
||||
vl_gpt = None
|
||||
vl_chat_processor = None
|
||||
loaded_repo = None
|
||||
|
||||
|
||||
class fake_attrdict():
|
||||
class AttrDict(dict): # dot notation access to dictionary attributes
|
||||
__getattr__ = dict.get
|
||||
__setattr__ = dict.__setitem__
|
||||
__delattr__ = dict.__delitem__
|
||||
|
||||
|
||||
def load(repo: str):
|
||||
"""Load DeepSeek VL2 model (experimental)."""
|
||||
global vl_gpt, vl_chat_processor, loaded_repo # pylint: disable=global-statement
|
||||
if not shared.cmd_opts.experimental:
|
||||
shared.log.error(f'Caption: type=vlm model="DeepSeek VL2" repo="{repo}" is experimental-only')
|
||||
return False
|
||||
folder = os.path.join(paths.script_path, 'repositories', 'deepseek-vl2')
|
||||
if not os.path.exists(folder):
|
||||
shared.log.error(f'Caption: type=vlm model="DeepSeek VL2" repo="{repo}" deepseek-vl2 repo not found')
|
||||
return False
|
||||
if vl_gpt is None or loaded_repo != repo:
|
||||
sys.modules['attrdict'] = fake_attrdict
|
||||
from transformers.models.llama import modeling_llama
|
||||
modeling_llama.LlamaFlashAttention2 = modeling_llama.LlamaAttention
|
||||
importlib.import_module('repositories.deepseek-vl2.deepseek_vl2')
|
||||
deekseek_vl_models = importlib.import_module('repositories.deepseek-vl2.deepseek_vl2.models')
|
||||
vl_chat_processor = deekseek_vl_models.DeepseekVLV2Processor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
|
||||
vl_gpt = AutoModelForCausalLM.from_pretrained(
|
||||
repo,
|
||||
trust_remote_code=True,
|
||||
cache_dir=shared.opts.hfcache_dir,
|
||||
)
|
||||
vl_gpt.to(dtype=devices.dtype)
|
||||
vl_gpt.eval()
|
||||
loaded_repo = repo
|
||||
shared.log.info(f'Caption: type=vlm model="DeepSeek VL2" repo="{repo}"')
|
||||
sd_models.move_model(vl_gpt, devices.device)
|
||||
return True
|
||||
|
||||
|
||||
def unload():
|
||||
"""Release DeepSeek VL2 model from GPU/memory."""
|
||||
global vl_gpt, vl_chat_processor, loaded_repo # pylint: disable=global-statement
|
||||
if vl_gpt is not None:
|
||||
shared.log.debug(f'DeepSeek unload: model="{loaded_repo}"')
|
||||
sd_models.move_model(vl_gpt, devices.cpu, force=True)
|
||||
vl_gpt = None
|
||||
vl_chat_processor = None
|
||||
loaded_repo = None
|
||||
devices.torch_gc(force=True)
|
||||
else:
|
||||
shared.log.debug('DeepSeek unload: no model loaded')
|
||||
|
||||
|
||||
def predict(question, image, repo):
|
||||
global vl_gpt # pylint: disable=global-statement
|
||||
if not load(repo):
|
||||
return ''
|
||||
|
||||
if len(question) < 2:
|
||||
question = "Describe the image."
|
||||
question = question.replace('<', '').replace('>', '')
|
||||
conversation = [
|
||||
{
|
||||
"role": "<|User|>",
|
||||
"content": f"<image>\n<|ref|>{question}<|/ref|>.",
|
||||
# "images": [image],
|
||||
},
|
||||
{"role": "<|Assistant|>", "content": ""},
|
||||
]
|
||||
|
||||
prepare_inputs = vl_chat_processor(
|
||||
conversations=conversation,
|
||||
images=[image],
|
||||
force_batchify=True,
|
||||
system_prompt=""
|
||||
).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)
|
||||
sd_models.move_model(vl_gpt, devices.device)
|
||||
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
|
||||
)
|
||||
vl_gpt = vl_gpt.to(devices.cpu)
|
||||
answer = vl_chat_processor.tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
|
||||
print('inputs', prepare_inputs['sft_format'][0])
|
||||
print('answer', answer)
|
||||
return answer
|
||||
@@ -0,0 +1,127 @@
|
||||
# based on <https://huggingface.co/fancyfeast/llama-joycaption-alpha-two-hf-llava>
|
||||
|
||||
from dataclasses import dataclass
|
||||
import torch
|
||||
from transformers import AutoProcessor, LlavaForConditionalGeneration
|
||||
from modules import shared, devices, sd_models, model_quant
|
||||
|
||||
|
||||
"""
|
||||
Example prompts
|
||||
Short description: Write a short description of the image.
|
||||
Detailed descriptive: Please provide a detailed description of the image.
|
||||
Descriptive: Write a descriptive caption for this image in a formal tone.
|
||||
Descriptive (Informal): Write a descriptive caption for this image in a casual tone.
|
||||
Training Prompt: Write a stable diffusion prompt for this image.
|
||||
MidJourney: Write a MidJourney prompt for this image.
|
||||
Booru tag list: Write a list of Booru tags for this image.
|
||||
Booru-like tag list: Write a list of Booru-like tags for this image.
|
||||
Art Critic: Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc.
|
||||
Product Listing: Write a caption for this image as though it were a product listing.
|
||||
Social Media Post: Write a caption for this image as if it were being used for a social media post.
|
||||
Extra Options:
|
||||
- If there is a person/character in the image you must refer to them as {name}.
|
||||
- Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).
|
||||
- Include information about lighting.
|
||||
- Include information about camera angle.
|
||||
- Include information about whether there is a watermark or not.
|
||||
- Include information about whether there are JPEG artifacts or not.
|
||||
- If it is a photo you MUST include information about what camera was likely used and details such as aperture, shutter speed, ISO, etc.
|
||||
- Do NOT include anything sexual; keep it PG.
|
||||
- Do NOT mention the image's resolution.
|
||||
- You MUST include information about the subjective aesthetic quality of the image from low to very high.
|
||||
- Include information on the image's composition style, such as leading lines, rule of thirds, or symmetry.
|
||||
- Do NOT mention any text that is in the image.
|
||||
- Specify the depth of field and whether the background is in focus or blurred.
|
||||
- If applicable, mention the likely use of artificial or natural lighting sources.
|
||||
- Do NOT use any ambiguous language.
|
||||
- Include whether the image is sfw, suggestive, or nsfw.
|
||||
- ONLY describe the most important elements of the image.
|
||||
"""
|
||||
|
||||
@dataclass
|
||||
class JoyOptions():
|
||||
repo: str = "fancyfeast/llama-joycaption-alpha-two-hf-llava"
|
||||
temp: float = 0.5
|
||||
top_k: float = 10
|
||||
top_p: float = 0.9
|
||||
max_new_tokens: int = 512
|
||||
sample: bool = True
|
||||
|
||||
def __str__(self):
|
||||
return f'repo="{self.repo}" temp={self.temp} top_k={self.top_k} top_p={self.top_p} sample={self.sample} tokens={self.max_new_tokens}'
|
||||
|
||||
|
||||
processor: AutoProcessor = None
|
||||
llava_model: LlavaForConditionalGeneration = None
|
||||
opts = JoyOptions()
|
||||
|
||||
|
||||
def load(repo: str = None):
|
||||
"""Load JoyCaption model."""
|
||||
global llava_model, processor # pylint: disable=global-statement
|
||||
repo = repo or opts.repo
|
||||
if llava_model is None or opts.repo != repo:
|
||||
opts.repo = repo
|
||||
llava_model = None
|
||||
shared.log.info(f'Caption: type=vlm model="JoyCaption" {str(opts)}')
|
||||
processor = AutoProcessor.from_pretrained(repo, max_pixels=1024*1024, cache_dir=shared.opts.hfcache_dir)
|
||||
quant_args = model_quant.create_config(module='LLM')
|
||||
llava_model = LlavaForConditionalGeneration.from_pretrained(
|
||||
repo,
|
||||
torch_dtype=devices.dtype,
|
||||
cache_dir=shared.opts.hfcache_dir,
|
||||
**quant_args,
|
||||
)
|
||||
llava_model.eval()
|
||||
sd_models.move_model(llava_model, devices.device)
|
||||
|
||||
|
||||
def unload():
|
||||
"""Release JoyCaption model from GPU/memory."""
|
||||
global llava_model, processor # pylint: disable=global-statement
|
||||
if llava_model is not None:
|
||||
shared.log.debug(f'JoyCaption unload: model="{opts.repo}"')
|
||||
sd_models.move_model(llava_model, devices.cpu, force=True)
|
||||
llava_model = None
|
||||
processor = None
|
||||
devices.torch_gc(force=True)
|
||||
else:
|
||||
shared.log.debug('JoyCaption unload: no model loaded')
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def predict(question: str, image, vqa_model: str = None) -> str:
|
||||
opts.max_new_tokens = shared.opts.caption_vlm_max_length
|
||||
load(vqa_model)
|
||||
|
||||
if len(question) < 2:
|
||||
question = "Describe the image."
|
||||
question = question.replace('<', '').replace('>', '')
|
||||
convo = [
|
||||
{ "role": "system", "content": "You are a helpful image captioner." },
|
||||
{ "role": "user", "content": question },
|
||||
]
|
||||
convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
|
||||
inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(devices.device)
|
||||
inputs['pixel_values'] = inputs['pixel_values'].to(devices.dtype)
|
||||
with devices.inference_context():
|
||||
generate_ids = llava_model.generate( # Generate the captions
|
||||
**inputs,
|
||||
# input_ids=inputs['input_ids'],
|
||||
# pixel_values=inputs['pixel_values'],
|
||||
# attention_mask=inputs['attention_mask'],
|
||||
max_new_tokens=opts.max_new_tokens,
|
||||
suppress_tokens=None,
|
||||
use_cache=True,
|
||||
do_sample=opts.sample,
|
||||
temperature=opts.temp,
|
||||
top_k=opts.top_k,
|
||||
top_p=opts.top_p,
|
||||
)[0]
|
||||
generate_ids = generate_ids[inputs['input_ids'].shape[1]:] # Trim off the prompt
|
||||
caption = processor.tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) # Decode the caption
|
||||
if shared.opts.caption_offload:
|
||||
sd_models.move_model(llava_model, devices.cpu, force=True)
|
||||
caption = caption.replace('\n\n', '\n').strip()
|
||||
return caption
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,408 @@
|
||||
# Moondream 3 Preview VLM Implementation
|
||||
# Source: https://huggingface.co/moondream/moondream3-preview
|
||||
# Model: 9.3GB, gated (requires HuggingFace authentication)
|
||||
# Architecture: Mixture-of-Experts (9B total params, 2B active)
|
||||
import os
|
||||
import re
|
||||
import transformers
|
||||
from PIL import Image
|
||||
from modules import shared, devices, sd_models
|
||||
from modules.caption import vqa_detection
|
||||
|
||||
|
||||
# Debug logging - function-based to avoid circular import
|
||||
debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None
|
||||
|
||||
def debug(*args, **kwargs):
|
||||
if debug_enabled:
|
||||
shared.log.trace(*args, **kwargs)
|
||||
|
||||
|
||||
# Global state
|
||||
moondream3_model = None
|
||||
loaded = None
|
||||
image_cache = {} # Cache encoded images for reuse
|
||||
|
||||
|
||||
def get_settings():
|
||||
"""
|
||||
Build settings dict for Moondream 3 API from global VQA options.
|
||||
Moondream 3 accepts: temperature, top_p, max_tokens
|
||||
"""
|
||||
settings = {}
|
||||
if shared.opts.caption_vlm_max_length > 0:
|
||||
settings['max_tokens'] = shared.opts.caption_vlm_max_length
|
||||
if shared.opts.caption_vlm_temperature > 0:
|
||||
settings['temperature'] = shared.opts.caption_vlm_temperature
|
||||
if shared.opts.caption_vlm_top_p > 0:
|
||||
settings['top_p'] = shared.opts.caption_vlm_top_p
|
||||
return settings if settings else None
|
||||
|
||||
|
||||
def load_model(repo: str):
|
||||
"""Load Moondream 3 model."""
|
||||
global moondream3_model, loaded # pylint: disable=global-statement
|
||||
|
||||
if moondream3_model is None or loaded != repo:
|
||||
shared.log.debug(f'Caption load: vlm="{repo}"')
|
||||
moondream3_model = None
|
||||
|
||||
moondream3_model = transformers.AutoModelForCausalLM.from_pretrained(
|
||||
repo,
|
||||
trust_remote_code=True,
|
||||
torch_dtype=devices.dtype,
|
||||
cache_dir=shared.opts.hfcache_dir,
|
||||
)
|
||||
|
||||
moondream3_model.eval()
|
||||
|
||||
# Initialize KV caches before moving to device (they're lazy by default)
|
||||
if hasattr(moondream3_model, '_setup_caches'):
|
||||
moondream3_model._setup_caches() # pylint: disable=protected-access
|
||||
|
||||
# Disable flex_attention decoding (can cause hangs due to torch.compile)
|
||||
if hasattr(moondream3_model, 'model') and hasattr(moondream3_model.model, 'use_flex_decoding'):
|
||||
moondream3_model.model.use_flex_decoding = False
|
||||
|
||||
loaded = repo
|
||||
devices.torch_gc()
|
||||
|
||||
# Move model to active device
|
||||
sd_models.move_model(moondream3_model, devices.device)
|
||||
return moondream3_model
|
||||
|
||||
|
||||
def encode_image(image: Image.Image, cache_key: str = None):
|
||||
"""
|
||||
Encode image for reuse across multiple queries.
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
cache_key: Optional cache key for storing encoded image
|
||||
|
||||
Returns:
|
||||
Encoded image tensor
|
||||
"""
|
||||
if cache_key and cache_key in image_cache:
|
||||
debug(f'VQA caption: handler=moondream3 using cached encoding for cache_key="{cache_key}"')
|
||||
return image_cache[cache_key]
|
||||
|
||||
model = load_model(loaded)
|
||||
|
||||
with devices.inference_context():
|
||||
encoded = model.encode_image(image)
|
||||
|
||||
if cache_key:
|
||||
image_cache[cache_key] = encoded
|
||||
debug(f'VQA caption: handler=moondream3 cached encoding cache_key="{cache_key}" cache_size={len(image_cache)}')
|
||||
|
||||
return encoded
|
||||
|
||||
|
||||
def query(image: Image.Image, question: str, repo: str, stream: bool = False,
|
||||
temperature: float = None, top_p: float = None, max_tokens: int = None,
|
||||
use_cache: bool = False, reasoning: bool = True):
|
||||
"""
|
||||
Visual question answering with optional streaming.
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
question: Question about the image
|
||||
repo: Model repository
|
||||
stream: Enable streaming output (generator)
|
||||
temperature: Sampling temperature (overrides global setting)
|
||||
top_p: Nucleus sampling parameter (overrides global setting)
|
||||
max_tokens: Maximum tokens to generate (overrides global setting)
|
||||
use_cache: Use cached image encoding if available
|
||||
|
||||
Returns:
|
||||
Answer dict or string (or generator if stream=True)
|
||||
"""
|
||||
model = load_model(repo)
|
||||
|
||||
# Build settings - per-call parameters override global settings
|
||||
settings = get_settings() or {}
|
||||
if temperature is not None:
|
||||
settings['temperature'] = temperature
|
||||
if top_p is not None:
|
||||
settings['top_p'] = top_p
|
||||
if max_tokens is not None:
|
||||
settings['max_tokens'] = max_tokens
|
||||
|
||||
debug(f'VQA caption: handler=moondream3 method=query question="{question}" stream={stream} settings={settings}')
|
||||
|
||||
# Use cached encoding if requested
|
||||
if use_cache:
|
||||
cache_key = f"{id(image)}_{question}"
|
||||
image_input = encode_image(image, cache_key)
|
||||
else:
|
||||
image_input = image
|
||||
|
||||
with devices.inference_context():
|
||||
response = model.query(
|
||||
image=image_input,
|
||||
question=question,
|
||||
stream=stream,
|
||||
settings=settings if settings else None,
|
||||
reasoning=reasoning
|
||||
)
|
||||
|
||||
# Log response structure (for non-streaming)
|
||||
if not stream:
|
||||
if isinstance(response, dict):
|
||||
debug(f'VQA caption: handler=moondream3 response_type=dict keys={list(response.keys())}')
|
||||
if 'reasoning' in response:
|
||||
reasoning_text = response['reasoning'].get('text', '')[:100] + '...' if len(response['reasoning'].get('text', '')) > 100 else response['reasoning'].get('text', '')
|
||||
debug(f'VQA caption: handler=moondream3 reasoning="{reasoning_text}"')
|
||||
if 'answer' in response:
|
||||
debug(f'VQA caption: handler=moondream3 answer="{response["answer"]}"')
|
||||
|
||||
return response
|
||||
|
||||
|
||||
def caption(image: Image.Image, repo: str, length: str = 'normal', stream: bool = False,
|
||||
temperature: float = None, top_p: float = None, max_tokens: int = None):
|
||||
"""
|
||||
Generate image captions at different lengths.
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
repo: Model repository
|
||||
length: Caption length - 'short', 'normal', or 'long'
|
||||
stream: Enable streaming output (generator)
|
||||
temperature: Sampling temperature (overrides global setting)
|
||||
top_p: Nucleus sampling parameter (overrides global setting)
|
||||
max_tokens: Maximum tokens to generate (overrides global setting)
|
||||
|
||||
Returns:
|
||||
Caption dict or string (or generator if stream=True)
|
||||
"""
|
||||
model = load_model(repo)
|
||||
|
||||
# Build settings - per-call parameters override global settings
|
||||
settings = get_settings() or {}
|
||||
if temperature is not None:
|
||||
settings['temperature'] = temperature
|
||||
if top_p is not None:
|
||||
settings['top_p'] = top_p
|
||||
if max_tokens is not None:
|
||||
settings['max_tokens'] = max_tokens
|
||||
|
||||
debug(f'VQA caption: handler=moondream3 method=caption length={length} stream={stream} settings={settings}')
|
||||
|
||||
with devices.inference_context():
|
||||
response = model.caption(
|
||||
image,
|
||||
length=length,
|
||||
stream=stream,
|
||||
settings=settings if settings else None
|
||||
)
|
||||
|
||||
# Log response structure (for non-streaming)
|
||||
if not stream and isinstance(response, dict):
|
||||
debug(f'VQA caption: handler=moondream3 response_type=dict keys={list(response.keys())}')
|
||||
|
||||
return response
|
||||
|
||||
|
||||
def point(image: Image.Image, object_name: str, repo: str):
|
||||
"""
|
||||
Identify coordinates of all instances of a specific object in the image.
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
object_name: Name of object to locate
|
||||
repo: Model repository
|
||||
|
||||
Returns:
|
||||
List of (x, y) tuples with coordinates normalized to 0-1 range, or None if not found
|
||||
Example: [(0.733, 0.442), (0.5, 0.6)] for 2 instances
|
||||
"""
|
||||
model = load_model(repo)
|
||||
|
||||
debug(f'VQA caption: handler=moondream3 method=point object_name="{object_name}"')
|
||||
|
||||
with devices.inference_context():
|
||||
result = model.point(image, object_name)
|
||||
|
||||
debug(f'VQA caption: handler=moondream3 point_raw_result="{result}" type={type(result)}')
|
||||
if isinstance(result, dict):
|
||||
debug(f'VQA caption: handler=moondream3 point_raw_result_keys={list(result.keys())}')
|
||||
|
||||
points = vqa_detection.parse_points(result)
|
||||
if points:
|
||||
debug(f'VQA caption: handler=moondream3 point_result={len(points)} points found')
|
||||
return points
|
||||
|
||||
debug('VQA caption: handler=moondream3 point_result=not found')
|
||||
return None
|
||||
|
||||
|
||||
def detect(image: Image.Image, object_name: str, repo: str, max_objects: int = 10):
|
||||
"""
|
||||
Detect all instances of a specific object with bounding boxes.
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
object_name: Name of object to detect
|
||||
repo: Model repository
|
||||
max_objects: Maximum number of objects to return
|
||||
|
||||
Returns:
|
||||
List of detection dicts with keys:
|
||||
- 'bbox': [x1, y1, x2, y2] normalized to 0-1
|
||||
- 'label': Object label
|
||||
- 'confidence': Detection confidence (0-1)
|
||||
Returns empty list if no objects found.
|
||||
"""
|
||||
model = load_model(repo)
|
||||
|
||||
debug(f'VQA caption: handler=moondream3 method=detect object_name="{object_name}" max_objects={max_objects}')
|
||||
|
||||
with devices.inference_context():
|
||||
result = model.detect(image, object_name)
|
||||
|
||||
debug(f'VQA caption: handler=moondream3 detect_raw_result="{result}" type={type(result)}')
|
||||
if isinstance(result, dict):
|
||||
debug(f'VQA caption: handler=moondream3 detect_raw_result_keys={list(result.keys())}')
|
||||
|
||||
detections = vqa_detection.parse_detections(result, object_name, max_objects)
|
||||
debug(f'VQA caption: handler=moondream3 detect_result={len(detections)} objects found')
|
||||
return detections
|
||||
|
||||
|
||||
def predict(question: str, image: Image.Image, repo: str, model_name: str = None, thinking_mode: bool = False,
|
||||
mode: str = None, stream: bool = False, use_cache: bool = False, **kwargs):
|
||||
"""
|
||||
Main entry point for Moondream 3 VQA - auto-detects mode from question.
|
||||
|
||||
Args:
|
||||
question: The question/prompt (e.g., "caption", "where is the cat?", "describe this")
|
||||
image: PIL Image
|
||||
repo: Model repository
|
||||
model_name: Display name for logging
|
||||
thinking_mode: Enable reasoning mode for query
|
||||
mode: Force specific mode ('query', 'caption', 'caption_short', 'caption_long', 'point', 'detect')
|
||||
stream: Enable streaming output (for query/caption)
|
||||
use_cache: Use cached image encoding (for query)
|
||||
**kwargs: Additional parameters (max_objects for detect, etc.)
|
||||
|
||||
Returns:
|
||||
Response string (detection data stored on VQA singleton instance.last_detection_data)
|
||||
(or generator if stream=True for query/caption modes)
|
||||
"""
|
||||
debug(f'VQA caption: handler=moondream3 model_name="{model_name}" repo="{repo}" question="{question}" image_size={image.size if image else None} mode={mode} stream={stream}')
|
||||
|
||||
# Clean question
|
||||
question = question.replace('<', '').replace('>', '').replace('_', ' ') if question else ''
|
||||
|
||||
# Auto-detect mode from question if not specified
|
||||
if mode is None:
|
||||
question_lower = question.lower()
|
||||
|
||||
# Caption detection
|
||||
if question in ['CAPTION', 'caption'] or 'caption' in question_lower:
|
||||
if 'more detailed' in question_lower or 'very long' in question_lower:
|
||||
mode = 'caption_long'
|
||||
elif 'detailed' in question_lower or 'long' in question_lower:
|
||||
mode = 'caption_normal'
|
||||
elif 'short' in question_lower or 'brief' in question_lower:
|
||||
mode = 'caption_short'
|
||||
else:
|
||||
# Default caption mode (matches vqa.py legacy behavior)
|
||||
if question == 'CAPTION':
|
||||
mode = 'caption_short'
|
||||
elif question == 'DETAILED CAPTION':
|
||||
mode = 'caption_normal'
|
||||
elif question == 'MORE DETAILED CAPTION':
|
||||
mode = 'caption_long'
|
||||
else:
|
||||
mode = 'caption_normal'
|
||||
|
||||
# Point detection
|
||||
elif 'where is' in question_lower or 'locate' in question_lower or 'find' in question_lower or 'point' in question_lower:
|
||||
mode = 'point'
|
||||
|
||||
# Object detection
|
||||
elif 'detect' in question_lower or 'bounding box' in question_lower or 'bbox' in question_lower:
|
||||
mode = 'detect'
|
||||
|
||||
# Default to query
|
||||
else:
|
||||
mode = 'query'
|
||||
|
||||
debug(f'VQA caption: handler=moondream3 mode_selected={mode}')
|
||||
|
||||
# Dispatch to appropriate method
|
||||
try:
|
||||
if mode == 'caption_short':
|
||||
response = caption(image, repo, length='short', stream=stream)
|
||||
elif mode == 'caption_long':
|
||||
response = caption(image, repo, length='long', stream=stream)
|
||||
elif mode in ['caption', 'caption_normal']:
|
||||
response = caption(image, repo, length='normal', stream=stream)
|
||||
elif mode == 'point':
|
||||
# Extract object name from question - case insensitive, preserve object names
|
||||
object_name = question
|
||||
for phrase in ['point at', 'where is', 'locate', 'find']:
|
||||
object_name = re.sub(rf'\b{phrase}\b', '', object_name, flags=re.IGNORECASE)
|
||||
object_name = re.sub(r'[?.!,]', '', object_name).strip()
|
||||
object_name = re.sub(r'^\s*the\s+', '', object_name, flags=re.IGNORECASE)
|
||||
debug(f'VQA caption: handler=moondream3 point_extracted_object="{object_name}"')
|
||||
result = point(image, object_name, repo)
|
||||
if result:
|
||||
from modules.caption import vqa
|
||||
vqa.get_instance().last_detection_data = {'points': result}
|
||||
return vqa_detection.format_points_text(result)
|
||||
return "Object not found"
|
||||
elif mode == 'detect':
|
||||
# Extract object name from question - case insensitive
|
||||
object_name = question
|
||||
for phrase in ['detect', 'find all', 'bounding box', 'bbox', 'find']:
|
||||
object_name = re.sub(rf'\b{phrase}\b', '', object_name, flags=re.IGNORECASE)
|
||||
object_name = re.sub(r'[?.!,]', '', object_name).strip()
|
||||
object_name = re.sub(r'^\s*the\s+', '', object_name, flags=re.IGNORECASE)
|
||||
if ' and ' in object_name.lower():
|
||||
object_name = re.split(r'\s+and\s+', object_name, flags=re.IGNORECASE)[0].strip()
|
||||
debug(f'VQA caption: handler=moondream3 detect_extracted_object="{object_name}"')
|
||||
|
||||
results = detect(image, object_name, repo, max_objects=kwargs.get('max_objects', 10))
|
||||
if results:
|
||||
from modules.caption import vqa
|
||||
vqa.get_instance().last_detection_data = {'detections': results}
|
||||
return vqa_detection.format_detections_text(results)
|
||||
return "No objects detected"
|
||||
else: # mode == 'query'
|
||||
if len(question) < 2:
|
||||
question = "Describe this image."
|
||||
response = query(image, question, repo, stream=stream, use_cache=use_cache, reasoning=thinking_mode)
|
||||
|
||||
debug(f'VQA caption: handler=moondream3 response_before_clean="{response}"')
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
from modules import errors
|
||||
errors.display(e, 'Moondream3')
|
||||
return f"Error: {str(e)}"
|
||||
|
||||
|
||||
def clear_cache():
|
||||
"""Clear image encoding cache."""
|
||||
cache_size = len(image_cache)
|
||||
image_cache.clear()
|
||||
debug(f'VQA caption: handler=moondream3 cleared image cache cache_size_was={cache_size}')
|
||||
shared.log.debug(f'Moondream3: Cleared image cache ({cache_size} entries)')
|
||||
|
||||
|
||||
def unload():
|
||||
"""Release Moondream 3 model from GPU/memory."""
|
||||
global moondream3_model, loaded # pylint: disable=global-statement
|
||||
if moondream3_model is not None:
|
||||
shared.log.debug(f'Moondream3 unload: model="{loaded}"')
|
||||
sd_models.move_model(moondream3_model, devices.cpu, force=True)
|
||||
moondream3_model = None
|
||||
loaded = None
|
||||
clear_cache()
|
||||
devices.torch_gc(force=True)
|
||||
else:
|
||||
shared.log.debug('Moondream3 unload: no model loaded')
|
||||
@@ -0,0 +1,333 @@
|
||||
import os
|
||||
import time
|
||||
from collections import namedtuple
|
||||
import threading
|
||||
import re
|
||||
import gradio as gr
|
||||
from PIL import Image
|
||||
from modules import devices, shared, errors, sd_models
|
||||
|
||||
|
||||
debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None
|
||||
debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
|
||||
|
||||
# Per-request overrides for API calls
|
||||
_clip_overrides = None
|
||||
|
||||
|
||||
def get_clip_setting(name):
|
||||
"""Get CLIP setting with per-request override support.
|
||||
|
||||
Args:
|
||||
name: Setting name without 'caption_openclip_' prefix (e.g., 'min_flavors', 'max_length')
|
||||
|
||||
Returns:
|
||||
Override value if set, otherwise the value from shared.opts
|
||||
"""
|
||||
if _clip_overrides is not None:
|
||||
value = _clip_overrides.get(name)
|
||||
if value is not None:
|
||||
return value
|
||||
return getattr(shared.opts, f'caption_openclip_{name}')
|
||||
|
||||
|
||||
def _apply_blip2_fix(model, processor):
|
||||
"""Apply compatibility fix for BLIP2 models with newer transformers versions."""
|
||||
from transformers import AddedToken
|
||||
if not hasattr(model.config, 'num_query_tokens'):
|
||||
return
|
||||
processor.num_query_tokens = model.config.num_query_tokens
|
||||
image_token = AddedToken("<image>", normalized=False, special=True)
|
||||
processor.tokenizer.add_tokens([image_token], special_tokens=True)
|
||||
model.resize_token_embeddings(len(processor.tokenizer), pad_to_multiple_of=64)
|
||||
model.config.image_token_index = len(processor.tokenizer) - 1
|
||||
debug_log(f'CLIP load: applied BLIP2 tokenizer fix num_query_tokens={model.config.num_query_tokens}')
|
||||
|
||||
|
||||
caption_models = {
|
||||
'blip-base': 'Salesforce/blip-image-captioning-base',
|
||||
'blip-large': 'Salesforce/blip-image-captioning-large',
|
||||
'blip2-opt-2.7b': 'Salesforce/blip2-opt-2.7b-coco',
|
||||
'blip2-opt-6.7b': 'Salesforce/blip2-opt-6.7b',
|
||||
'blip2-flip-t5-xl': 'Salesforce/blip2-flan-t5-xl',
|
||||
'blip2-flip-t5-xxl': 'Salesforce/blip2-flan-t5-xxl',
|
||||
}
|
||||
caption_types = [
|
||||
'best',
|
||||
'fast',
|
||||
'classic',
|
||||
'caption',
|
||||
'negative',
|
||||
]
|
||||
clip_models = []
|
||||
ci = None
|
||||
blip_image_eval_size = 384
|
||||
clip_model_name = 'ViT-L/14'
|
||||
Category = namedtuple("Category", ["name", "topn", "items"])
|
||||
re_topn = re.compile(r"\.top(\d+)\.")
|
||||
load_lock = threading.Lock()
|
||||
|
||||
|
||||
class BatchWriter:
|
||||
def __init__(self, folder, mode='w'):
|
||||
self.folder = folder
|
||||
self.csv = None
|
||||
self.file = None
|
||||
self.mode = mode
|
||||
|
||||
def add(self, file, prompt):
|
||||
txt_file = os.path.splitext(file)[0] + ".txt"
|
||||
if self.mode == 'a':
|
||||
prompt = '\n' + prompt
|
||||
with open(os.path.join(self.folder, txt_file), self.mode, encoding='utf-8') as f:
|
||||
f.write(prompt)
|
||||
|
||||
def close(self):
|
||||
if self.file is not None:
|
||||
self.file.close()
|
||||
|
||||
|
||||
def update_caption_params():
|
||||
if ci is not None:
|
||||
ci.caption_max_length = get_clip_setting('max_length')
|
||||
ci.chunk_size = get_clip_setting('chunk_size')
|
||||
ci.flavor_intermediate_count = get_clip_setting('flavor_count')
|
||||
ci.clip_offload = shared.opts.caption_offload
|
||||
ci.caption_offload = shared.opts.caption_offload
|
||||
|
||||
|
||||
|
||||
def get_clip_models():
|
||||
return clip_models
|
||||
|
||||
|
||||
def refresh_clip_models():
|
||||
global clip_models # pylint: disable=global-statement
|
||||
import open_clip
|
||||
models = sorted(open_clip.list_pretrained())
|
||||
shared.log.debug(f'Caption: pkg=openclip version={open_clip.__version__} models={len(models)}')
|
||||
clip_models = ['/'.join(x) for x in models]
|
||||
return clip_models
|
||||
|
||||
|
||||
def load_captioner(clip_model, blip_model):
|
||||
from installer import install
|
||||
install('clip_interrogator==0.6.0')
|
||||
import clip_interrogator
|
||||
clip_interrogator.clip_interrogator.CAPTION_MODELS = caption_models
|
||||
global ci # pylint: disable=global-statement
|
||||
if ci is None:
|
||||
t0 = time.time()
|
||||
device = devices.get_optimal_device()
|
||||
cache_path = shared.opts.clip_models_path
|
||||
shared.log.info(f'CLIP load: clip="{clip_model}" blip="{blip_model}" device={device}')
|
||||
debug_log(f'CLIP load: cache_path="{cache_path}" max_length={shared.opts.caption_openclip_max_length} chunk_size={shared.opts.caption_openclip_chunk_size} flavor_count={shared.opts.caption_openclip_flavor_count} offload={shared.opts.caption_offload}')
|
||||
captioner_config = clip_interrogator.Config(
|
||||
device=device,
|
||||
cache_path=cache_path,
|
||||
clip_model_name=clip_model,
|
||||
caption_model_name=blip_model,
|
||||
quiet=True,
|
||||
caption_max_length=shared.opts.caption_openclip_max_length,
|
||||
chunk_size=shared.opts.caption_openclip_chunk_size,
|
||||
flavor_intermediate_count=shared.opts.caption_openclip_flavor_count,
|
||||
clip_offload=shared.opts.caption_offload,
|
||||
caption_offload=shared.opts.caption_offload,
|
||||
)
|
||||
ci = clip_interrogator.Interrogator(captioner_config)
|
||||
|
||||
if blip_model.startswith('blip2-'):
|
||||
_apply_blip2_fix(ci.caption_model, ci.caption_processor)
|
||||
shared.log.debug(f'CLIP load: time={time.time()-t0:.2f}')
|
||||
elif clip_model != ci.config.clip_model_name or blip_model != ci.config.caption_model_name:
|
||||
t0 = time.time()
|
||||
if clip_model != ci.config.clip_model_name:
|
||||
shared.log.info(f'CLIP load: clip="{clip_model}" reloading')
|
||||
debug_log(f'CLIP load: previous clip="{ci.config.clip_model_name}"')
|
||||
ci.config.clip_model_name = clip_model
|
||||
ci.config.clip_model = None
|
||||
ci.load_clip_model()
|
||||
ci.clip_offloaded = True # Reset flag so _prepare_clip() will move model to device
|
||||
if blip_model != ci.config.caption_model_name:
|
||||
shared.log.info(f'CLIP load: blip="{blip_model}" reloading')
|
||||
debug_log(f'CLIP load: previous blip="{ci.config.caption_model_name}"')
|
||||
ci.config.caption_model_name = blip_model
|
||||
ci.config.caption_model = None
|
||||
ci.load_caption_model()
|
||||
ci.caption_offloaded = True # Reset flag so _prepare_caption() will move model to device
|
||||
if blip_model.startswith('blip2-'):
|
||||
_apply_blip2_fix(ci.caption_model, ci.caption_processor)
|
||||
shared.log.debug(f'CLIP load: time={time.time()-t0:.2f}')
|
||||
else:
|
||||
debug_log(f'CLIP: models already loaded clip="{clip_model}" blip="{blip_model}"')
|
||||
|
||||
|
||||
def unload_clip_model():
|
||||
if ci is not None and shared.opts.caption_offload:
|
||||
shared.log.debug('CLIP unload: offloading models to CPU')
|
||||
sd_models.move_model(ci.caption_model, devices.cpu)
|
||||
sd_models.move_model(ci.clip_model, devices.cpu)
|
||||
ci.caption_offloaded = True
|
||||
ci.clip_offloaded = True
|
||||
devices.torch_gc()
|
||||
debug_log('CLIP unload: complete')
|
||||
|
||||
|
||||
def caption(image, mode, base_caption=None):
|
||||
if isinstance(image, list):
|
||||
image = image[0] if len(image) > 0 else None
|
||||
if isinstance(image, dict) and 'name' in image:
|
||||
image = Image.open(image['name'])
|
||||
if image is None:
|
||||
return ''
|
||||
image = image.convert("RGB")
|
||||
t0 = time.time()
|
||||
min_flavors = get_clip_setting('min_flavors')
|
||||
max_flavors = get_clip_setting('max_flavors')
|
||||
debug_log(f'CLIP: mode="{mode}" image_size={image.size} caption={base_caption is not None} min_flavors={min_flavors} max_flavors={max_flavors}')
|
||||
# NOTE: Method names like .interrogate(), .interrogate_classic(), etc. come from the external
|
||||
# clip-interrogator library (https://github.com/pharmapsychotic/clip-interrogator) and cannot be renamed.
|
||||
if mode == 'best':
|
||||
prompt = ci.interrogate(image, caption=base_caption, min_flavors=min_flavors, max_flavors=max_flavors)
|
||||
elif mode == 'caption':
|
||||
prompt = ci.generate_caption(image) if base_caption is None else base_caption
|
||||
elif mode == 'classic':
|
||||
prompt = ci.interrogate_classic(image, caption=base_caption, max_flavors=max_flavors)
|
||||
elif mode == 'fast':
|
||||
prompt = ci.interrogate_fast(image, caption=base_caption, max_flavors=max_flavors)
|
||||
elif mode == 'negative':
|
||||
prompt = ci.interrogate_negative(image, max_flavors=max_flavors)
|
||||
else:
|
||||
raise RuntimeError(f"Unknown mode {mode}")
|
||||
debug_log(f'CLIP: mode="{mode}" time={time.time()-t0:.2f} result="{prompt[:100]}..."' if len(prompt) > 100 else f'CLIP: mode="{mode}" time={time.time()-t0:.2f} result="{prompt}"')
|
||||
return prompt
|
||||
|
||||
|
||||
|
||||
def caption_image(image, clip_model, blip_model, mode, overrides=None):
|
||||
global _clip_overrides # pylint: disable=global-statement
|
||||
jobid = shared.state.begin('Caption CLiP')
|
||||
t0 = time.time()
|
||||
shared.log.info(f'CLIP: mode="{mode}" clip="{clip_model}" blip="{blip_model}" image_size={image.size if image else None}')
|
||||
if overrides:
|
||||
debug_log(f'CLIP: overrides={overrides}')
|
||||
try:
|
||||
# Set per-request overrides
|
||||
_clip_overrides = overrides
|
||||
if shared.sd_loaded:
|
||||
from modules.sd_models import apply_balanced_offload # prevent circular import
|
||||
apply_balanced_offload(shared.sd_model)
|
||||
debug_log('CLIP: applied balanced offload to sd_model')
|
||||
load_captioner(clip_model, blip_model)
|
||||
# Apply overrides to loaded captioner
|
||||
update_caption_params()
|
||||
image = image.convert('RGB')
|
||||
prompt = caption(image, mode)
|
||||
if shared.opts.caption_offload:
|
||||
unload_clip_model()
|
||||
devices.torch_gc()
|
||||
shared.log.debug(f'CLIP: complete time={time.time()-t0:.2f}')
|
||||
except Exception as e:
|
||||
prompt = f"Exception {type(e)}"
|
||||
shared.log.error(f'CLIP: {e}')
|
||||
errors.display(e, 'Caption')
|
||||
finally:
|
||||
# Clear per-request overrides
|
||||
_clip_overrides = None
|
||||
shared.state.end(jobid)
|
||||
return prompt
|
||||
|
||||
|
||||
|
||||
def caption_batch(batch_files, batch_folder, batch_str, clip_model, blip_model, mode, write, append, recursive):
|
||||
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):
|
||||
from modules.files_cache import list_files
|
||||
files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp', '.jxl'], recursive=recursive))
|
||||
if len(files) == 0:
|
||||
shared.log.warning('CLIP batch: no images found')
|
||||
return ''
|
||||
t0 = time.time()
|
||||
shared.log.info(f'CLIP batch: mode="{mode}" images={len(files)} clip="{clip_model}" blip="{blip_model}" write={write} append={append}')
|
||||
debug_log(f'CLIP batch: recursive={recursive} files={files[:5]}{"..." if len(files) > 5 else ""}')
|
||||
jobid = shared.state.begin('Caption batch')
|
||||
prompts = []
|
||||
|
||||
load_captioner(clip_model, blip_model)
|
||||
if write:
|
||||
file_mode = 'w' if not append else 'a'
|
||||
writer = BatchWriter(os.path.dirname(files[0]), mode=file_mode)
|
||||
debug_log(f'CLIP batch: writing to "{os.path.dirname(files[0])}" mode="{file_mode}"')
|
||||
import rich.progress as rp
|
||||
pbar = rp.Progress(rp.TextColumn('[cyan]Caption:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
|
||||
with pbar:
|
||||
task = pbar.add_task(total=len(files), description='starting...')
|
||||
for file in files:
|
||||
pbar.update(task, advance=1, description=file)
|
||||
try:
|
||||
if shared.state.interrupted:
|
||||
shared.log.info('CLIP batch: interrupted')
|
||||
break
|
||||
image = Image.open(file).convert('RGB')
|
||||
prompt = caption(image, mode)
|
||||
prompts.append(prompt)
|
||||
if write:
|
||||
writer.add(file, prompt)
|
||||
except OSError as e:
|
||||
shared.log.error(f'CLIP batch: file="{file}" error={e}')
|
||||
if write:
|
||||
writer.close()
|
||||
ci.config.quiet = False
|
||||
unload_clip_model()
|
||||
shared.state.end(jobid)
|
||||
shared.log.info(f'CLIP batch: complete images={len(prompts)} time={time.time()-t0:.2f}')
|
||||
return '\n\n'.join(prompts)
|
||||
|
||||
|
||||
|
||||
def analyze_image(image, clip_model, blip_model):
|
||||
t0 = time.time()
|
||||
shared.log.info(f'CLIP analyze: clip="{clip_model}" blip="{blip_model}" image_size={image.size if image else None}')
|
||||
load_captioner(clip_model, blip_model)
|
||||
image = image.convert('RGB')
|
||||
image_features = ci.image_to_features(image)
|
||||
debug_log(f'CLIP analyze: features shape={image_features.shape if hasattr(image_features, "shape") else "unknown"}')
|
||||
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(sorted(zip(top_mediums, ci.similarities(image_features, top_mediums)), key=lambda x: x[1], reverse=True))
|
||||
artist_ranks = dict(sorted(zip(top_artists, ci.similarities(image_features, top_artists)), key=lambda x: x[1], reverse=True))
|
||||
movement_ranks = dict(sorted(zip(top_movements, ci.similarities(image_features, top_movements)), key=lambda x: x[1], reverse=True))
|
||||
trending_ranks = dict(sorted(zip(top_trendings, ci.similarities(image_features, top_trendings)), key=lambda x: x[1], reverse=True))
|
||||
flavor_ranks = dict(sorted(zip(top_flavors, ci.similarities(image_features, top_flavors)), key=lambda x: x[1], reverse=True))
|
||||
shared.log.debug(f'CLIP analyze: complete time={time.time()-t0:.2f}')
|
||||
|
||||
# Format labels as text
|
||||
def format_category(name, ranks):
|
||||
lines = [f"{name}:"]
|
||||
for item, score in ranks.items():
|
||||
lines.append(f" • {item} - {score*100:.1f}%")
|
||||
return '\n'.join(lines)
|
||||
|
||||
formatted_text = '\n\n'.join([
|
||||
format_category("Medium", medium_ranks),
|
||||
format_category("Artist", artist_ranks),
|
||||
format_category("Movement", movement_ranks),
|
||||
format_category("Trending", trending_ranks),
|
||||
format_category("Flavor", flavor_ranks),
|
||||
])
|
||||
|
||||
return [
|
||||
gr.update(value=medium_ranks, visible=True),
|
||||
gr.update(value=artist_ranks, visible=True),
|
||||
gr.update(value=movement_ranks, visible=True),
|
||||
gr.update(value=trending_ranks, visible=True),
|
||||
gr.update(value=flavor_ranks, visible=True),
|
||||
gr.update(value=formatted_text, visible=True), # New text output for the textbox
|
||||
]
|
||||
@@ -0,0 +1,79 @@
|
||||
# Unified Tagger Interface - Dispatches to WaifuDiffusion or DeepBooru based on model selection
|
||||
# Provides a common interface for the Booru Tags tab
|
||||
|
||||
from modules import shared
|
||||
|
||||
DEEPBOORU_MODEL = "DeepBooru"
|
||||
|
||||
|
||||
def get_models() -> list:
|
||||
"""Return combined list: DeepBooru + WaifuDiffusion models."""
|
||||
from modules.caption import waifudiffusion
|
||||
return [DEEPBOORU_MODEL] + waifudiffusion.get_models()
|
||||
|
||||
|
||||
def refresh_models() -> list:
|
||||
"""Refresh and return all models."""
|
||||
return get_models()
|
||||
|
||||
|
||||
def is_deepbooru(model_name: str) -> bool:
|
||||
"""Check if selected model is DeepBooru."""
|
||||
return model_name == DEEPBOORU_MODEL
|
||||
|
||||
|
||||
def load_model(model_name: str) -> bool:
|
||||
"""Load appropriate backend."""
|
||||
if is_deepbooru(model_name):
|
||||
from modules.caption import deepbooru
|
||||
return deepbooru.load_model()
|
||||
else:
|
||||
from modules.caption import waifudiffusion
|
||||
return waifudiffusion.load_model(model_name)
|
||||
|
||||
|
||||
def unload_model():
|
||||
"""Unload both backends to ensure memory is freed."""
|
||||
from modules.caption import deepbooru, waifudiffusion
|
||||
deepbooru.unload_model()
|
||||
waifudiffusion.unload_model()
|
||||
|
||||
|
||||
def tag(image, model_name: str = None, **kwargs) -> str:
|
||||
"""Unified tagging - dispatch to correct backend.
|
||||
|
||||
Args:
|
||||
image: PIL Image to tag
|
||||
model_name: Model to use (DeepBooru or WaifuDiffusion model name)
|
||||
**kwargs: Additional arguments passed to the backend
|
||||
|
||||
Returns:
|
||||
Formatted tag string
|
||||
"""
|
||||
if model_name is None:
|
||||
model_name = shared.opts.waifudiffusion_model
|
||||
|
||||
if is_deepbooru(model_name):
|
||||
from modules.caption import deepbooru
|
||||
return deepbooru.tag(image, **kwargs)
|
||||
else:
|
||||
from modules.caption import waifudiffusion
|
||||
return waifudiffusion.tag(image, model_name=model_name, **kwargs)
|
||||
|
||||
|
||||
def batch(model_name: str, **kwargs) -> str:
|
||||
"""Unified batch processing.
|
||||
|
||||
Args:
|
||||
model_name: Model to use (DeepBooru or WaifuDiffusion model name)
|
||||
**kwargs: Additional arguments passed to the backend
|
||||
|
||||
Returns:
|
||||
Combined tag results
|
||||
"""
|
||||
if is_deepbooru(model_name):
|
||||
from modules.caption import deepbooru
|
||||
return deepbooru.batch(model_name=model_name, **kwargs)
|
||||
else:
|
||||
from modules.caption import waifudiffusion
|
||||
return waifudiffusion.batch(model_name=model_name, **kwargs)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,207 @@
|
||||
# VQA Detection Utilities
|
||||
# Parsing, formatting, and drawing functions for detection results (points, bboxes, gaze)
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from modules import shared
|
||||
|
||||
|
||||
def parse_points(result) -> list:
|
||||
"""Parse and validate point coordinates from model result.
|
||||
|
||||
Args:
|
||||
result: Model output, typically dict with 'points' key or list of coordinates
|
||||
|
||||
Returns:
|
||||
List of (x, y) tuples with coordinates clamped to 0-1 range.
|
||||
"""
|
||||
points = []
|
||||
|
||||
# Dict format: {'points': [{'x': 0.5, 'y': 0.5}, ...]}
|
||||
if isinstance(result, dict) and 'points' in result:
|
||||
points_list = result['points']
|
||||
if points_list and len(points_list) > 0:
|
||||
for point_data in points_list:
|
||||
if isinstance(point_data, dict) and 'x' in point_data and 'y' in point_data:
|
||||
x = max(0.0, min(1.0, float(point_data['x'])))
|
||||
y = max(0.0, min(1.0, float(point_data['y'])))
|
||||
points.append((x, y))
|
||||
|
||||
# Fallback for simple [x, y] format
|
||||
elif isinstance(result, (list, tuple)) and len(result) == 2:
|
||||
try:
|
||||
x = max(0.0, min(1.0, float(result[0])))
|
||||
y = max(0.0, min(1.0, float(result[1])))
|
||||
points.append((x, y))
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
return points
|
||||
|
||||
|
||||
def parse_detections(result, label: str, max_objects: int = None) -> list:
|
||||
"""Parse and validate detection bboxes from model result.
|
||||
|
||||
Args:
|
||||
result: Model output, typically dict with 'objects' key
|
||||
label: Label to assign to detected objects
|
||||
max_objects: Maximum number of objects to return (None for all)
|
||||
|
||||
Returns:
|
||||
List of {'bbox': [x1,y1,x2,y2], 'label': str, 'confidence': float}
|
||||
with coordinates clamped to 0-1 range.
|
||||
"""
|
||||
detections = []
|
||||
|
||||
if isinstance(result, dict) and 'objects' in result:
|
||||
objects = result['objects']
|
||||
if max_objects is not None:
|
||||
objects = objects[:max_objects]
|
||||
|
||||
for obj in objects:
|
||||
if all(k in obj for k in ['x_min', 'y_min', 'x_max', 'y_max']):
|
||||
bbox = [
|
||||
max(0.0, min(1.0, float(obj['x_min']))),
|
||||
max(0.0, min(1.0, float(obj['y_min']))),
|
||||
max(0.0, min(1.0, float(obj['x_max']))),
|
||||
max(0.0, min(1.0, float(obj['y_max'])))
|
||||
]
|
||||
detections.append({
|
||||
'bbox': bbox,
|
||||
'label': label,
|
||||
'confidence': obj.get('confidence', 1.0)
|
||||
})
|
||||
|
||||
return detections
|
||||
|
||||
|
||||
def format_points_text(points: list) -> str:
|
||||
"""Format point coordinates as human-readable text.
|
||||
|
||||
Args:
|
||||
points: List of (x, y) tuples with normalized coordinates
|
||||
|
||||
Returns:
|
||||
Formatted text string describing the points.
|
||||
"""
|
||||
if not points:
|
||||
return "Object not found"
|
||||
|
||||
if len(points) == 1:
|
||||
return f"Found at: ({points[0][0]:.3f}, {points[0][1]:.3f})"
|
||||
|
||||
lines = [f"Found {len(points)} instances:"]
|
||||
for i, (x, y) in enumerate(points, 1):
|
||||
lines.append(f" {i}. ({x:.3f}, {y:.3f})")
|
||||
return '\n'.join(lines)
|
||||
|
||||
|
||||
def format_detections_text(detections: list, include_confidence: bool = True) -> str:
|
||||
"""Format detections with bboxes as human-readable text.
|
||||
|
||||
Args:
|
||||
detections: List of detection dicts with 'bbox', 'label', 'confidence'
|
||||
include_confidence: Whether to include confidence scores in output
|
||||
|
||||
Returns:
|
||||
Formatted text string describing the detections.
|
||||
"""
|
||||
if not detections:
|
||||
return "No objects detected"
|
||||
|
||||
lines = []
|
||||
for det in detections:
|
||||
bbox = det['bbox']
|
||||
label = det.get('label', 'object')
|
||||
confidence = det.get('confidence', 1.0)
|
||||
|
||||
if include_confidence and confidence < 1.0:
|
||||
lines.append(f"{label}: [{bbox[0]:.3f}, {bbox[1]:.3f}, {bbox[2]:.3f}, {bbox[3]:.3f}] (confidence: {confidence:.2f})")
|
||||
else:
|
||||
lines.append(f"{label}: [{bbox[0]:.3f}, {bbox[1]:.3f}, {bbox[2]:.3f}, {bbox[3]:.3f}]")
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
|
||||
def calculate_eye_position(face_bbox: dict) -> tuple:
|
||||
"""Calculate approximate eye position from face bounding box.
|
||||
|
||||
Args:
|
||||
face_bbox: Dict with 'x_min', 'y_min', 'x_max', 'y_max' keys
|
||||
|
||||
Returns:
|
||||
(eye_x, eye_y) tuple with normalized coordinates.
|
||||
"""
|
||||
eye_x = (face_bbox['x_min'] + face_bbox['x_max']) / 2
|
||||
eye_y = face_bbox['y_min'] + (face_bbox['y_max'] - face_bbox['y_min']) * 0.3 # Approximate eye level
|
||||
return (eye_x, eye_y)
|
||||
|
||||
|
||||
def draw_bounding_boxes(image: Image.Image, detections: list, points: list = None) -> Image.Image:
|
||||
"""
|
||||
Draw bounding boxes and/or points on an image.
|
||||
|
||||
Args:
|
||||
image: PIL Image to annotate
|
||||
detections: List of detection dicts with format:
|
||||
[{'label': str, 'bbox': [x1, y1, x2, y2], 'confidence': float}, ...]
|
||||
where coordinates are normalized 0-1
|
||||
points: Optional list of (x, y) tuples with normalized 0-1 coordinates
|
||||
|
||||
Returns:
|
||||
Annotated PIL Image with boxes and labels drawn, or None if no annotations
|
||||
"""
|
||||
if not detections and not points:
|
||||
return None
|
||||
|
||||
# Create a copy to avoid modifying original
|
||||
annotated = image.copy()
|
||||
draw = ImageDraw.Draw(annotated)
|
||||
width, height = image.size
|
||||
|
||||
# Try to load a font, fall back to default if unavailable
|
||||
try:
|
||||
font_size = max(12, int(min(width, height) * 0.02))
|
||||
font_path = shared.opts.font or "javascript/notosans-nerdfont-regular.ttf"
|
||||
font = ImageFont.truetype(font_path, size=font_size)
|
||||
except Exception:
|
||||
font = ImageFont.load_default()
|
||||
|
||||
# Draw bounding boxes
|
||||
if detections:
|
||||
colors = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#FF00FF', '#00FFFF', '#FFA500', '#800080']
|
||||
for idx, det in enumerate(detections):
|
||||
bbox = det['bbox']
|
||||
label = det.get('label', 'object')
|
||||
confidence = det.get('confidence', 1.0)
|
||||
|
||||
# Convert normalized coordinates to pixel coordinates
|
||||
x1 = int(bbox[0] * width)
|
||||
y1 = int(bbox[1] * height)
|
||||
x2 = int(bbox[2] * width)
|
||||
y2 = int(bbox[3] * height)
|
||||
|
||||
# Choose color
|
||||
color = colors[idx % len(colors)]
|
||||
|
||||
# Draw box
|
||||
draw.rectangle([x1, y1, x2, y2], outline=color, width=max(2, int(min(width, height) * 0.003)))
|
||||
|
||||
# Draw label with background
|
||||
label_text = f"{label} {confidence:.2f}" if confidence < 1.0 else label
|
||||
bbox_font = draw.textbbox((x1, y1), label_text, font=font)
|
||||
text_width = bbox_font[2] - bbox_font[0]
|
||||
text_height = bbox_font[3] - bbox_font[1]
|
||||
draw.rectangle([x1, y1 - text_height - 4, x1 + text_width + 4, y1], fill=color)
|
||||
draw.text((x1 + 2, y1 - text_height - 2), label_text, fill='white', font=font)
|
||||
|
||||
# Draw points
|
||||
if points:
|
||||
point_radius = max(3, int(min(width, height) * 0.01))
|
||||
for px, py in points:
|
||||
x = int(px * width)
|
||||
y = int(py * height)
|
||||
# Draw point as a circle
|
||||
draw.ellipse([x - point_radius, y - point_radius, x + point_radius, y + point_radius],
|
||||
fill='#FF0000', outline='#FFFFFF', width=2)
|
||||
|
||||
return annotated
|
||||
@@ -0,0 +1,544 @@
|
||||
# WaifuDiffusion Tagger - ONNX-based anime/illustration tagging
|
||||
# Based on SmilingWolf's tagger models: https://huggingface.co/SmilingWolf
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import threading
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from modules import shared, devices, errors
|
||||
|
||||
|
||||
# Debug logging - enable with SD_CAPTION_DEBUG environment variable
|
||||
debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None
|
||||
debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
|
||||
|
||||
re_special = re.compile(r'([\\()])')
|
||||
load_lock = threading.Lock()
|
||||
|
||||
# WaifuDiffusion model repository mappings
|
||||
WAIFUDIFFUSION_MODELS = {
|
||||
# v3 models (latest, recommended)
|
||||
"wd-eva02-large-tagger-v3": "SmilingWolf/wd-eva02-large-tagger-v3",
|
||||
"wd-vit-tagger-v3": "SmilingWolf/wd-vit-tagger-v3",
|
||||
"wd-convnext-tagger-v3": "SmilingWolf/wd-convnext-tagger-v3",
|
||||
"wd-swinv2-tagger-v3": "SmilingWolf/wd-swinv2-tagger-v3",
|
||||
# v2 models
|
||||
"wd-v1-4-moat-tagger-v2": "SmilingWolf/wd-v1-4-moat-tagger-v2",
|
||||
"wd-v1-4-swinv2-tagger-v2": "SmilingWolf/wd-v1-4-swinv2-tagger-v2",
|
||||
"wd-v1-4-convnext-tagger-v2": "SmilingWolf/wd-v1-4-convnext-tagger-v2",
|
||||
"wd-v1-4-convnextv2-tagger-v2": "SmilingWolf/wd-v1-4-convnextv2-tagger-v2",
|
||||
"wd-v1-4-vit-tagger-v2": "SmilingWolf/wd-v1-4-vit-tagger-v2",
|
||||
}
|
||||
|
||||
# Tag categories from selected_tags.csv
|
||||
CATEGORY_GENERAL = 0
|
||||
CATEGORY_CHARACTER = 4
|
||||
CATEGORY_RATING = 9
|
||||
|
||||
|
||||
class WaifuDiffusionTagger:
|
||||
"""WaifuDiffusion Tagger using ONNX inference."""
|
||||
|
||||
def __init__(self):
|
||||
self.session = None
|
||||
self.tags = None
|
||||
self.tag_categories = None
|
||||
self.model_name = None
|
||||
self.model_path = None
|
||||
self.image_size = 448 # Standard for WD models
|
||||
|
||||
def load(self, model_name: str = None):
|
||||
"""Load the ONNX model and tags from HuggingFace."""
|
||||
import huggingface_hub
|
||||
|
||||
if model_name is None:
|
||||
model_name = shared.opts.waifudiffusion_model
|
||||
if model_name not in WAIFUDIFFUSION_MODELS:
|
||||
shared.log.error(f'WaifuDiffusion: unknown model "{model_name}"')
|
||||
return False
|
||||
|
||||
with load_lock:
|
||||
if self.session is not None and self.model_name == model_name:
|
||||
debug_log(f'WaifuDiffusion: model already loaded model="{model_name}"')
|
||||
return True # Already loaded
|
||||
|
||||
# Unload previous model if different
|
||||
if self.model_name != model_name and self.session is not None:
|
||||
debug_log(f'WaifuDiffusion: switching model from "{self.model_name}" to "{model_name}"')
|
||||
self.unload()
|
||||
|
||||
repo_id = WAIFUDIFFUSION_MODELS[model_name]
|
||||
t0 = time.time()
|
||||
shared.log.info(f'WaifuDiffusion load: model="{model_name}" repo="{repo_id}"')
|
||||
|
||||
try:
|
||||
# Download only ONNX model and tags CSV (skip safetensors/msgpack variants)
|
||||
debug_log(f'WaifuDiffusion load: downloading from HuggingFace cache_dir="{shared.opts.hfcache_dir}"')
|
||||
self.model_path = huggingface_hub.snapshot_download(
|
||||
repo_id,
|
||||
cache_dir=shared.opts.hfcache_dir,
|
||||
allow_patterns=["model.onnx", "selected_tags.csv"],
|
||||
)
|
||||
debug_log(f'WaifuDiffusion load: model_path="{self.model_path}"')
|
||||
|
||||
# Load ONNX model
|
||||
model_file = os.path.join(self.model_path, "model.onnx")
|
||||
if not os.path.exists(model_file):
|
||||
shared.log.error(f'WaifuDiffusion load: model file not found: {model_file}')
|
||||
return False
|
||||
|
||||
import onnxruntime as ort
|
||||
|
||||
debug_log(f'WaifuDiffusion load: onnxruntime version={ort.__version__}')
|
||||
|
||||
self.session = ort.InferenceSession(model_file, providers=devices.onnx)
|
||||
self.model_name = model_name
|
||||
|
||||
# Get actual providers used
|
||||
actual_providers = self.session.get_providers()
|
||||
debug_log(f'WaifuDiffusion load: active providers={actual_providers}')
|
||||
|
||||
# Load tags from CSV
|
||||
self._load_tags()
|
||||
|
||||
load_time = time.time() - t0
|
||||
shared.log.debug(f'WaifuDiffusion load: time={load_time:.2f} tags={len(self.tags)}')
|
||||
debug_log(f'WaifuDiffusion load: input_name={self.session.get_inputs()[0].name} output_name={self.session.get_outputs()[0].name}')
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
shared.log.error(f'WaifuDiffusion load: failed error={e}')
|
||||
errors.display(e, 'WaifuDiffusion load')
|
||||
self.unload()
|
||||
return False
|
||||
|
||||
def _load_tags(self):
|
||||
"""Load tags and categories from selected_tags.csv."""
|
||||
import csv
|
||||
|
||||
csv_path = os.path.join(self.model_path, "selected_tags.csv")
|
||||
if not os.path.exists(csv_path):
|
||||
shared.log.error(f'WaifuDiffusion load: tags file not found: {csv_path}')
|
||||
return
|
||||
|
||||
self.tags = []
|
||||
self.tag_categories = []
|
||||
|
||||
with open(csv_path, 'r', encoding='utf-8') as f:
|
||||
reader = csv.DictReader(f)
|
||||
for row in reader:
|
||||
self.tags.append(row['name'])
|
||||
self.tag_categories.append(int(row['category']))
|
||||
|
||||
# Count tags by category
|
||||
category_counts = {}
|
||||
for cat in self.tag_categories:
|
||||
category_counts[cat] = category_counts.get(cat, 0) + 1
|
||||
debug_log(f'WaifuDiffusion load: tag categories={category_counts}')
|
||||
|
||||
def unload(self):
|
||||
"""Unload the model and free resources."""
|
||||
if self.session is not None:
|
||||
shared.log.debug(f'WaifuDiffusion unload: model="{self.model_name}"')
|
||||
self.session = None
|
||||
self.tags = None
|
||||
self.tag_categories = None
|
||||
self.model_name = None
|
||||
self.model_path = None
|
||||
devices.torch_gc(force=True)
|
||||
debug_log('WaifuDiffusion unload: complete')
|
||||
else:
|
||||
debug_log('WaifuDiffusion unload: no model loaded')
|
||||
|
||||
def preprocess_image(self, image: Image.Image) -> np.ndarray:
|
||||
"""Preprocess image for WaifuDiffusion model input.
|
||||
|
||||
- Resize to 448x448 (standard for WD models)
|
||||
- Pad to square with white background
|
||||
- Normalize to [0, 1] range
|
||||
- BGR channel order (as used by these models)
|
||||
"""
|
||||
original_size = image.size
|
||||
original_mode = image.mode
|
||||
|
||||
# Convert to RGB if needed
|
||||
if image.mode != 'RGB':
|
||||
image = image.convert('RGB')
|
||||
|
||||
# Pad to square with white background
|
||||
w, h = image.size
|
||||
max_dim = max(w, h)
|
||||
pad_left = (max_dim - w) // 2
|
||||
pad_top = (max_dim - h) // 2
|
||||
|
||||
padded = Image.new('RGB', (max_dim, max_dim), (255, 255, 255))
|
||||
padded.paste(image, (pad_left, pad_top))
|
||||
|
||||
# Resize to model input size
|
||||
if max_dim != self.image_size:
|
||||
padded = padded.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
|
||||
|
||||
# Convert to numpy array and normalize
|
||||
img_array = np.array(padded, dtype=np.float32)
|
||||
|
||||
# Convert RGB to BGR (model expects BGR)
|
||||
img_array = img_array[:, :, ::-1]
|
||||
|
||||
# Add batch dimension
|
||||
img_array = np.expand_dims(img_array, axis=0)
|
||||
|
||||
debug_log(f'WaifuDiffusion preprocess: original_size={original_size} mode={original_mode} padded_size={max_dim} output_shape={img_array.shape}')
|
||||
return img_array
|
||||
|
||||
def predict(
|
||||
self,
|
||||
image: Image.Image,
|
||||
general_threshold: float = None,
|
||||
character_threshold: float = None,
|
||||
include_rating: bool = None,
|
||||
exclude_tags: str = None,
|
||||
max_tags: int = None,
|
||||
sort_alpha: bool = None,
|
||||
use_spaces: bool = None,
|
||||
escape_brackets: bool = None,
|
||||
) -> str:
|
||||
"""Run inference and return formatted tag string.
|
||||
|
||||
Args:
|
||||
image: PIL Image to tag
|
||||
general_threshold: Threshold for general tags (0-1)
|
||||
character_threshold: Threshold for character tags (0-1)
|
||||
include_rating: Whether to include rating tags
|
||||
exclude_tags: Comma-separated tags to exclude
|
||||
max_tags: Maximum number of tags to return
|
||||
sort_alpha: Sort tags alphabetically vs by confidence
|
||||
use_spaces: Use spaces instead of underscores
|
||||
escape_brackets: Escape parentheses/brackets in tags
|
||||
|
||||
Returns:
|
||||
Formatted tag string
|
||||
"""
|
||||
t0 = time.time()
|
||||
|
||||
# Use settings defaults if not specified
|
||||
general_threshold = general_threshold or shared.opts.tagger_threshold
|
||||
character_threshold = character_threshold or shared.opts.waifudiffusion_character_threshold
|
||||
include_rating = include_rating if include_rating is not None else shared.opts.tagger_include_rating
|
||||
exclude_tags = exclude_tags or shared.opts.tagger_exclude_tags
|
||||
max_tags = max_tags or shared.opts.tagger_max_tags
|
||||
sort_alpha = sort_alpha if sort_alpha is not None else shared.opts.tagger_sort_alpha
|
||||
use_spaces = use_spaces if use_spaces is not None else shared.opts.tagger_use_spaces
|
||||
escape_brackets = escape_brackets if escape_brackets is not None else shared.opts.tagger_escape_brackets
|
||||
|
||||
debug_log(f'WaifuDiffusion predict: general_threshold={general_threshold} character_threshold={character_threshold} max_tags={max_tags} include_rating={include_rating} sort_alpha={sort_alpha}')
|
||||
|
||||
# Handle input variations
|
||||
if isinstance(image, list):
|
||||
image = image[0] if len(image) > 0 else None
|
||||
if isinstance(image, dict) and 'name' in image:
|
||||
image = Image.open(image['name'])
|
||||
if image is None:
|
||||
shared.log.error('WaifuDiffusion predict: no image provided')
|
||||
return ''
|
||||
|
||||
# Load model if needed
|
||||
if self.session is None:
|
||||
if not self.load():
|
||||
return ''
|
||||
|
||||
# Preprocess image
|
||||
img_input = self.preprocess_image(image)
|
||||
|
||||
# Run inference
|
||||
t_infer = time.time()
|
||||
input_name = self.session.get_inputs()[0].name
|
||||
output_name = self.session.get_outputs()[0].name
|
||||
probs = self.session.run([output_name], {input_name: img_input})[0][0]
|
||||
infer_time = time.time() - t_infer
|
||||
debug_log(f'WaifuDiffusion predict: inference time={infer_time:.3f}s output_shape={probs.shape}')
|
||||
|
||||
# Build tag list with probabilities
|
||||
tag_probs = {}
|
||||
exclude_set = {x.strip().replace(' ', '_').lower() for x in exclude_tags.split(',') if x.strip()}
|
||||
if exclude_set:
|
||||
debug_log(f'WaifuDiffusion predict: exclude_tags={exclude_set}')
|
||||
|
||||
general_count = 0
|
||||
character_count = 0
|
||||
rating_count = 0
|
||||
|
||||
for i, (tag_name, prob) in enumerate(zip(self.tags, probs)):
|
||||
category = self.tag_categories[i]
|
||||
tag_lower = tag_name.lower()
|
||||
|
||||
# Skip excluded tags
|
||||
if tag_lower in exclude_set:
|
||||
continue
|
||||
|
||||
# Apply category-specific thresholds
|
||||
if category == CATEGORY_RATING:
|
||||
if not include_rating:
|
||||
continue
|
||||
# Always include rating if enabled
|
||||
tag_probs[tag_name] = float(prob)
|
||||
rating_count += 1
|
||||
elif category == CATEGORY_CHARACTER:
|
||||
if prob >= character_threshold:
|
||||
tag_probs[tag_name] = float(prob)
|
||||
character_count += 1
|
||||
elif category == CATEGORY_GENERAL:
|
||||
if prob >= general_threshold:
|
||||
tag_probs[tag_name] = float(prob)
|
||||
general_count += 1
|
||||
else:
|
||||
# Other categories use general threshold
|
||||
if prob >= general_threshold:
|
||||
tag_probs[tag_name] = float(prob)
|
||||
|
||||
debug_log(f'WaifuDiffusion predict: matched tags general={general_count} character={character_count} rating={rating_count} total={len(tag_probs)}')
|
||||
|
||||
# Sort tags
|
||||
if sort_alpha:
|
||||
sorted_tags = sorted(tag_probs.keys())
|
||||
else:
|
||||
sorted_tags = [t for t, _ in sorted(tag_probs.items(), key=lambda x: -x[1])]
|
||||
|
||||
# Limit number of tags
|
||||
if max_tags > 0 and len(sorted_tags) > max_tags:
|
||||
sorted_tags = sorted_tags[:max_tags]
|
||||
debug_log(f'WaifuDiffusion predict: limited to max_tags={max_tags}')
|
||||
|
||||
# Format output
|
||||
result = []
|
||||
for tag_name in sorted_tags:
|
||||
formatted_tag = tag_name
|
||||
if use_spaces:
|
||||
formatted_tag = formatted_tag.replace('_', ' ')
|
||||
if escape_brackets:
|
||||
formatted_tag = re.sub(re_special, r'\\\1', formatted_tag)
|
||||
if shared.opts.tagger_show_scores:
|
||||
formatted_tag = f"({formatted_tag}:{tag_probs[tag_name]:.2f})"
|
||||
result.append(formatted_tag)
|
||||
|
||||
output = ", ".join(result)
|
||||
total_time = time.time() - t0
|
||||
debug_log(f'WaifuDiffusion predict: complete tags={len(result)} time={total_time:.2f} result="{output[:100]}..."' if len(output) > 100 else f'WaifuDiffusion predict: complete tags={len(result)} time={total_time:.2f} result="{output}"')
|
||||
|
||||
return output
|
||||
|
||||
def tag(self, image: Image.Image, **kwargs) -> str:
|
||||
"""Alias for predict() to match deepbooru interface."""
|
||||
return self.predict(image, **kwargs)
|
||||
|
||||
|
||||
# Global tagger instance
|
||||
tagger = WaifuDiffusionTagger()
|
||||
|
||||
|
||||
def _save_tags_to_file(img_path, tags_str: str, save_append: bool) -> bool:
|
||||
"""Save tags to a text file with error handling.
|
||||
|
||||
Args:
|
||||
img_path: Path to the image file
|
||||
tags_str: Tags string to save
|
||||
save_append: If True, append to existing file; otherwise overwrite
|
||||
|
||||
Returns:
|
||||
True if save succeeded, False otherwise
|
||||
"""
|
||||
try:
|
||||
txt_path = img_path.with_suffix('.txt')
|
||||
if save_append and txt_path.exists():
|
||||
with open(txt_path, 'a', encoding='utf-8') as f:
|
||||
f.write(f', {tags_str}')
|
||||
else:
|
||||
with open(txt_path, 'w', encoding='utf-8') as f:
|
||||
f.write(tags_str)
|
||||
return True
|
||||
except Exception as e:
|
||||
shared.log.error(f'WaifuDiffusion batch: failed to save file="{img_path}" error={e}')
|
||||
return False
|
||||
|
||||
|
||||
def get_models() -> list:
|
||||
"""Return list of available WaifuDiffusion model names."""
|
||||
return list(WAIFUDIFFUSION_MODELS.keys())
|
||||
|
||||
|
||||
def refresh_models() -> list:
|
||||
"""Refresh and return list of available models."""
|
||||
# For now, just return the static list
|
||||
# Could be extended to check for locally cached models
|
||||
return get_models()
|
||||
|
||||
|
||||
def load_model(model_name: str = None) -> bool:
|
||||
"""Load the specified WaifuDiffusion model."""
|
||||
return tagger.load(model_name)
|
||||
|
||||
|
||||
def unload_model():
|
||||
"""Unload the current WaifuDiffusion model."""
|
||||
tagger.unload()
|
||||
|
||||
|
||||
def tag(image: Image.Image, model_name: str = None, **kwargs) -> str:
|
||||
"""Tag an image using WaifuDiffusion tagger.
|
||||
|
||||
Args:
|
||||
image: PIL Image to tag
|
||||
model_name: Model to use (loads if needed)
|
||||
**kwargs: Additional arguments passed to predict()
|
||||
|
||||
Returns:
|
||||
Formatted tag string
|
||||
"""
|
||||
t0 = time.time()
|
||||
jobid = shared.state.begin('WaifuDiffusion Tag')
|
||||
shared.log.info(f'WaifuDiffusion: model="{model_name or tagger.model_name or shared.opts.waifudiffusion_model}" image_size={image.size if image else None}')
|
||||
|
||||
try:
|
||||
if model_name and model_name != tagger.model_name:
|
||||
tagger.load(model_name)
|
||||
result = tagger.predict(image, **kwargs)
|
||||
shared.log.debug(f'WaifuDiffusion: complete time={time.time()-t0:.2f} tags={len(result.split(", ")) if result else 0}')
|
||||
# Offload model if setting enabled
|
||||
if shared.opts.caption_offload:
|
||||
tagger.unload()
|
||||
except Exception as e:
|
||||
result = f"Exception {type(e)}"
|
||||
shared.log.error(f'WaifuDiffusion: {e}')
|
||||
errors.display(e, 'WaifuDiffusion Tag')
|
||||
|
||||
shared.state.end(jobid)
|
||||
return result
|
||||
|
||||
|
||||
def batch(
|
||||
model_name: str,
|
||||
batch_files: list,
|
||||
batch_folder: str,
|
||||
batch_str: str,
|
||||
save_output: bool = True,
|
||||
save_append: bool = False,
|
||||
recursive: bool = False,
|
||||
**kwargs
|
||||
) -> str:
|
||||
"""Process multiple images in batch mode.
|
||||
|
||||
Args:
|
||||
model_name: Model to use
|
||||
batch_files: List of file paths
|
||||
batch_folder: Folder path from file picker
|
||||
batch_str: Folder path as string
|
||||
save_output: Save caption to .txt files
|
||||
save_append: Append to existing caption files
|
||||
recursive: Recursively process subfolders
|
||||
**kwargs: Additional arguments passed to predict()
|
||||
|
||||
Returns:
|
||||
Combined tag results
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
# Load model
|
||||
if model_name:
|
||||
tagger.load(model_name)
|
||||
elif tagger.session is None:
|
||||
tagger.load()
|
||||
|
||||
# Collect image files
|
||||
image_files = []
|
||||
image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
|
||||
|
||||
# From file picker
|
||||
if batch_files:
|
||||
for f in batch_files:
|
||||
if isinstance(f, dict):
|
||||
image_files.append(Path(f['name']))
|
||||
elif hasattr(f, 'name'):
|
||||
image_files.append(Path(f.name))
|
||||
else:
|
||||
image_files.append(Path(f))
|
||||
|
||||
# From folder picker
|
||||
if batch_folder:
|
||||
folder_path = None
|
||||
if isinstance(batch_folder, list) and len(batch_folder) > 0:
|
||||
f = batch_folder[0]
|
||||
if isinstance(f, dict):
|
||||
folder_path = Path(f['name']).parent
|
||||
elif hasattr(f, 'name'):
|
||||
folder_path = Path(f.name).parent
|
||||
if folder_path and folder_path.is_dir():
|
||||
if recursive:
|
||||
for ext in image_extensions:
|
||||
image_files.extend(folder_path.rglob(f'*{ext}'))
|
||||
else:
|
||||
for ext in image_extensions:
|
||||
image_files.extend(folder_path.glob(f'*{ext}'))
|
||||
|
||||
# From string path
|
||||
if batch_str and batch_str.strip():
|
||||
folder_path = Path(batch_str.strip())
|
||||
if folder_path.is_dir():
|
||||
if recursive:
|
||||
for ext in image_extensions:
|
||||
image_files.extend(folder_path.rglob(f'*{ext}'))
|
||||
else:
|
||||
for ext in image_extensions:
|
||||
image_files.extend(folder_path.glob(f'*{ext}'))
|
||||
|
||||
# Remove duplicates while preserving order
|
||||
seen = set()
|
||||
unique_files = []
|
||||
for f in image_files:
|
||||
f_resolved = f.resolve()
|
||||
if f_resolved not in seen:
|
||||
seen.add(f_resolved)
|
||||
unique_files.append(f)
|
||||
image_files = unique_files
|
||||
|
||||
if not image_files:
|
||||
shared.log.warning('WaifuDiffusion batch: no images found')
|
||||
return ''
|
||||
|
||||
t0 = time.time()
|
||||
jobid = shared.state.begin('WaifuDiffusion Batch')
|
||||
shared.log.info(f'WaifuDiffusion batch: model="{tagger.model_name}" images={len(image_files)} write={save_output} append={save_append} recursive={recursive}')
|
||||
debug_log(f'WaifuDiffusion batch: files={[str(f) for f in image_files[:5]]}{"..." if len(image_files) > 5 else ""}')
|
||||
|
||||
results = []
|
||||
|
||||
# Progress bar
|
||||
import rich.progress as rp
|
||||
pbar = rp.Progress(rp.TextColumn('[cyan]WaifuDiffusion:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
|
||||
|
||||
with pbar:
|
||||
task = pbar.add_task(total=len(image_files), description='starting...')
|
||||
for img_path in image_files:
|
||||
pbar.update(task, advance=1, description=str(img_path.name))
|
||||
try:
|
||||
if shared.state.interrupted:
|
||||
shared.log.info('WaifuDiffusion batch: interrupted')
|
||||
break
|
||||
|
||||
image = Image.open(img_path)
|
||||
tags_str = tagger.predict(image, **kwargs)
|
||||
|
||||
if save_output:
|
||||
_save_tags_to_file(img_path, tags_str, save_append)
|
||||
|
||||
results.append(f'{img_path.name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{img_path.name}: {tags_str}')
|
||||
|
||||
except Exception as e:
|
||||
shared.log.error(f'WaifuDiffusion batch: file="{img_path}" error={e}')
|
||||
results.append(f'{img_path.name}: ERROR - {e}')
|
||||
|
||||
elapsed = time.time() - t0
|
||||
shared.log.info(f'WaifuDiffusion batch: complete images={len(results)} time={elapsed:.1f}s')
|
||||
shared.state.end(jobid)
|
||||
|
||||
return '\n'.join(results)
|
||||
Reference in New Issue
Block a user