新增：实现模型的自定义NSFW状态切换功能，更新相关API和前端组件，优化模型列表和详情展示

c1921 · c1921 · commit d9bb979549d7 · 2025-04-22T17:59:52.000+08:00
diff --git a/frontend/src/App.vue b/frontend/src/App.vue
@@ -44,6 +44,7 @@
             :blur-nsfw="blurNsfw"
             @model-click="openModelDetails"
             @open-settings="openSettings"
+            @model-updated="handleModelUpdated"
           />
         </div>
       </div>
@@ -309,6 +310,15 @@ function closeModelDetail() {
   selectedModel.value = null;
 }
 
+// 处理模型更新
+function handleModelUpdated(updatedModel: Model) {
+  // 查找并更新模型列表中的对应模型
+  const index = models.value.findIndex(model => model.id === updatedModel.id);
+  if (index !== -1) {
+    models.value[index] = { ...models.value[index], ...updatedModel };
+  }
+}
+
 // 监听模型列表变化
 watch(models, () => {
   updateFilterOptions();
diff --git a/frontend/src/api/models.ts b/frontend/src/api/models.ts
@@ -18,6 +18,8 @@ export interface Model {
   size?: number;
   preview?: string;
   nsfw?: boolean;
+  custom_nsfw?: boolean;
+  original_nsfw?: boolean;
   created_at?: string;
   hash?: string;
   tags?: string[];
@@ -32,6 +34,8 @@ interface BackendModel {
   type: string;
   preview_url?: string;
   nsfw?: boolean;
+  custom_nsfw?: boolean;
+  original_nsfw?: boolean;
   baseModel?: string;
   url?: string;
   nsfwLevel?: number;
@@ -49,6 +53,8 @@ function convertModel(backendModel: BackendModel): Model {
     type: backendModel.type,
     preview: backendModel.preview_url,
     nsfw: backendModel.nsfw || false,
+    custom_nsfw: backendModel.custom_nsfw || false,
+    original_nsfw: backendModel.original_nsfw || false,
     base_model: backendModel.baseModel,
     url: backendModel.url
   };
@@ -184,5 +190,11 @@ export const ModelsAPI = {
   getVersion: async (): Promise<{version: string; company: string; copyright: string}> => {
     const response = await apiClient.get('/version');
     return response.data;
+  },
+
+  // 切换模型的自定义NSFW状态
+  toggleModelNsfw: async (modelId: string): Promise<{nsfw: boolean}> => {
+    const response = await apiClient.post('/toggle-nsfw', { model_id: modelId });
+    return { nsfw: response.data.nsfw };
   }
 }; 
diff --git a/frontend/src/components/ModelDetailModal.vue b/frontend/src/components/ModelDetailModal.vue
@@ -19,7 +19,14 @@
     <div class="modal-dialog overlay-open:opacity-100 overlay-open:duration-300 max-w-4xl">
       <div class="modal-content bg-base-100">
         <div class="modal-header border-b border-base-200">
-          <h3 class="modal-title text-base-content">{{ model?.name || '' }}</h3>
+          <h3 class="modal-title text-base-content">
+            {{ model?.name || '' }}
+            <span 
+              v-if="model?.nsfw" 
+              class="badge ms-2 align-middle" 
+              :class="model.custom_nsfw ? 'badge-warning' : 'badge-error'"
+            >NSFW</span>
+          </h3>
           <button 
             type="button" 
             class="btn btn-text btn-circle btn-sm absolute end-3 top-3 text-base-content/70 hover:text-base-content" 
diff --git a/frontend/src/components/ModelList.vue b/frontend/src/components/ModelList.vue
@@ -44,8 +44,20 @@
             </div>
             <div 
               v-if="model.nsfw" 
-              class="badge badge-error absolute top-2 right-2 z-10"
+              class="badge absolute top-2 right-2 z-10"
+              :class="model.custom_nsfw ? 'badge-warning' : 'badge-error'"
             >NSFW</div>
+            
+            <!-- NSFW标记按钮 - 只在非原始NSFW模型上显示 -->
+            <button 
+              v-if="!model.original_nsfw"
+              class="absolute top-2 left-2 z-10 btn btn-circle btn-xs" 
+              :class="[model.custom_nsfw ? 'btn-warning' : 'btn-outline btn-neutral']"
+              title="标记/取消标记为NSFW内容"
+              @click.stop="toggleNsfw(model)"
+            >
+              <span class="icon-[tabler--eye-off] size-3.5"></span>
+            </button>
           </div>
           <div 
             class="p-4 flex-1 cursor-pointer"
@@ -102,6 +114,7 @@
 
 <script setup lang="ts">
 import type { Model } from '../api/models';
+import { ModelsAPI } from '../api/models';
 import toast from '../utils/toast';
 
 defineProps<{
@@ -118,6 +131,7 @@ defineProps<{
 const emit = defineEmits<{
   'open-settings': [];
   'model-click': [model: Model];
+  'model-updated': [model: Model];
 }>();
 
 function onOpenSettings() {
@@ -128,6 +142,32 @@ function onModelClick(model: Model) {
   emit('model-click', model);
 }
 
+async function toggleNsfw(model: Model) {
+  // 如果是原始NSFW模型，不允许更改
+  if (model.original_nsfw) {
+    toast.error('无法修改原始NSFW模型的状态');
+    return;
+  }
+  
+  try {
+    // 调用API切换NSFW状态
+    const result = await ModelsAPI.toggleModelNsfw(model.id);
+    
+    // 更新模型状态
+    model.custom_nsfw = result.nsfw;
+    model.nsfw = result.nsfw || model.original_nsfw; // 保持与原始NSFW状态一致
+    
+    // 通知父组件模型已更新
+    emit('model-updated', model);
+    
+    // 显示成功提示
+    toast.success(`已${result.nsfw ? '标记' : '取消标记'}为NSFW内容`);
+  } catch (error) {
+    console.error('切换NSFW状态失败:', error);
+    toast.error('切换NSFW状态失败');
+  }
+}
+
 function copyFileName(filename: string) {
   navigator.clipboard.writeText(filename)
     .then(() => {
diff --git a/src/api/model_api.py b/src/api/model_api.py
@@ -11,6 +11,9 @@
 class PathUpdate(BaseModel):
     path: str
 
+class ModelIdParam(BaseModel):
+    model_id: str
+
 def create_api(manager, frontend_url=None):
     """创建并配置FastAPI应用
     
@@ -140,6 +143,15 @@ async def get_config():
             "is_path_valid": os.path.exists(manager.models_path) if manager.models_path else False
         }
 
+    @app.post("/api/toggle-nsfw")
+    async def toggle_model_nsfw(model_param: ModelIdParam):
+        """切换模型的NSFW状态"""
+        try:
+            new_state = manager.toggle_custom_nsfw(model_param.model_id)
+            return {"success": True, "model_id": model_param.model_id, "nsfw": new_state}
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"设置NSFW状态失败: {str(e)}")
+
     @app.get("/api/select-path")
     async def select_path_endpoint():
         """选择目录"""
diff --git a/src/core/model_manager.py b/src/core/model_manager.py
@@ -25,13 +25,15 @@ def __init__(self, config_file="config.json"):
         # 添加并发限制和超时设置
         self.semaphore = asyncio.Semaphore(5)  # 限制并发请求数
         self.timeout = ClientTimeout(total=10)  # 10秒超时
+        # 添加自定义NSFW模型ID列表
+        self.custom_nsfw_models = self.config.get("custom_nsfw_models", [])
         
     def load_config(self) -> dict:
         """加载配置文件"""
         if os.path.exists(self.config_file):
             with open(self.config_file, "r", encoding="utf-8") as f:
                 return json.load(f)
-        return {"models_path": "", "output_file": "models_info.json"}
+        return {"models_path": "", "output_file": "models_info.json", "custom_nsfw_models": []}
     
     def save_config(self):
         """保存配置文件"""
@@ -212,7 +214,9 @@ def get_model_display_info(self, model_path: str) -> dict:
                 "preview_url": None,
                 "description": "未找到模型信息",
                 "baseModel": "未知",
-                "nsfw": False,
+                "nsfw": str(model_path) in self.custom_nsfw_models,  # 检查是否在自定义NSFW列表中
+                "custom_nsfw": str(model_path) in self.custom_nsfw_models,  # 新增自定义NSFW标记
+                "original_nsfw": False,  # 新增原始NSFW标记
                 "nsfwLevel": 0,
             }
         
@@ -224,13 +228,19 @@ def get_model_display_info(self, model_path: str) -> dict:
         # 添加本地图片路径
         local_preview = info.get("local_preview")
         
+        # 如果在自定义NSFW列表中，覆盖API返回的nsfw值
+        is_custom_nsfw = str(model_path) in self.custom_nsfw_models
+        is_original_nsfw = model_data.get("nsfw", False)
+        
         return {
             "name": model_data.get("name", Path(model_path).name),
             "type": model_data.get("type", "未知"),
             "preview_url": local_preview or preview_url,
             "baseModel": info.get("baseModel", "未知"),
-            "url": f"https://civitai.com/models/{info['modelId']}?modelVersionId={info['id']}",
-            "nsfw": model_data.get("nsfw", False),
+            "url": f"https://civitai.com/models/{info['modelId']}?modelVersionId={info['id']}" if 'modelId' in info and 'id' in info else None,
+            "nsfw": is_custom_nsfw or is_original_nsfw,  # 自定义NSFW或API返回的NSFW
+            "custom_nsfw": is_custom_nsfw,  # 新增自定义NSFW标记
+            "original_nsfw": is_original_nsfw,  # 新增原始NSFW标记
             "nsfwLevel": preview_image.get("nsfwLevel", 0),
         }
 
@@ -244,4 +254,41 @@ def get_all_models_info(self) -> list:
             }
             for model_path in self.models_info.keys()
             if str(model_path).startswith(current_path)
-        ] 
+        ] 
+
+    def toggle_custom_nsfw(self, model_path: str) -> bool:
+        """切换模型的自定义NSFW状态
+        
+        Args:
+            model_path: 模型路径
+            
+        Returns:
+            bool: 更新后的NSFW状态
+        """
+        model_path = str(model_path)  # 确保为字符串
+        
+        # 检查是否为原始NSFW模型
+        model_info = self.models_info.get(model_path, {})
+        if model_info:
+            info = model_info.get("info", {})
+            model_data = info.get("model", {})
+            is_original_nsfw = model_data.get("nsfw", False)
+            
+            # 如果是原始NSFW模型，不允许更改
+            if is_original_nsfw:
+                return True  # 保持NSFW状态
+        
+        if model_path in self.custom_nsfw_models:
+            # 如果模型已在自定义NSFW列表中，则移除
+            self.custom_nsfw_models.remove(model_path)
+            current_state = False
+        else:
+            # 否则添加到列表中
+            self.custom_nsfw_models.append(model_path)
+            current_state = True
+            
+        # 更新配置并保存
+        self.config["custom_nsfw_models"] = self.custom_nsfw_models
+        self.save_config()
+        
+        return current_state 
