Jiao77/RoRD-Layout-Recognation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

change some problem 4

Browse Source
This commit is contained in:
Jiao77
2025-11-20 03:04:56 +08:00
parent ebda75fa5e
commit f8975b26b4
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
tools/diffusion/ic_layout_diffusion_optimized.py
View File

@@ -207,6 +207,14 @@ class ManhattanAwareUNet(nn.Module):
self._make_block(512, 1024, stride=2), self._make_block(512, 1024, stride=2),
]) ])
# 残差投影层 - 确保残差连接的通道数匹配
self.residual_projections = nn.ModuleList([
nn.Conv2d(64, 128, 1), # 64 -> 128
nn.Conv2d(128, 256, 1), # 128 -> 256
nn.Conv2d(256, 512, 1), # 256 -> 512
nn.Conv2d(512, 1024, 1), # 512 -> 1024
])
# 中间层 # 中间层
self.middle = nn.Sequential( self.middle = nn.Sequential(
nn.Conv2d(1024, 1024, 3, padding=1), nn.Conv2d(1024, 1024, 3, padding=1),
@@ -283,16 +291,16 @@ class ManhattanAwareUNet(nn.Module):
# 编码器路径 # 编码器路径
skips = [] skips = []
for i, (encoder, fusion) in enumerate(zip(self.encoder, self.time_fusion)): for i, (encoder, fusion) in enumerate(zip(self.encoder, self.time_fusion)):
# 残差连接 # 保存残差连接并使用投影层匹配通道数
residual = x residual = self.residual_projections[i](x)
x = encoder(x) x = encoder(x)
# 融合时间信息 # 融合时间信息
t_feat = fusion(t_emb).unsqueeze(-1).unsqueeze(-1) t_feat = fusion(t_emb).unsqueeze(-1).unsqueeze(-1)
x = x + t_feat x = x + t_feat
# 跳跃连接 # 跳跃连接 - 添加投影后的残差
skips.append(x + residual if i == 0 else x) skips.append(x + residual)
# 中间层 # 中间层
x = self.middle(x) x = self.middle(x)