From 3d75ed722a48ac3b954fe61e3ecc63e41d2f9e8a Mon Sep 17 00:00:00 2001 From: Jiao77 Date: Thu, 20 Nov 2025 03:06:20 +0800 Subject: [PATCH] change some problem 6 --- tools/diffusion/ic_layout_diffusion_optimized.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tools/diffusion/ic_layout_diffusion_optimized.py b/tools/diffusion/ic_layout_diffusion_optimized.py index 4cc0997..f855962 100644 --- a/tools/diffusion/ic_layout_diffusion_optimized.py +++ b/tools/diffusion/ic_layout_diffusion_optimized.py @@ -308,6 +308,21 @@ class ManhattanAwareUNet(nn.Module): # 解码器路径 for i, (decoder, skip) in enumerate(zip(self.decoder, reversed(skips))): x = decoder(x) + + # 确保跳跃连接尺寸匹配 - 如果尺寸不匹配则进行裁剪或填充 + if x.shape[2:] != skip.shape[2:]: + # 裁剪到最小尺寸 + h_min = min(x.shape[2], skip.shape[2]) + w_min = min(x.shape[3], skip.shape[3]) + if x.shape[2] > h_min: + x = x[:, :, :h_min, :] + if x.shape[3] > w_min: + x = x[:, :, :, :w_min] + if skip.shape[2] > h_min: + skip = skip[:, :, :h_min, :] + if skip.shape[3] > w_min: + skip = skip[:, :, :, :w_min] + x = x + skip # 跳跃连接 # 输出