update

app.py

@@ -9,6 +9,8 @@ import spaces
 from PIL import Image
 from huggingface_hub import snapshot_download
 import gc
+import psutil
+import os
 try:
     import pynvml
     pynvml.nvmlInit()
@@ -38,14 +40,14 @@ SUGGESTED_PROMPTS = [
 ]
 CUSTOM_OPTION = "✍️ Enter custom prompt..."
 
-def check_memory_usage():
+def check_memory_usage(tag):
     process = psutil.Process(os.getpid())
     memory_mb = process.memory_info().rss / 1024 / 1024
-    print(f"🖥️ 当前内存使用: {memory_mb:.2f} MB")
+    print(f"{tag}\t🖥️ 当前内存使用: {memory_mb:.2f} MB")
     
     # 系统总内存
     total_memory = psutil.virtual_memory().total / 1024 / 1024 / 1024
-    print(f"💾 系统总内存: {total_memory:.2f} GB")
+    print(f"{tag}\t💾 系统总内存: {total_memory:.2f} GB")
 
 def load_models():
     global pipeline, style_generator, unitok, processor, code_freq, local_repo_dir
@@ -99,6 +101,7 @@ def load_models():
     unitok_state_dict = torch.load(os.path.join(codebook_path, "model.pth"), map_location="cpu")
     unitok.load_state_dict(unitok_state_dict)
     unitok.to(device, dtype=weight_type)
+    check_memory_usage('before pipeline')
     print('='*10, 'before pipeline')
     pipeline = CoTylePipeline.from_pretrained(
         local_repo_dir,
@@ -109,7 +112,7 @@ def load_models():
         requires_safety_checker=False,
         
     )
-
+    check_memory_usage('before qwen')
     qwen_text_visual_encoder = Qwen2_5_VLForConditionalGeneration_Quant.from_pretrained(
         local_repo_dir,
         subfolder="text_encoder",
@@ -117,7 +120,7 @@ def load_models():
     qwen_text_visual_encoder = Qwen2_5_VL_Quant(unitok, qwen_text_visual_encoder)
     qwen_text_visual_encoder.to(device, dtype=weight_type)
     pipeline.text_encoder = qwen_text_visual_encoder
-
+    check_memory_usage('after qwen')
     processor = Qwen2VLProcessor.from_pretrained(
         local_repo_dir,
         subfolder="processor",