33统一管理编辑器的所有资源，包括图片、蒙版、区域等。
44"""
55
6+ import gc
67import logging
78import os
89from typing import Dict , List , Optional
1415from .types import MaskType
1516
1617
18+ def _release_gpu_memory ():
19+ """释放GPU显存"""
20+ try :
21+ import torch
22+ if torch .cuda .is_available ():
23+ torch .cuda .empty_cache ()
24+ torch .cuda .synchronize ()
25+ except ImportError :
26+ pass
27+ except Exception :
28+ pass
29+
30+
1731class ResourceManager :
1832 """资源管理器
1933
@@ -68,7 +82,7 @@ def load_image(self, image_path: str, json_data: Optional[Dict] = None) -> Image
6882
6983 # 加载图片
7084 try :
71- self .logger .info (f"Loading image: { image_path } )
85+ self .logger .debug (f"Loading image: { image_path } )
7286 image = Image .open (image_path )
7387
7488 # 创建资源对象
@@ -86,7 +100,7 @@ def load_image(self, image_path: str, json_data: Optional[Dict] = None) -> Image
86100 # 设置为当前图片
87101 self ._current_image = resource
88102
89- self .logger .info (f"Image loaded successfully: { image_path } { image .width } { image .height } )
103+ self .logger .debug (f"Image loaded successfully: { image_path } { image .width } { image .height } )
90104 return resource
91105
92106 except Exception as e :
@@ -107,18 +121,66 @@ def _add_to_cache(self, path: str, resource: ImageResource) -> None:
107121 old_resource = self ._image_cache .pop (oldest_path )
108122 old_resource .release ()
109123 self .logger .debug (f"Removed oldest image from cache: { oldest_path } )
124+
125+ # 释放内存
126+ gc .collect ()
110127
111128 self ._image_cache [path ] = resource
112129
113- def unload_image (self ) -> None :
114- """卸载当前图片及所有关联资源"""
130+ def release_image_from_cache (self , path : str ) -> bool :
131+ """从缓存中释放指定图片
132+
133+ Args:
134+ path: 图片路径
135+
136+ Returns:
137+ bool: 是否成功释放
138+ """
139+ path = os .path .normpath (path )
140+ if path in self ._image_cache :
141+ resource = self ._image_cache .pop (path )
142+ resource .release ()
143+ gc .collect ()
144+ self .logger .debug (f"Released image from cache: { path } )
145+ return True
146+ return False
147+
148+ def clear_image_cache (self ) -> None :
149+ """清空所有图片缓存"""
150+ for resource in self ._image_cache .values ():
151+ resource .release ()
152+ self ._image_cache .clear ()
153+ gc .collect ()
154+ _release_gpu_memory ()
155+ self .logger .info ("Cleared all image cache" )
156+
157+ def unload_image (self , release_from_cache : bool = False ) -> None :
158+ """卸载当前图片及所有关联资源
159+
160+ Args:
161+ release_from_cache: 是否同时从缓存中释放该图片
162+ """
115163 if self ._current_image :
116- self .logger .info (f"Unloading image: { self ._current_image .path } )
164+ current_path = self ._current_image .path
165+
166+ # 如果需要从缓存中释放
167+ if release_from_cache and current_path in self ._image_cache :
168+ resource = self ._image_cache .pop (current_path )
169+ resource .release ()
170+ self .logger .debug (f"Released image from cache: { current_path } )
171+
117172 self ._current_image = None
118173
119174 # 清空所有关联资源
120175 self .clear_masks ()
121176 self .clear_regions ()
177+ self .clear_cache ()
178+
179+ # 强制垃圾回收
180+ gc .collect ()
181+ _release_gpu_memory ()
182+
183+ self .logger .debug ("Image unloaded and memory released" )
122184
123185 def get_current_image (self ) -> Optional [ImageResource ]:
124186 """获取当前图片资源
@@ -249,12 +311,13 @@ def get_region(self, region_id: int) -> Optional[RegionResource]:
249311 return self ._regions .get (region_id )
250312
251313 def get_all_regions (self ) -> List [RegionResource ]:
252- """获取所有区域
314+ """获取所有区域（按region_id排序）
253315
254316 Returns:
255- List[RegionResource]: 区域列表
317+ List[RegionResource]: 区域列表，按region_id升序排列
256318 """
257- return list (self ._regions .values ())
319+ # 按region_id排序，确保顺序正确
320+ return [self ._regions [rid ] for rid in sorted (self ._regions .keys ())]
258321
259322 def clear_regions (self ) -> None :
260323 """清空所有区域"""
@@ -306,16 +369,48 @@ def cleanup_all(self) -> None:
306369 """清理所有资源"""
307370 self .logger .info ("Cleaning up all resources" )
308371
309- # 卸载当前图片
310- self .unload_image ()
372+ # 卸载当前图片（不从缓存释放，因为下面会清空缓存）
373+ if self ._current_image :
374+ self ._current_image = None
375+
376+ # 清空蒙版
377+ self .clear_masks ()
378+
379+ # 清空区域
380+ self .clear_regions ()
381+
382+ # 清空临时缓存
383+ self ._temp_cache .clear ()
311384
312- # 清理缓存
385+ # 清理图片缓存
313386 for resource in self ._image_cache .values ():
314387 resource .release ()
315388 self ._image_cache .clear ()
316389
390+ # 强制垃圾回收和GPU显存释放
391+ gc .collect ()
392+ _release_gpu_memory ()
393+
317394 self .logger .info ("All resources cleaned up" )
318395
396+ def release_memory_after_export (self ) -> None :
397+ """导出后释放内存
398+
399+ 清理临时缓存和GPU显存，但保留图片缓存以便快速切换
400+ """
401+ self .logger .info ("Releasing memory after export" )
402+
403+ # 清空临时缓存（inpainted图片等）
404+ self ._temp_cache .clear ()
405+
406+ # 强制垃圾回收
407+ gc .collect ()
408+
409+ # 释放GPU显存
410+ _release_gpu_memory ()
411+
412+ self .logger .info ("Memory released after export" )
413+
319414 def get_memory_usage_estimate (self ) -> int :
320415 """估算内存使用量（字节）
321416
0 commit comments