@lobehub/chat 1.98.2 → 1.99.0

package/src/store/image/slices/generationConfig/action.ts

@@ -0,0 +1,295 @@
+import { StateCreator } from 'zustand/vanilla';
+
+import {
+  ModelParamsSchema,
+  RuntimeImageGenParams,
+  RuntimeImageGenParamsKeys,
+  RuntimeImageGenParamsValue,
+  extractDefaultValues,
+} from '@/libs/standard-parameters/meta-schema';
+import { aiProviderSelectors, getAiInfraStoreState } from '@/store/aiInfra';
+import { AIImageModelCard } from '@/types/aiModel';
+
+import type { ImageStore } from '../../store';
+import { adaptSizeToRatio, parseRatio } from '../../utils/size';
+
+export interface GenerationConfigAction {
+  setParamOnInput<K extends RuntimeImageGenParamsKeys>(
+    paramName: K,
+    value: RuntimeImageGenParamsValue,
+  ): void;
+
+  setModelAndProviderOnSelect(model: string, provider: string): void;
+
+  setImageNum: (imageNum: number) => void;
+
+  reuseSettings: (
+    model: string,
+    provider: string,
+    settings: Partial<RuntimeImageGenParams>,
+  ) => void;
+  reuseSeed: (seed: number) => void;
+
+  setWidth(width: number): void;
+  setHeight(height: number): void;
+  toggleAspectRatioLock(): void;
+  setAspectRatio(aspectRatio: string): void;
+}
+
+/**
+ * @internal
+ * This function is exported only for testing purposes.
+ * Do not use this function directly in application code.
+ */
+export function getModelAndDefaults(model: string, provider: string) {
+  const enabledImageModelList = aiProviderSelectors.enabledImageModelList(getAiInfraStoreState());
+  const activeModel = enabledImageModelList
+    .find((providerItem) => providerItem.id === provider)
+    ?.children.find((modelItem) => modelItem.id === model) as unknown as AIImageModelCard;
+
+  const parametersSchema = activeModel.parameters as ModelParamsSchema;
+  const defaultValues = extractDefaultValues(parametersSchema);
+
+  return { defaultValues, activeModel, parametersSchema };
+}
+
+export const createGenerationConfigSlice: StateCreator<
+  ImageStore,
+  [['zustand/devtools', never]],
+  [],
+  GenerationConfigAction
+> = (set, get) => ({
+  setParamOnInput: (paramName, value) => {
+    set(
+      (state) => {
+        const { parameters } = state;
+        return { parameters: { ...parameters, [paramName]: value } };
+      },
+      false,
+      `setParamOnInput/${paramName}`,
+    );
+  },
+
+  setWidth: (width) => {
+    set(
+      (state) => {
+        const {
+          parameters,
+          isAspectRatioLocked,
+          activeAspectRatio,
+          parametersSchema: parametersSchema,
+        } = state;
+
+        const newParams = { ...parameters, width };
+        if (isAspectRatioLocked && activeAspectRatio) {
+          const ratio = parseRatio(activeAspectRatio);
+          const heightSchema = parametersSchema?.height;
+          if (
+            heightSchema &&
+            typeof heightSchema.max === 'number' &&
+            typeof heightSchema.min === 'number'
+          ) {
+            const newHeight = Math.round(width / ratio);
+            newParams.height = Math.max(Math.min(newHeight, heightSchema.max), heightSchema.min);
+          }
+        }
+
+        return { parameters: newParams };
+      },
+      false,
+      `setWidth`,
+    );
+  },
+
+  setHeight: (height) => {
+    set(
+      (state) => {
+        const {
+          parameters,
+          isAspectRatioLocked,
+          activeAspectRatio,
+          parametersSchema: parametersSchema,
+        } = state;
+        const newParams = { ...parameters, height };
+
+        if (isAspectRatioLocked && activeAspectRatio) {
+          const ratio = parseRatio(activeAspectRatio);
+          const widthSchema = parametersSchema?.width;
+          if (
+            widthSchema &&
+            typeof widthSchema.max === 'number' &&
+            typeof widthSchema.min === 'number'
+          ) {
+            const newWidth = Math.round(height * ratio);
+            newParams.width = Math.max(Math.min(newWidth, widthSchema.max), widthSchema.min);
+          }
+        }
+
+        return { parameters: newParams };
+      },
+      false,
+      `setHeight`,
+    );
+  },
+
+  toggleAspectRatioLock: () => {
+    set(
+      (state) => {
+        const {
+          isAspectRatioLocked,
+          activeAspectRatio,
+          parameters,
+          parametersSchema: parametersSchema,
+        } = state;
+        const newLockState = !isAspectRatioLocked;
+
+        // 如果是从解锁变为锁定，且有活动的宽高比，则立即调整尺寸
+        if (newLockState && activeAspectRatio && parameters && parametersSchema) {
+          const currentWidth = parameters.width;
+          const currentHeight = parameters.height;
+
+          // 只有当width和height都存在时才进行调整
+          if (
+            typeof currentWidth === 'number' &&
+            typeof currentHeight === 'number' &&
+            parametersSchema?.width &&
+            parametersSchema?.height
+          ) {
+            const targetRatio = parseRatio(activeAspectRatio);
+            const currentRatio = currentWidth / currentHeight;
+
+            // 如果当前比例与目标比例不匹配，则需要调整
+            if (Math.abs(currentRatio - targetRatio) > 0.01) {
+              // 允许小误差
+              const widthSchema = parametersSchema.width;
+              const heightSchema = parametersSchema.height;
+
+              if (
+                widthSchema &&
+                heightSchema &&
+                typeof widthSchema.max === 'number' &&
+                typeof widthSchema.min === 'number' &&
+                typeof heightSchema.max === 'number' &&
+                typeof heightSchema.min === 'number'
+              ) {
+                // 优先保持宽度，调整高度
+                let newWidth = currentWidth;
+                let newHeight = Math.round(currentWidth / targetRatio);
+
+                // 如果计算出的高度超出范围，则改为保持高度，调整宽度
+                if (newHeight > heightSchema.max || newHeight < heightSchema.min) {
+                  newHeight = currentHeight;
+                  newWidth = Math.round(currentHeight * targetRatio);
+
+                  // 确保宽度也在范围内
+                  newWidth = Math.max(Math.min(newWidth, widthSchema.max), widthSchema.min);
+                } else {
+                  // 确保高度在范围内
+                  newHeight = Math.max(Math.min(newHeight, heightSchema.max), heightSchema.min);
+                }
+
+                return {
+                  isAspectRatioLocked: newLockState,
+                  parameters: { ...parameters, width: newWidth, height: newHeight },
+                };
+              }
+            }
+          }
+        }
+
+        return { isAspectRatioLocked: newLockState };
+      },
+      false,
+      'toggleAspectRatioLock',
+    );
+  },
+
+  setAspectRatio: (aspectRatio) => {
+    const { parameters, parametersSchema: parametersSchema } = get();
+    if (!parameters || !parametersSchema) return;
+
+    const defaultValues = extractDefaultValues(parametersSchema);
+    const newParams = { ...parameters };
+
+    // 如果模型支持 width/height，则计算新尺寸
+    if (
+      parametersSchema?.width &&
+      parametersSchema?.height &&
+      typeof defaultValues.width === 'number' &&
+      typeof defaultValues.height === 'number'
+    ) {
+      const ratio = parseRatio(aspectRatio);
+      const { width, height } = adaptSizeToRatio(ratio, defaultValues.width, defaultValues.height);
+      newParams.width = width;
+      newParams.height = height;
+    }
+
+    // 如果模型本身支持 aspectRatio，则更新它
+    if (parametersSchema?.aspectRatio) {
+      newParams.aspectRatio = aspectRatio;
+    }
+
+    set(
+      { activeAspectRatio: aspectRatio, parameters: newParams },
+      false,
+      `setAspectRatio/${aspectRatio}`,
+    );
+  },
+
+  setModelAndProviderOnSelect: (model, provider) => {
+    const { defaultValues, activeModel } = getModelAndDefaults(model, provider);
+    const parametersSchema = activeModel.parameters;
+
+    let initialActiveRatio: string | null = null;
+
+    // 如果模型没有原生比例或尺寸参数，但有宽高，则启用虚拟比例控制
+    if (
+      !parametersSchema?.aspectRatio &&
+      !parametersSchema?.size &&
+      parametersSchema?.width &&
+      parametersSchema?.height
+    ) {
+      const { width, height } = defaultValues;
+      if (typeof width === 'number' && typeof height === 'number' && width > 0 && height > 0) {
+        initialActiveRatio = `${width}:${height}`;
+      } else {
+        initialActiveRatio = '1:1';
+      }
+    }
+
+    set(
+      {
+        model,
+        provider,
+        parameters: defaultValues,
+        parametersSchema: parametersSchema,
+        isAspectRatioLocked: false,
+        activeAspectRatio: initialActiveRatio,
+      },
+      false,
+      `setModelAndProviderOnSelect/${model}/${provider}`,
+    );
+  },
+
+  setImageNum: (imageNum) => {
+    set(() => ({ imageNum }), false, `setImageNum/${imageNum}`);
+  },
+
+  reuseSettings: (model: string, provider: string, settings: Partial<RuntimeImageGenParams>) => {
+    const { defaultValues, parametersSchema } = getModelAndDefaults(model, provider);
+    set(
+      () => ({
+        model,
+        provider,
+        parameters: { ...defaultValues, ...settings },
+        parametersSchema: parametersSchema,
+      }),
+      false,
+      `reuseSettings/${model}/${provider}`,
+    );
+  },
+
+  reuseSeed: (seed: number) => {
+    set((state) => ({ parameters: { ...state.parameters, seed } }), false, `reuseSeed/${seed}`);
+  },
+});

