IOPaint/inpaint/plugins/segment_anything2/build_sam.py
root 70af4845af new file: inpaint/__init__.py
new file:   inpaint/__main__.py
	new file:   inpaint/api.py
	new file:   inpaint/batch_processing.py
	new file:   inpaint/benchmark.py
	new file:   inpaint/cli.py
	new file:   inpaint/const.py
	new file:   inpaint/download.py
	new file:   inpaint/file_manager/__init__.py
	new file:   inpaint/file_manager/file_manager.py
	new file:   inpaint/file_manager/storage_backends.py
	new file:   inpaint/file_manager/utils.py
	new file:   inpaint/helper.py
	new file:   inpaint/installer.py
	new file:   inpaint/model/__init__.py
	new file:   inpaint/model/anytext/__init__.py
	new file:   inpaint/model/anytext/anytext_model.py
	new file:   inpaint/model/anytext/anytext_pipeline.py
	new file:   inpaint/model/anytext/anytext_sd15.yaml
	new file:   inpaint/model/anytext/cldm/__init__.py
	new file:   inpaint/model/anytext/cldm/cldm.py
	new file:   inpaint/model/anytext/cldm/ddim_hacked.py
	new file:   inpaint/model/anytext/cldm/embedding_manager.py
	new file:   inpaint/model/anytext/cldm/hack.py
	new file:   inpaint/model/anytext/cldm/model.py
	new file:   inpaint/model/anytext/cldm/recognizer.py
	new file:   inpaint/model/anytext/ldm/__init__.py
	new file:   inpaint/model/anytext/ldm/models/__init__.py
	new file:   inpaint/model/anytext/ldm/models/autoencoder.py
	new file:   inpaint/model/anytext/ldm/models/diffusion/__init__.py
	new file:   inpaint/model/anytext/ldm/models/diffusion/ddim.py
	new file:   inpaint/model/anytext/ldm/models/diffusion/ddpm.py
	new file:   inpaint/model/anytext/ldm/models/diffusion/dpm_solver/__init__.py
	new file:   inpaint/model/anytext/ldm/models/diffusion/dpm_solver/dpm_solver.py
	new file:   inpaint/model/anytext/ldm/models/diffusion/dpm_solver/sampler.py
	new file:   inpaint/model/anytext/ldm/models/diffusion/plms.py
	new file:   inpaint/model/anytext/ldm/models/diffusion/sampling_util.py
	new file:   inpaint/model/anytext/ldm/modules/__init__.py
	new file:   inpaint/model/anytext/ldm/modules/attention.py
	new file:   inpaint/model/anytext/ldm/modules/diffusionmodules/__init__.py
	new file:   inpaint/model/anytext/ldm/modules/diffusionmodules/model.py
	new file:   inpaint/model/anytext/ldm/modules/diffusionmodules/openaimodel.py
	new file:   inpaint/model/anytext/ldm/modules/diffusionmodules/upscaling.py
	new file:   inpaint/model/anytext/ldm/modules/diffusionmodules/util.py
	new file:   inpaint/model/anytext/ldm/modules/distributions/__init__.py
	new file:   inpaint/model/anytext/ldm/modules/distributions/distributions.py
	new file:   inpaint/model/anytext/ldm/modules/ema.py
	new file:   inpaint/model/anytext/ldm/modules/encoders/__init__.py
	new file:   inpaint/model/anytext/ldm/modules/encoders/modules.py
	new file:   inpaint/model/anytext/ldm/util.py
	new file:   inpaint/model/anytext/main.py
	new file:   inpaint/model/anytext/ocr_recog/RNN.py
	new file:   inpaint/model/anytext/ocr_recog/RecCTCHead.py
	new file:   inpaint/model/anytext/ocr_recog/RecModel.py
	new file:   inpaint/model/anytext/ocr_recog/RecMv1_enhance.py
	new file:   inpaint/model/anytext/ocr_recog/RecSVTR.py
	new file:   inpaint/model/anytext/ocr_recog/__init__.py
	new file:   inpaint/model/anytext/ocr_recog/common.py
	new file:   inpaint/model/anytext/ocr_recog/en_dict.txt
	new file:   inpaint/model/anytext/ocr_recog/ppocr_keys_v1.txt
	new file:   inpaint/model/anytext/utils.py
	new file:   inpaint/model/base.py
	new file:   inpaint/model/brushnet/__init__.py
	new file:   inpaint/model/brushnet/brushnet.py
	new file:   inpaint/model/brushnet/brushnet_unet_forward.py
	new file:   inpaint/model/brushnet/brushnet_wrapper.py
	new file:   inpaint/model/brushnet/pipeline_brushnet.py
	new file:   inpaint/model/brushnet/unet_2d_blocks.py
	new file:   inpaint/model/controlnet.py
	new file:   inpaint/model/ddim_sampler.py
	new file:   inpaint/model/fcf.py
	new file:   inpaint/model/helper/__init__.py
	new file:   inpaint/model/helper/controlnet_preprocess.py
	new file:   inpaint/model/helper/cpu_text_encoder.py
	new file:   inpaint/model/helper/g_diffuser_bot.py
	new file:   inpaint/model/instruct_pix2pix.py
	new file:   inpaint/model/kandinsky.py
	new file:   inpaint/model/lama.py
	new file:   inpaint/model/ldm.py
	new file:   inpaint/model/manga.py
	new file:   inpaint/model/mat.py
	new file:   inpaint/model/mi_gan.py
	new file:   inpaint/model/opencv2.py
	new file:   inpaint/model/original_sd_configs/__init__.py
	new file:   inpaint/model/original_sd_configs/sd_xl_base.yaml
	new file:   inpaint/model/original_sd_configs/sd_xl_refiner.yaml
	new file:   inpaint/model/original_sd_configs/v1-inference.yaml
	new file:   inpaint/model/original_sd_configs/v2-inference-v.yaml
	new file:   inpaint/model/paint_by_example.py
	new file:   inpaint/model/plms_sampler.py
	new file:   inpaint/model/power_paint/__init__.py
	new file:   inpaint/model/power_paint/pipeline_powerpaint.py
	new file:   inpaint/model/power_paint/power_paint.py
	new file:   inpaint/model/power_paint/power_paint_v2.py
	new file:   inpaint/model/power_paint/powerpaint_tokenizer.py
