1b87a98261
## 主要更新 - ✨ 更新所有依赖到最新稳定版本 - 📝 添加详细的项目文档和模型推荐 - 🔧 配置 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
81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
from typing import Type, List, Union
|
|
|
|
import torch
|
|
from torch import nn as nn
|
|
from torch.nn import init as init
|
|
from torch.nn.modules.batchnorm import _BatchNorm
|
|
|
|
|
|
@torch.no_grad()
|
|
def default_init_weights(
|
|
module_list: Union[List[nn.Module], nn.Module],
|
|
scale: float = 1,
|
|
bias_fill: float = 0,
|
|
**kwargs,
|
|
) -> None:
|
|
"""Initialize network weights.
|
|
|
|
Args:
|
|
module_list (list[nn.Module] | nn.Module): Modules to be initialized.
|
|
scale (float): Scale initialized weights, especially for residual
|
|
blocks. Default: 1.
|
|
bias_fill (float): The value to fill bias. Default: 0
|
|
kwargs (dict): Other arguments for initialization function.
|
|
"""
|
|
if not isinstance(module_list, list):
|
|
module_list = [module_list]
|
|
for module in module_list:
|
|
for m in module.modules():
|
|
if isinstance(m, nn.Conv2d):
|
|
init.kaiming_normal_(m.weight, **kwargs)
|
|
m.weight.data *= scale
|
|
if m.bias is not None:
|
|
m.bias.data.fill_(bias_fill)
|
|
elif isinstance(m, nn.Linear):
|
|
init.kaiming_normal_(m.weight, **kwargs)
|
|
m.weight.data *= scale
|
|
if m.bias is not None:
|
|
m.bias.data.fill_(bias_fill)
|
|
elif isinstance(m, _BatchNorm):
|
|
init.constant_(m.weight, 1)
|
|
if m.bias is not None:
|
|
m.bias.data.fill_(bias_fill)
|
|
|
|
|
|
def make_layer(
|
|
basic_block: Type[nn.Module], num_basic_block: int, **kwarg
|
|
) -> nn.Sequential:
|
|
"""Make layers by stacking the same blocks.
|
|
|
|
Args:
|
|
basic_block (Type[nn.Module]): nn.Module class for basic block.
|
|
num_basic_block (int): number of blocks.
|
|
|
|
Returns:
|
|
nn.Sequential: Stacked blocks in nn.Sequential.
|
|
"""
|
|
layers = []
|
|
for _ in range(num_basic_block):
|
|
layers.append(basic_block(**kwarg))
|
|
return nn.Sequential(*layers)
|
|
|
|
|
|
# TODO: may write a cpp file
|
|
def pixel_unshuffle(x: torch.Tensor, scale: int) -> torch.Tensor:
|
|
"""Pixel unshuffle.
|
|
|
|
Args:
|
|
x (Tensor): Input feature with shape (b, c, hh, hw).
|
|
scale (int): Downsample ratio.
|
|
|
|
Returns:
|
|
Tensor: the pixel unshuffled feature.
|
|
"""
|
|
b, c, hh, hw = x.size()
|
|
out_channel = c * (scale**2)
|
|
assert hh % scale == 0 and hw % scale == 0
|
|
h = hh // scale
|
|
w = hw // scale
|
|
x_view = x.view(b, c, h, scale, w, scale)
|
|
return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)
|