🎨 完整的 IOPaint 项目更新
## 主要更新 - ✨ 更新所有依赖到最新稳定版本 - 📝 添加详细的项目文档和模型推荐 - 🔧 配置 VSCode Cloud Studio 预览功能 - 🐛 修复 PyTorch API 弃用警告 ## 依赖更新 - diffusers: 0.27.2 → 0.35.2 - gradio: 4.21.0 → 5.46.0 - peft: 0.7.1 → 0.18.0 - Pillow: 9.5.0 → 11.3.0 - fastapi: 0.108.0 → 0.116.2 ## 新增文件 - CLAUDE.md - 项目架构和开发指南 - UPGRADE_NOTES.md - 详细的升级说明 - .vscode/preview.yml - 预览配置 - .vscode/LAUNCH_GUIDE.md - 启动指南 - .gitignore - 更新的忽略规则 ## 代码修复 - 修复 iopaint/model/ldm.py 中的 torch.cuda.amp.autocast() 弃用警告 ## 文档更新 - README.md - 添加模型推荐和使用指南 - 完整的项目源码（iopaint/） - Web 前端源码（web_app/） 🤖 Generated with Claude Code
This commit is contained in:
let5sne
342
iopaint/plugins/segment_anything2/build_sam.py
Normal file
342
iopaint/plugins/segment_anything2/build_sam.py
Normal file
@@ -0,0 +1,342 @@
|
||||
# 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
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
common_kwargs_for_2_1 = 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,
|
||||
no_obj_embed_spatial=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=True,
|
||||
proj_tpos_enc_in_obj_ptrs=True,
|
||||
use_signed_tpos_enc_to_obj_ptrs=True,
|
||||
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_image_encoder_tiny():
|
||||
return 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",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_image_encoder_small():
|
||||
return 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",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_image_encoder_base():
|
||||
return 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",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_image_encoder_large():
|
||||
return 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",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_sam2_tiny():
|
||||
return SAM2Base(
|
||||
**common_kwargs,
|
||||
image_encoder=build_image_encoder_tiny(),
|
||||
memory_attention=build_memory_attention(),
|
||||
memory_encoder=build_memory_encoder(),
|
||||
)
|
||||
|
||||
|
||||
def build_sam2_small():
|
||||
return SAM2Base(
|
||||
**common_kwargs,
|
||||
image_encoder=build_image_encoder_small(),
|
||||
memory_attention=build_memory_attention(),
|
||||
memory_encoder=build_memory_encoder(),
|
||||
)
|
||||
|
||||
|
||||
def build_sam2_base():
|
||||
return SAM2Base(
|
||||
**common_kwargs,
|
||||
image_encoder=build_image_encoder_base(),
|
||||
memory_attention=build_memory_attention(),
|
||||
memory_encoder=build_memory_encoder(),
|
||||
)
|
||||
|
||||
|
||||
def build_sam2_large():
|
||||
return SAM2Base(
|
||||
**common_kwargs,
|
||||
image_encoder=build_image_encoder_large(),
|
||||
memory_attention=build_memory_attention(),
|
||||
memory_encoder=build_memory_encoder(),
|
||||
)
|
||||
|
||||
|
||||
def build_sam2_1_tiny():
|
||||
return SAM2Base(
|
||||
**common_kwargs_for_2_1,
|
||||
image_encoder=build_image_encoder_tiny(),
|
||||
memory_attention=build_memory_attention(),
|
||||
memory_encoder=build_memory_encoder(),
|
||||
)
|
||||
|
||||
|
||||
def build_sam2_1_small():
|
||||
return SAM2Base(
|
||||
**common_kwargs_for_2_1,
|
||||
image_encoder=build_image_encoder_small(),
|
||||
memory_attention=build_memory_attention(),
|
||||
memory_encoder=build_memory_encoder(),
|
||||
)
|
||||
|
||||
|
||||
def build_sam2_1_base():
|
||||
return SAM2Base(
|
||||
**common_kwargs_for_2_1,
|
||||
image_encoder=build_image_encoder_base(),
|
||||
memory_attention=build_memory_attention(),
|
||||
memory_encoder=build_memory_encoder(),
|
||||
)
|
||||
|
||||
|
||||
def build_sam2_1_large():
|
||||
return SAM2Base(
|
||||
**common_kwargs_for_2_1,
|
||||
image_encoder=build_image_encoder_large(),
|
||||
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,
|
||||
"sam2_1_tiny": build_sam2_1_tiny,
|
||||
"sam2_1_small": build_sam2_1_small,
|
||||
"sam2_1_base": build_sam2_1_base,
|
||||
"sam2_1_large": build_sam2_1_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")
|
||||
Reference in New Issue
Block a user