Update app.py

app.py

@@ -44,8 +44,9 @@ STYLE_PRESETS = {
 }
 
 # 固定模型配置 - 使用NSFW版本的safetensors文件
-FIXED_MODEL = "https://huggingface.co/NSFW-API/NSFW_Wan_14b/resolve/main/wan_14B_e15.safetensors"
-BASE_CONFIG_MODEL = "Wan-AI/Wan2.1-T2V-14B"  # 用于获取配置信息
+NSFW_CHECKPOINT_URL = "https://huggingface.co/NSFW-API/NSFW_Wan_14b/resolve/main/wan_14B_e15.safetensors"
+# 使用CogVideoX作为基础配置，因为Wan很可能基于CogVideoX架构
+BASE_CONFIG_MODEL = "THUDM/CogVideoX-5b"  # 用于获取配置信息
 
 # 质量增强提示词 - 适配视频
 QUALITY_ENHANCERS = [
@@ -116,46 +117,48 @@ def initialize_model():
         print(f"🖥️ Using device: {device}")
         
         print(f"📦 Loading NSFW Wan T2V checkpoint: wan_14B_e15.safetensors")
-        print(f"🔧 Using config from: {BASE_CONFIG_MODEL}")
+        print(f"🔧 Using CogVideoX config from: {BASE_CONFIG_MODEL}")
         
-        # 使用from_single_file方法加载NSFW检查点
+        # 方法1: 尝试使用CogVideoXPipeline加载（最可能成功）
         try:
-            from diffusers import WanPipeline
-            pipeline = WanPipeline.from_single_file(
-                FIXED_MODEL,
-                config=BASE_CONFIG_MODEL,  # 使用官方配置
+            from diffusers import CogVideoXPipeline
+            pipeline = CogVideoXPipeline.from_single_file(
+                NSFW_CHECKPOINT_URL,
+                config=BASE_CONFIG_MODEL,
                 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                 use_safetensors=True,
                 safety_checker=None,
                 requires_safety_checker=False
             )
-            print("✅ Successfully loaded with WanPipeline.from_single_file()")
-        except ImportError:
-            # 如果WanPipeline不可用，使用通用DiffusionPipeline
-            print("⚠️ WanPipeline not found, trying DiffusionPipeline")
-            pipeline = DiffusionPipeline.from_single_file(
-                FIXED_MODEL,
-                config=BASE_CONFIG_MODEL,  # 使用官方配置
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                use_safetensors=True,
-                safety_checker=None,
-                requires_safety_checker=False,
-                trust_remote_code=True
-            )
-            print("✅ Successfully loaded with DiffusionPipeline.from_single_file()")
-        except Exception as single_file_error:
-            # 最后备选：尝试直接从官方模型仓库加载（但这不是NSFW版本）
-            print(f"⚠️ Single file loading failed: {single_file_error}")
-            print("🔄 Falling back to official model repository...")
-            pipeline = DiffusionPipeline.from_pretrained(
-                BASE_CONFIG_MODEL,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                use_safetensors=True,
-                safety_checker=None,
-                requires_safety_checker=False,
-                trust_remote_code=True
-            )
-            print("✅ Loaded official model as fallback")
+            print("✅ Successfully loaded with CogVideoXPipeline.from_single_file()")
+        except Exception as cogvideo_error:
+            print(f"⚠️ CogVideoX loading failed: {cogvideo_error}")
+            
+            # 方法2: 尝试通用DiffusionPipeline
+            try:
+                pipeline = DiffusionPipeline.from_single_file(
+                    NSFW_CHECKPOINT_URL,
+                    config=BASE_CONFIG_MODEL,
+                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                    use_safetensors=True,
+                    safety_checker=None,
+                    requires_safety_checker=False,
+                    trust_remote_code=True
+                )
+                print("✅ Successfully loaded with DiffusionPipeline.from_single_file()")
+            except Exception as diffusion_error:
+                print(f"⚠️ DiffusionPipeline single file loading failed: {diffusion_error}")
+                
+                # 方法3: 最后备选 - 直接使用官方CogVideoX（非NSFW）
+                print("🔄 Falling back to official CogVideoX model (non-NSFW)...")
+                pipeline = CogVideoXPipeline.from_pretrained(
+                    BASE_CONFIG_MODEL,
+                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                    use_safetensors=True,
+                    safety_checker=None,
+                    requires_safety_checker=False
+                )
+                print("✅ Loaded official CogVideoX as fallback (non-NSFW version)")
         
         pipeline = pipeline.to(device)
         
@@ -635,7 +638,7 @@ def create_interface():
             # 简化警告信息
             gr.HTML('''
                 <div class="warning-box">
-                    ⚠️ 18+ CONTENT WARNING ⚠️ | T2V Model: NSFW_Wan_14b (wan_14B_e15.safetensors)
+                    ⚠️ 18+ CONTENT WARNING ⚠️ | Model: NSFW_Wan_14b → CogVideoX Architecture
                 </div>
             ''')
             
@@ -863,9 +866,9 @@ def create_interface():
 # ===== 启动应用 =====
 if __name__ == "__main__":
     print("🎬 Starting NSFW Video Generator...")
-    print(f"🔧 NSFW Wan T2V Checkpoint: wan_14B_e15.safetensors")
-    print(f"🔧 Base Config Model: {BASE_CONFIG_MODEL}")
-    print("🔧 Using from_single_file() method for NSFW checkpoint")
+    print(f"🔧 NSFW Checkpoint: {NSFW_CHECKPOINT_URL}")
+    print(f"🔧 Base Config: {BASE_CONFIG_MODEL} (CogVideoX-5b)")
+    print("🔧 Using CogVideoXPipeline.from_single_file() method")
     print(f"🔧 Default Duration: {VIDEO_CONFIG['default_duration']}s")
     print(f"🔧 Default Resolution: {VIDEO_CONFIG['default_width']}×{VIDEO_CONFIG['default_height']}")
     print(f"🔧 Spaces GPU: {'✅ Available' if SPACES_AVAILABLE else '❌ Not Available'}")