refactor(ui): make layer adjustments schemas/types composable

Author: psychedelicious

invokeai/frontend/web/src/features/controlLayers/components/RasterLayer/RasterLayerCurvesAdjustmentsEditor.tsx

@@ -5,20 +5,19 @@ import { useEntityAdapterContext } from 'features/controlLayers/contexts/EntityA
 import { useEntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
 import { rasterLayerAdjustmentsCurvesUpdated } from 'features/controlLayers/store/canvasSlice';
 import { selectCanvasSlice, selectEntity } from 'features/controlLayers/store/selectors';
+import type { ChannelName, ChannelPoints } from 'features/controlLayers/store/types';
 import React, { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiArrowCounterClockwiseBold } from 'react-icons/pi';
 
-const DEFAULT_POINTS: Array<[number, number]> = [
+const DEFAULT_POINTS: ChannelPoints = [
   [0, 0],
   [255, 255],
 ];
 
-type Channel = 'master' | 'r' | 'g' | 'b';
+type ChannelHistograms = Record<ChannelName, number[] | null>;
 
-type ChannelHistograms = Record<Channel, number[] | null>;
-
-const channelColor: Record<Channel, string> = {
+const channelColor: Record<ChannelName, string> = {
   master: '#888',
   r: '#e53e3e',
   g: '#38a169',
@@ -27,7 +26,7 @@ const channelColor: Record<Channel, string> = {
 
 const clamp = (v: number, min: number, max: number) => (v < min ? min : v > max ? max : v);
 
-const sortPoints = (pts: Array<[number, number]>) =>
+const sortPoints = (pts: ChannelPoints) =>
   [...pts]
     .sort((a, b) => {
       const xDiff = a[0] - b[0];
@@ -63,17 +62,17 @@ const CANVAS_STYLE: React.CSSProperties = {
 
 type CurveGraphProps = {
   title: string;
-  channel: Channel;
-  points: Array<[number, number]> | undefined;
+  channel: ChannelName;
+  points: ChannelPoints | undefined;
   histogram: number[] | null;
-  onChange: (pts: Array<[number, number]>) => void;
+  onChange: (pts: ChannelPoints) => void;
 };
 
 const drawHistogram = (
   c: HTMLCanvasElement,
-  channel: Channel,
+  channel: ChannelName,
   histogram: number[] | null,
-  points: Array<[number, number]>
+  points: ChannelPoints
 ) => {
   // Use device pixel ratio for crisp rendering on HiDPI displays.
   const dpr = window.devicePixelRatio || 1;
@@ -207,7 +206,7 @@ const drawHistogram = (
   }
 };
 
-const getNearestPointIndex = (c: HTMLCanvasElement, points: Array<[number, number]>, mx: number, my: number) => {
+const getNearestPointIndex = (c: HTMLCanvasElement, points: ChannelPoints, mx: number, my: number) => {
   const cssWidth = c.clientWidth || CANVAS_WIDTH;
   const cssHeight = c.clientHeight || CANVAS_HEIGHT;
   const innerWidth = cssWidth - MARGIN_LEFT - MARGIN_RIGHT;
@@ -249,7 +248,7 @@ const canvasYToValueY = (c: HTMLCanvasElement, cy: number) => {
 const CurveGraph = memo(function CurveGraph(props: CurveGraphProps) {
   const { title, channel, points, histogram, onChange } = props;
   const canvasRef = useRef<HTMLCanvasElement | null>(null);
-  const [localPoints, setLocalPoints] = useState<Array<[number, number]>>(sortPoints(points ?? DEFAULT_POINTS));
+  const [localPoints, setLocalPoints] = useState<ChannelPoints>(sortPoints(points ?? DEFAULT_POINTS));
   const [dragIndex, setDragIndex] = useState<number | null>(null);
 
   useEffect(() => {
@@ -333,7 +332,7 @@ const CurveGraph = memo(function CurveGraph(props: CurveGraphProps) {
   );
 
   const commit = useCallback(
-    (pts: Array<[number, number]>) => {
+    (pts: ChannelPoints) => {
       onChange(sortPoints(pts));
     },
     [onChange]
@@ -534,17 +533,17 @@ export const RasterLayerCurvesAdjustmentsEditor = memo(() => {
   }, [layer?.objects, layer?.adjustments, recalcHistogram]);
 
   const onChangePoints = useCallback(
-    (channel: Channel, pts: Array<[number, number]>) => {
+    (channel: ChannelName, pts: ChannelPoints) => {
       dispatch(rasterLayerAdjustmentsCurvesUpdated({ entityIdentifier, channel, points: pts }));
     },
     [dispatch, entityIdentifier]
   );
 
   // Memoize per-channel change handlers to avoid inline lambdas in JSX
-  const onChangeMaster = useCallback((pts: Array<[number, number]>) => onChangePoints('master', pts), [onChangePoints]);
-  const onChangeR = useCallback((pts: Array<[number, number]>) => onChangePoints('r', pts), [onChangePoints]);
-  const onChangeG = useCallback((pts: Array<[number, number]>) => onChangePoints('g', pts), [onChangePoints]);
-  const onChangeB = useCallback((pts: Array<[number, number]>) => onChangePoints('b', pts), [onChangePoints]);
+  const onChangeMaster = useCallback((pts: ChannelPoints) => onChangePoints('master', pts), [onChangePoints]);
+  const onChangeR = useCallback((pts: ChannelPoints) => onChangePoints('r', pts), [onChangePoints]);
+  const onChangeG = useCallback((pts: ChannelPoints) => onChangePoints('g', pts), [onChangePoints]);
+  const onChangeB = useCallback((pts: ChannelPoints) => onChangePoints('b', pts), [onChangePoints]);
 
   return (
     <Flex

invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts

@@ -19,12 +19,16 @@ import type {
   CanvasEntityType,
   CanvasInpaintMaskState,
   CanvasMetadata,
+  ChannelName,
+  ChannelPoints,
   ControlLoRAConfig,
   EntityMovedByPayload,
   FillStyle,
   FLUXReduxImageInfluence,
+  RasterLayerAdjustments,
   RegionalGuidanceRefImageState,
   RgbColor,
+  SimpleConfig,
 } from 'features/controlLayers/store/types';
 import {
   calculateNewSize,
@@ -97,7 +101,6 @@ import {
   initialIPAdapter,
   initialT2IAdapter,
   makeDefaultRasterLayerAdjustments,
-  type RasterLayerAdjustments,
 } from './util';
 
 const slice = createSlice({
@@ -108,14 +111,7 @@ const slice = createSlice({
     //#region Raster layers
     rasterLayerAdjustmentsSet: (
       state,
-      action: PayloadAction<
-        EntityIdentifierPayload<
-          {
-            adjustments: RasterLayerAdjustments | null;
-          },
-          'raster_layer'
-        >
-      >
+      action: PayloadAction<EntityIdentifierPayload<{ adjustments: RasterLayerAdjustments | null }, 'raster_layer'>>
     ) => {
       const { entityIdentifier, adjustments } = action.payload;
       const layer = selectEntity(state, entityIdentifier);
@@ -152,14 +148,7 @@ const slice = createSlice({
     },
     rasterLayerAdjustmentsSimpleUpdated: (
       state,
-      action: PayloadAction<
-        EntityIdentifierPayload<
-          {
-            simple: Partial<RasterLayerAdjustments['simple']>;
-          },
-          'raster_layer'
-        >
-      >
+      action: PayloadAction<EntityIdentifierPayload<{ simple: Partial<SimpleConfig> }, 'raster_layer'>>
     ) => {
       const { entityIdentifier, simple } = action.payload;
       const layer = selectEntity(state, entityIdentifier);
@@ -173,15 +162,7 @@ const slice = createSlice({
     },
     rasterLayerAdjustmentsCurvesUpdated: (
       state,
-      action: PayloadAction<
-        EntityIdentifierPayload<
-          {
-            channel: 'master' | 'r' | 'g' | 'b';
-            points: Array<[number, number]>;
-          },
-          'raster_layer'
-        >
-      >
+      action: PayloadAction<EntityIdentifierPayload<{ channel: ChannelName; points: ChannelPoints }, 'raster_layer'>>
     ) => {
       const { entityIdentifier, channel, points } = action.payload;
       const layer = selectEntity(state, entityIdentifier);

invokeai/frontend/web/src/features/controlLayers/store/types.ts

@@ -378,36 +378,41 @@ const zControlLoRAConfig = z.object({
 });
 export type ControlLoRAConfig = z.infer<typeof zControlLoRAConfig>;
 
+const zSimpleConfig = z.object({
+  // All simple params normalized to [-1, 1] except sharpness [0, 1]
+  brightness: z.number().gte(-1).lte(1),
+  contrast: z.number().gte(-1).lte(1),
+  saturation: z.number().gte(-1).lte(1),
+  temperature: z.number().gte(-1).lte(1),
+  tint: z.number().gte(-1).lte(1),
+  sharpness: z.number().gte(0).lte(1),
+});
+export type SimpleConfig = z.infer<typeof zSimpleConfig>;
+
+const zUint8 = z.number().int().min(0).max(255);
+const zChannelPoints = z.array(z.tuple([zUint8, zUint8])).min(2);
+const zChannelName = z.enum(['master', 'r', 'g', 'b']);
+const zCurvesConfig = z.record(zChannelName, zChannelPoints);
+export type ChannelName = z.infer<typeof zChannelName>;
+export type ChannelPoints = z.infer<typeof zChannelPoints>;
+
+const zRasterLayerAdjustments = z.object({
+  version: z.literal(1),
+  enabled: z.boolean(),
+  collapsed: z.boolean(),
+  mode: z.enum(['simple', 'curves']),
+  simple: zSimpleConfig,
+  curves: zCurvesConfig,
+});
+export type RasterLayerAdjustments = z.infer<typeof zRasterLayerAdjustments>;
+
 const zCanvasRasterLayerState = zCanvasEntityBase.extend({
   type: z.literal('raster_layer'),
   position: zCoordinate,
   opacity: zOpacity,
   objects: z.array(zCanvasObjectState),
   // Optional per-layer color adjustments (simple + curves). When undefined, no adjustments are applied.
-  adjustments: z
-    .object({
-      version: z.literal(1),
-      enabled: z.boolean(),
-      collapsed: z.boolean(),
-      mode: z.enum(['simple', 'curves']),
-      simple: z.object({
-        // All simple params normalized to [-1, 1] except sharpness [0, 1]
-        brightness: z.number().gte(-1).lte(1),
-        contrast: z.number().gte(-1).lte(1),
-        saturation: z.number().gte(-1).lte(1),
-        temperature: z.number().gte(-1).lte(1),
-        tint: z.number().gte(-1).lte(1),
-        sharpness: z.number().gte(0).lte(1),
-      }),
-      curves: z.object({
-        // Curves are arrays of [x, y] control points in 0..255 space (no strict monotonic checks here)
-        master: z.array(z.tuple([z.number().int().gte(0).lte(255), z.number().int().gte(0).lte(255)])).min(2),
-        r: z.array(z.tuple([z.number().int().gte(0).lte(255), z.number().int().gte(0).lte(255)])).min(2),
-        g: z.array(z.tuple([z.number().int().gte(0).lte(255), z.number().int().gte(0).lte(255)])).min(2),
-        b: z.array(z.tuple([z.number().int().gte(0).lte(255), z.number().int().gte(0).lte(255)])).min(2),
-      }),
-    })
-    .optional(),
+  adjustments: zRasterLayerAdjustments.optional(),
 });
 export type CanvasRasterLayerState = z.infer<typeof zCanvasRasterLayerState>;
 

invokeai/frontend/web/src/features/controlLayers/store/util.ts

@@ -15,6 +15,7 @@ import type {
   Gemini2_5ReferenceImageConfig,
   ImageWithDims,
   IPAdapterConfig,
+  RasterLayerAdjustments,
   RefImageState,
   RgbColor,
   T2IAdapterConfig,
@@ -118,8 +119,6 @@ export const initialControlLoRA: ControlLoRAConfig = {
   weight: 0.75,
 };
 
-export type RasterLayerAdjustments = NonNullable<CanvasRasterLayerState['adjustments']>;
-
 export const makeDefaultRasterLayerAdjustments = (mode: 'simple' | 'curves' = 'simple'): RasterLayerAdjustments => ({
   version: 1,
   enabled: true,