Hugging Face's logo

Spaces:
Jack698
/
velo
Sleeping

App Files Community
velo
File size: 4,830 Bytes
adc1e1c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
/*
Copyright (c) 2025 Tethys Plex

This file is part of Veloera.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/

import { useState, useCallback } from 'react';
import { ModelMappingUtils } from '../utils/modelMappingUtils';

/**
 * 模型映射管理 Hook
 * 提供模型映射相关的状态管理和操作方法
 */
export const useModelMapping = (initialMapping = {}) => {
  // 用于追踪模型的原始名称映射关系 { displayName: originalName }
  const [modelOriginalMapping, setModelOriginalMapping] = useState(initialMapping);

  /**
   * 更新模型列表和映射关系
   * @param {Array} newModels - 新的模型列表
   * @param {Object} newMapping - 新的映射关系
   * @param {Function} updateModelsCallback - 更新模型列表的回调函数
   */
  const updateModelsList = useCallback((newModels, newMapping, updateModelsCallback) => {
    const uniqueModels = ModelMappingUtils.filterAndDedupeModels(newModels);
    
    if (updateModelsCallback) {
      updateModelsCallback(uniqueModels);
    }
    
    setModelOriginalMapping(newMapping);
  }, []);

  /**
   * 同步模型映射到模型配置
   * @param {string} mappingValue - JSON 格式的映射配置
   * @param {Array} currentModels - 当前模型列表
   * @param {Function} updateModelsCallback - 更新模型列表的回调函数
   */
  const syncModelMappingToModels = useCallback((mappingValue, currentModels, updateModelsCallback) => {
    const mapping = ModelMappingUtils.parseModelMapping(mappingValue);

    if (!mapping || Object.keys(mapping).length === 0) {
      // 当映射为空时，恢复为原始名称
      if (!currentModels || currentModels.length === 0) {
        return;
      }

      // 没有映射关系，则不需要恢复
      if (!modelOriginalMapping || Object.keys(modelOriginalMapping).length === 0) {
        updateModelsList(currentModels, {}, updateModelsCallback);
        return;
      }

      // 恢复模型到原始名称
      const restoredModels = ModelMappingUtils.restoreModelsToOriginal(
        currentModels, 
        modelOriginalMapping
      );

      // 清空原始映射关系，因为已经恢复完成
      updateModelsList(restoredModels, {}, updateModelsCallback);
      return;
    }

    // 在应用新映射之前，先恢复当前模型到原始名称
    const restoredModels = ModelMappingUtils.restoreModelsToOriginal(
      currentModels, 
      modelOriginalMapping
    );

    const { updatedModels, newMapping, hasChanges } = ModelMappingUtils.applyModelMapping(
      mapping,
      restoredModels,
      {}
    );

    if (hasChanges || !ModelMappingUtils.areModelsEqual(currentModels, updatedModels)) {
      updateModelsList(updatedModels, newMapping, updateModelsCallback);
    }
  }, [modelOriginalMapping, updateModelsList]);

  /**
   * 初始化模型映射关系
   * @param {string} mappingValue - JSON 格式的映射配置
   * @param {Array} models - 模型列表
   */
  const initializeModelMapping = useCallback((mappingValue, models) => {
    const mapping = ModelMappingUtils.parseModelMapping(mappingValue);
    
    if (mapping && Object.keys(mapping).length > 0) {
    const initialMapping = Object.create(null);
    const DANGEROUS_KEYS = new Set(['__proto__', 'prototype', 'constructor']);  
      
      // 根据当前的模型映射和模型列表，建立原始映射关系
      Object.entries(mapping).forEach(([displayName, originalName]) => {
        const displayNameTrimmed = displayName.trim();
        const originalNameTrimmed = typeof originalName === 'string' ? originalName.trim() : null;
        
        if (displayNameTrimmed && originalNameTrimmed && !DANGEROUS_KEYS.has(displayNameTrimmed) && models.includes(displayNameTrimmed)) {
          initialMapping[displayNameTrimmed] = originalNameTrimmed;
        }
      });
      
      setModelOriginalMapping(initialMapping);
    } else {
      setModelOriginalMapping({});
    }
  }, []);

  /**
   * 重置映射关系
   */
  const resetModelMapping = useCallback(() => {
    setModelOriginalMapping({});
  }, []);

  return {
    modelOriginalMapping,
    updateModelsList,
    syncModelMappingToModels,
    initializeModelMapping,
    resetModelMapping,
    setModelOriginalMapping
  };
};