package/src/store/image/slices/generationConfig/hooks.test.ts

@@ -0,0 +1,304 @@
+import { act, renderHook } from '@testing-library/react';
+import { describe, expect, it, vi } from 'vitest';
+
+import { fluxSchnellParamsSchema } from '@/config/paramsSchemas/fal/flux-schnell';
+import { ModelParamsSchema, RuntimeImageGenParams } from '@/libs/standard-parameters/meta-schema';
+import { useImageStore } from '@/store/image';
+import { AIImageModelCard } from '@/types/aiModel';
+
+import { useGenerationConfigParam } from './hooks';
+
+// Mock external dependencies
+vi.mock('@/store/aiInfra', () => ({
+  aiProviderSelectors: {
+    enabledImageModelList: vi.fn(() => [
+      {
+        id: 'fal',
+        name: 'Fal',
+        children: [
+          {
+            id: 'flux/schnell',
+            displayName: 'FLUX.1 Schnell',
+            type: 'image',
+            parameters: fluxSchnellParamsSchema,
+            releasedAt: '2024-08-01',
+          } as AIImageModelCard,
+        ],
+      },
+    ]),
+  },
+  getAiInfraStoreState: vi.fn(() => ({})),
+}));
+
+const testModelSchema: ModelParamsSchema = {
+  prompt: { default: '', description: 'The text prompt for image generation' },
+  width: { default: 1024, min: 512, max: 2048, step: 64, description: 'Image width' },
+  height: { default: 768, min: 256, max: 1536, step: 32, description: 'Image height' },
+  steps: { default: 20, min: 1, max: 50, description: 'Number of inference steps' },
+  seed: { default: null, min: 0, max: 2147483647, description: 'Random seed' },
+  cfg: { default: 7.5, min: 1, max: 20, step: 0.5, description: 'CFG scale' },
+  aspectRatio: {
+    default: '16:9',
+    enum: ['1:1', '16:9', '4:3', '9:16'],
+    description: 'Aspect ratio',
+  },
+};
+
+const testParameters: RuntimeImageGenParams = {
+  prompt: 'test prompt',
+  width: 1024,
+  height: 768,
+  steps: 25,
+  seed: 12345,
+  cfg: 8.0,
+  aspectRatio: '16:9',
+};
+
+describe('useGenerationConfigParam', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    // Reset store state before each test
+    useImageStore.setState({
+      parametersSchema: testModelSchema,
+      parameters: testParameters,
+    });
+  });
+
+  describe('value and setValue', () => {
+    it('should return current parameter value', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('width'));
+      expect(result.current.value).toBe(1024);
+    });
+
+    it('should return current string parameter value', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('prompt'));
+      expect(result.current.value).toBe('test prompt');
+    });
+
+    it('should return null for null values', () => {
+      useImageStore.setState({
+        parameters: { ...testParameters, seed: null },
+      });
+
+      const { result } = renderHook(() => useGenerationConfigParam('seed'));
+      expect(result.current.value).toBeNull();
+    });
+
+    it('should update parameter value using setValue', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('width'));
+
+      act(() => {
+        result.current.setValue(2048);
+      });
+
+      const updatedParameters = useImageStore.getState().parameters;
+      expect(updatedParameters?.width).toBe(2048);
+    });
+
+    it('should update string parameter value using setValue', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('prompt'));
+
+      act(() => {
+        result.current.setValue('new test prompt');
+      });
+
+      const updatedParameters = useImageStore.getState().parameters;
+      expect(updatedParameters?.prompt).toBe('new test prompt');
+    });
+
+    it('should handle array parameter updates', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('imageUrls'));
+
+      act(() => {
+        result.current.setValue(['image1.jpg', 'image2.jpg']);
+      });
+
+      const updatedParameters = useImageStore.getState().parameters;
+      expect(updatedParameters?.imageUrls).toEqual(['image1.jpg', 'image2.jpg']);
+    });
+  });
+
+  describe('parameter constraints', () => {
+    it('should return min and max constraints for numeric parameter', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('width'));
+
+      expect(result.current.min).toBe(512);
+      expect(result.current.max).toBe(2048);
+      expect(result.current.step).toBe(64);
+      expect(result.current.description).toBe('Image width');
+      expect(result.current.enumValues).toBeUndefined();
+    });
+
+    it('should return step constraint for decimal parameter', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('cfg'));
+
+      expect(result.current.min).toBe(1);
+      expect(result.current.max).toBe(20);
+      expect(result.current.step).toBe(0.5);
+      expect(result.current.description).toBe('CFG scale');
+    });
+
+    it('should return enum values for enum parameter', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('aspectRatio'));
+
+      expect(result.current.enumValues).toEqual(['1:1', '16:9', '4:3', '9:16']);
+      expect(result.current.description).toBe('Aspect ratio');
+      expect(result.current.min).toBeUndefined();
+      expect(result.current.max).toBeUndefined();
+      expect(result.current.step).toBeUndefined();
+    });
+
+    it('should return undefined constraints for parameter without constraints', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('prompt'));
+
+      expect(result.current.min).toBeUndefined();
+      expect(result.current.max).toBeUndefined();
+      expect(result.current.step).toBeUndefined();
+      expect(result.current.enumValues).toBeUndefined();
+      expect(result.current.description).toBe('The text prompt for image generation');
+    });
+
+    it('should handle seed parameter with large range', () => {
+      const { result } = renderHook(() => useGenerationConfigParam('seed'));
+
+      expect(result.current.min).toBe(0);
+      expect(result.current.max).toBe(2147483647);
+      expect(result.current.step).toBeUndefined();
+      expect(result.current.description).toBe('Random seed');
+    });
+  });
+
+  describe('edge cases', () => {
+    it('should handle undefined parameters', () => {
+      useImageStore.setState({ parameters: undefined });
+
+      const { result } = renderHook(() => useGenerationConfigParam('width'));
+      expect(result.current.value).toBeUndefined();
+    });
+
+    it('should handle undefined parametersSchema', () => {
+      useImageStore.setState({ parametersSchema: undefined });
+
+      const { result } = renderHook(() => useGenerationConfigParam('width'));
+
+      expect(result.current.min).toBeUndefined();
+      expect(result.current.max).toBeUndefined();
+      expect(result.current.step).toBeUndefined();
+      expect(result.current.description).toBeUndefined();
+      expect(result.current.enumValues).toBeUndefined();
+    });
+
+    it('should handle parameter not in current parametersSchema', () => {
+      useImageStore.setState({
+        parametersSchema: { prompt: { default: '' } }, // Only prompt defined
+      });
+
+      const { result } = renderHook(() => useGenerationConfigParam('width'));
+
+      expect(result.current.min).toBeUndefined();
+      expect(result.current.max).toBeUndefined();
+      expect(result.current.step).toBeUndefined();
+      expect(result.current.enumValues).toBeUndefined();
+    });
+  });
+
+  describe('flux/schnell real-world parameters', () => {
+    it('should work with actual flux/schnell parameters', () => {
+      useImageStore.setState({
+        parametersSchema: fluxSchnellParamsSchema,
+        parameters: {
+          prompt: 'A beautiful landscape',
+          width: 1024,
+          height: 1024,
+          steps: 4,
+          seed: null,
+        },
+      });
+
+      // Test prompt parameter
+      const { result: promptResult } = renderHook(() => useGenerationConfigParam('prompt'));
+      expect(promptResult.current.value).toBe('A beautiful landscape');
+      // flux-schnell's prompt parameter doesn't have description
+
+      // Test width parameter
+      const { result: widthResult } = renderHook(() => useGenerationConfigParam('width'));
+      expect(widthResult.current.value).toBe(1024);
+      expect(widthResult.current.min).toBeDefined();
+      expect(widthResult.current.max).toBeDefined();
+
+      // Test steps parameter
+      const { result: stepsResult } = renderHook(() => useGenerationConfigParam('steps'));
+      expect(stepsResult.current.value).toBe(4);
+      expect(stepsResult.current.min).toBeDefined();
+      expect(stepsResult.current.max).toBeDefined();
+
+      // Test seed parameter
+      const { result: seedResult } = renderHook(() => useGenerationConfigParam('seed'));
+      expect(seedResult.current.value).toBeNull();
+    });
+
+    it('should update flux/schnell parameters correctly', () => {
+      useImageStore.setState({
+        parametersSchema: fluxSchnellParamsSchema,
+        parameters: {
+          prompt: 'original prompt',
+          width: 512,
+          height: 512,
+          steps: 1,
+          seed: null,
+        },
+      });
+
+      const { result: promptResult } = renderHook(() => useGenerationConfigParam('prompt'));
+      const { result: widthResult } = renderHook(() => useGenerationConfigParam('width'));
+
+      // Update prompt
+      act(() => {
+        promptResult.current.setValue('updated prompt');
+      });
+
+      // Update width
+      act(() => {
+        widthResult.current.setValue(1024);
+      });
+
+      const updatedParameters = useImageStore.getState().parameters;
+      expect(updatedParameters?.prompt).toBe('updated prompt');
+      expect(updatedParameters?.width).toBe(1024);
+    });
+  });
+
+  describe('setValue callback stability', () => {
+    it('should maintain setValue callback reference when value changes', () => {
+      const { result, rerender } = renderHook(() => useGenerationConfigParam('width'));
+
+      const initialSetValue = result.current.setValue;
+
+      // Change the parameter value
+      act(() => {
+        result.current.setValue(2048);
+      });
+
+      rerender();
+
+      // setValue callback should remain the same reference
+      expect(result.current.setValue).toBe(initialSetValue);
+    });
+
+    it('should maintain setValue callback reference when other parameters change', () => {
+      const { result: widthResult } = renderHook(() => useGenerationConfigParam('width'));
+      const { result: promptResult } = renderHook(() => useGenerationConfigParam('prompt'));
+
+      const initialWidthSetValue = widthResult.current.setValue;
+
+      // Change a different parameter
+      act(() => {
+        promptResult.current.setValue('new prompt');
+      });
+
+      // setValue callback for width should remain the same reference
+      expect(widthResult.current.setValue).toBe(initialWidthSetValue);
+    });
+  });
+});