Update app.py

app.py

@@ -263,6 +263,9 @@ class FluxLatentDINOFlow(nn.Module):
         self.type_emb_pixel = nn.Parameter(torch.randn(1, 1, hidden_size) * 0.02)
         self.type_emb_dino = nn.Parameter(torch.randn(1, 1, hidden_size) * 0.02)
         self.initialize_weights()
+        # 新增缓存
+        self._cached_dino_map = None
+        self._cached_lq_hash = None  # 可选：缓存输入哈希
 
     def initialize_weights(self):
         for name, m in self.named_modules():
@@ -301,15 +304,30 @@ class FluxLatentDINOFlow(nn.Module):
         pixel_tokens = pixel_tokens.flatten(2).transpose(1, 2)
         pixel_tokens = pixel_tokens + self.type_emb_pixel
 
-        with torch.no_grad():
-            mean = torch.tensor([0.485, 0.456, 0.406], device=lq_img.device).view(1, 3, 1, 1)
-            std = torch.tensor([0.229, 0.224, 0.225], device=lq_img.device).view(1, 3, 1, 1)
-            dino_in = (lq_img * 0.5 + 0.5 - mean) / std
-            dino_feats = self.dino.forward_features(dino_in)
-            if getattr(self.dino, "num_prefix_tokens", 0) > 0:
-                dino_feats = dino_feats[:, self.dino.num_prefix_tokens:]
-            d_h = d_w = int(dino_feats.shape[1] ** 0.5)
-            dino_map = dino_feats.transpose(1, 2).reshape(B, -1, d_h, d_w)
+        # 计算输入 hash
+        lq_hash = hash(lq_img.data_ptr())  # 简单用指针做哈希，也可用 tensor.sum().item() 更精确
+        if self._cached_dino_map is None or self._cached_lq_hash != lq_hash:
+            print("recalculating hash...")
+            # 只在缓存不存在或输入变化时计算 DINO
+            with torch.no_grad():
+                mean = torch.tensor([0.485, 0.456, 0.406], device=lq_img.device).view(1, 3, 1, 1)
+                std = torch.tensor([0.229, 0.224, 0.225], device=lq_img.device).view(1, 3, 1, 1)
+                dino_in = (lq_img * 0.5 + 0.5 - mean) / std
+                dino_feats = self.dino.forward_features(dino_in)
+                if getattr(self.dino, "num_prefix_tokens", 0) > 0:
+                    dino_feats = dino_feats[:, self.dino.num_prefix_tokens:]
+                d_h = d_w = int(dino_feats.shape[1] ** 0.5)
+                dino_map = dino_feats.transpose(1, 2).reshape(B, -1, d_h, d_w)
+                dino_map_resized = F.interpolate(dino_map, size=(self.grid_size, self.grid_size), mode='bilinear', align_corners=False)
+                dino_tokens = self.dino_adapter(dino_map_resized)
+                dino_tokens = dino_tokens.flatten(2).transpose(1, 2)
+                dino_tokens = dino_tokens + self.type_emb_dino
+
+            # 更新缓存
+            self._cached_dino_map = dino_tokens
+            self._cached_lq_hash = lq_hash
+        else:
+            dino_tokens = self._cached_dino_map
 
         dino_map_resized = F.interpolate(dino_map, size=(self.grid_size, self.grid_size), mode='bilinear', align_corners=False)
         dino_tokens = self.dino_adapter(dino_map_resized)