Update app.py

app.py

@@ -127,9 +127,14 @@ def initialize_model():
                 PRIVATE_MODEL,
                 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                 use_safetensors=True,
-                trust_remote_code=True  # 私人仓库需要
+                trust_remote_code=True,
+                # 更激进的内存优化
+                text_encoder_dtype=torch.float32,
+                device_map="balanced",
+                load_in_8bit=True,  # 8bit量化
+                low_cpu_mem_usage=True  # 低CPU内存使用
             )
-            print("Successfully loaded private NSFW Wan model!")
+            print("Successfully loaded private NSFW Wan model with memory optimization!")
             
         except Exception as private_error:
             print(f"Private Wan model loading failed: {private_error}")
@@ -140,7 +145,9 @@ def initialize_model():
                 pipeline = WanPipeline.from_pretrained(
                     "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
                     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                    use_safetensors=True
+                    use_safetensors=True,
+                    text_encoder_dtype=torch.float32,
+                    device_map="balanced"
                 )
                 print("Loaded official Wan2.2-Diffusers model")
             except Exception as wan_error:
@@ -156,22 +163,24 @@ def initialize_model():
         
         pipeline = pipeline.to(device)
         
-        # GPU优化 - CogVideoX专用
+        # GPU优化 - Wan模型专用内存管理
         if torch.cuda.is_available():
             try:
-                # CogVideoX特有的优化方法
-                if hasattr(pipeline, 'enable_vae_tiling'):
-                    pipeline.enable_vae_tiling()
+                # Wan模型特有的优化方法
                 if hasattr(pipeline, 'enable_model_cpu_offload'):
-                    pipeline.enable_model_cpu_offload()
+                    pipeline.enable_model_cpu_offload()  # 将部分组件移至CPU
+                if hasattr(pipeline, 'enable_vae_tiling'):
+                    pipeline.enable_vae_tiling()  # VAE分块处理
+                if hasattr(pipeline, 'enable_sequential_cpu_offload'):
+                    pipeline.enable_sequential_cpu_offload()  # 顺序CPU卸载
                 # 通用内存优化
                 try:
                     pipeline.enable_xformers_memory_efficient_attention()
                 except:
                     pass
-                print("✅ CogVideoX memory optimizations applied")
+                print("Wan model memory optimizations applied")
             except Exception as mem_error:
-                print(f"⚠️ Memory optimization warning: {mem_error}")
+                print(f"Memory optimization warning: {mem_error}")
         
         # 初始化Compel
         if COMPEL_AVAILABLE and hasattr(pipeline, 'tokenizer'):