2024-08-20 21:17:33 +02:00

263 lines
7.7 KiB
Python

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import logging
import torch
from pathlib import Path
from .modeling.backbones.hieradet import Hiera
from .modeling.backbones.image_encoder import ImageEncoder, FpnNeck
from .modeling.memory_attention import MemoryAttention, MemoryAttentionLayer
from .modeling.memory_encoder import MemoryEncoder, MaskDownSampler, Fuser, CXBlock
from .modeling.position_encoding import PositionEmbeddingSine
from .modeling.sam.transformer import RoPEAttention
from .modeling.sam2_base import SAM2Base
CURRENT_DIR = Path(__file__).parent
CONFIG_DIR = CURRENT_DIR / "sam2_configs"
common_kwargs = dict(
num_maskmem=7,
image_size=1024,
sigmoid_scale_for_mem_enc=20.0,
sigmoid_bias_for_mem_enc=-10.0,
use_mask_input_as_output_without_sam=True,
directly_add_no_mem_embed=True,
use_high_res_features_in_sam=True,
multimask_output_in_sam=True,
iou_prediction_use_sigmoid=True,
use_obj_ptrs_in_encoder=True,
add_tpos_enc_to_obj_ptrs=False,
only_obj_ptrs_in_the_past_for_eval=True,
pred_obj_scores=True,
pred_obj_scores_mlp=True,
fixed_no_obj_ptr=True,
multimask_output_for_tracking=True,
use_multimask_token_for_obj_ptr=True,
multimask_min_pt_num=0,
multimask_max_pt_num=1,
use_mlp_for_obj_ptr_proj=True,
compile_image_encoder=False,
)
def build_memory_attention():
return MemoryAttention(
d_model=256,
pos_enc_at_input=True,
layer=MemoryAttentionLayer(
activation="relu",
dim_feedforward=2048,
dropout=0.1,
pos_enc_at_attn=False,
self_attention=RoPEAttention(
rope_theta=10000.0,
feat_sizes=[32, 32],
embedding_dim=256,
num_heads=1,
downsample_rate=1,
dropout=0.1,
),
d_model=256,
pos_enc_at_cross_attn_keys=True,
pos_enc_at_cross_attn_queries=False,
cross_attention=RoPEAttention(
rope_theta=10000.0,
feat_sizes=[32, 32],
embedding_dim=256,
num_heads=1,
downsample_rate=1,
dropout=0.1,
kv_in_dim=64,
),
),
num_layers=4,
)
def build_memory_encoder():
return MemoryEncoder(
out_dim=64,
position_encoding=PositionEmbeddingSine(
num_pos_feats=64, normalize=True, scale=None, temperature=10000
),
mask_downsampler=MaskDownSampler(
kernel_size=3,
stride=2,
padding=1,
),
fuser=Fuser(
layer=CXBlock(
dim=256,
kernel_size=7,
padding=3,
layer_scale_init_value=1e-6,
use_dwconv=True,
),
num_layers=2,
),
)
def build_sam2_tiny():
return SAM2Base(
**common_kwargs,
image_encoder=ImageEncoder(
scalp=1,
trunk=Hiera(
embed_dim=96,
num_heads=1,
stages=(1, 2, 7, 2),
global_att_blocks=(5, 7, 9),
window_pos_embed_bkg_spatial_size=(7, 7),
window_spec=(8, 4, 14, 7),
),
neck=FpnNeck(
position_encoding=PositionEmbeddingSine(
num_pos_feats=256,
normalize=True,
scale=None,
temperature=10000,
),
d_model=256,
backbone_channel_list=[768, 384, 192, 96],
fpn_top_down_levels=[2, 3],
fpn_interp_model="nearest",
),
),
memory_attention=build_memory_attention(),
memory_encoder=build_memory_encoder(),
)
def build_sam2_small():
return SAM2Base(
**common_kwargs,
image_encoder=ImageEncoder(
scalp=1,
trunk=Hiera(
embed_dim=96,
num_heads=1,
stages=(1, 2, 11, 2),
global_att_blocks=(7, 10, 13),
window_pos_embed_bkg_spatial_size=(7, 7),
window_spec=(8, 4, 14, 7),
),
neck=FpnNeck(
position_encoding=PositionEmbeddingSine(
num_pos_feats=256,
normalize=True,
scale=None,
temperature=10000,
),
d_model=256,
backbone_channel_list=[768, 384, 192, 96],
fpn_top_down_levels=[2, 3],
fpn_interp_model="nearest",
),
),
memory_attention=build_memory_attention(),
memory_encoder=build_memory_encoder(),
)
def build_sam2_base():
return SAM2Base(
**common_kwargs,
image_encoder=ImageEncoder(
scalp=1,
trunk=Hiera(
embed_dim=112,
num_heads=2,
stages=(2, 3, 16, 3),
global_att_blocks=(12, 16, 20),
window_pos_embed_bkg_spatial_size=(14, 14),
window_spec=(8, 4, 14, 7),
),
neck=FpnNeck(
position_encoding=PositionEmbeddingSine(
num_pos_feats=256,
normalize=True,
scale=None,
temperature=10000,
),
d_model=256,
backbone_channel_list=[896, 448, 224, 112],
fpn_top_down_levels=[2, 3],
fpn_interp_model="nearest",
),
),
memory_attention=build_memory_attention(),
memory_encoder=build_memory_encoder(),
)
def build_sam2_large():
return SAM2Base(
**common_kwargs,
image_encoder=ImageEncoder(
scalp=1,
trunk=Hiera(
embed_dim=144,
num_heads=2,
stages=(2, 6, 36, 4),
global_att_blocks=(23, 33, 43),
window_pos_embed_bkg_spatial_size=(7, 7),
window_spec=(8, 4, 16, 8),
),
neck=FpnNeck(
position_encoding=PositionEmbeddingSine(
num_pos_feats=256,
normalize=True,
scale=None,
temperature=10000,
),
d_model=256,
backbone_channel_list=[1152, 576, 288, 144],
fpn_top_down_levels=[2, 3],
fpn_interp_model="nearest",
),
),
memory_attention=build_memory_attention(),
memory_encoder=build_memory_encoder(),
)
sam2_model_registry = {
"sam2_tiny": build_sam2_tiny,
"sam2_small": build_sam2_small,
"sam2_base": build_sam2_base,
"sam2_large": build_sam2_large,
}
def build_sam2(
name,
ckpt_path=None,
device="cuda",
mode="eval",
):
model = sam2_model_registry[name]()
_load_checkpoint(model, ckpt_path)
model = model.to(device)
if mode == "eval":
model.eval()
return model
def _load_checkpoint(model, ckpt_path):
if ckpt_path is not None:
sd = torch.load(ckpt_path, map_location="cpu")["model"]
missing_keys, unexpected_keys = model.load_state_dict(sd)
if missing_keys:
logging.error(missing_keys)
raise RuntimeError()
if unexpected_keys:
logging.error(unexpected_keys)
raise RuntimeError()
logging.info("Loaded checkpoint sucessfully")