From 65fb6af01f59af6fe27416875e7408c00f40dd22 Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Wed, 30 Aug 2023 10:15:26 +1200
Subject: [PATCH] ui: Make aspect ratio logic more robust

---
 .../BoundingBox/ParamBoundingBoxSize.tsx      | 54 +++++++++++++------
 .../Parameters/Core/ParamAspectRatio.tsx      |  2 +
 .../components/Parameters/Core/ParamSize.tsx  | 25 +++++++--
 .../parameters/store/generationSlice.ts       |  9 ----
 4 files changed, 62 insertions(+), 28 deletions(-)

diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/BoundingBox/ParamBoundingBoxSize.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/BoundingBox/ParamBoundingBoxSize.tsx
index 7725461fd3..1c1f60bc09 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/BoundingBox/ParamBoundingBoxSize.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/BoundingBox/ParamBoundingBoxSize.tsx
@@ -13,7 +13,9 @@ import { useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { FaLock } from 'react-icons/fa';
 import { MdOutlineSwapVert } from 'react-icons/md';
-import ParamAspectRatio from '../../Core/ParamAspectRatio';
+import ParamAspectRatio, {
+  mappedAspectRatios,
+} from '../../Core/ParamAspectRatio';
 import ParamBoundingBoxHeight from './ParamBoundingBoxHeight';
 import ParamBoundingBoxWidth from './ParamBoundingBoxWidth';
 
@@ -39,12 +41,43 @@ export default function ParamBoundingBoxSize() {
     useAppSelector(sizeOptsSelector);
 
   const handleLockRatio = useCallback(() => {
-    dispatch(
-      setAspectRatio(boundingBoxDimensions.width / boundingBoxDimensions.height)
-    );
-    dispatch(setShouldLockAspectRatio(!shouldLockAspectRatio));
+    if (shouldLockAspectRatio) {
+      dispatch(setShouldLockAspectRatio(false));
+      if (
+        !mappedAspectRatios.includes(
+          boundingBoxDimensions.width / boundingBoxDimensions.height
+        )
+      ) {
+        dispatch(setAspectRatio(null));
+      } else {
+        dispatch(
+          setAspectRatio(
+            boundingBoxDimensions.width / boundingBoxDimensions.height
+          )
+        );
+      }
+    } else {
+      dispatch(setShouldLockAspectRatio(true));
+      dispatch(
+        setAspectRatio(
+          boundingBoxDimensions.width / boundingBoxDimensions.height
+        )
+      );
+    }
   }, [shouldLockAspectRatio, boundingBoxDimensions, dispatch]);
 
+  const handleToggleSize = useCallback(() => {
+    dispatch(flipBoundingBoxAxes());
+    dispatch(setAspectRatio(null));
+    if (shouldLockAspectRatio) {
+      dispatch(
+        setAspectRatio(
+          boundingBoxDimensions.height / boundingBoxDimensions.width
+        )
+      );
+    }
+  }, [dispatch, shouldLockAspectRatio, boundingBoxDimensions]);
+
   return (
     <Flex
       sx={{
@@ -80,16 +113,7 @@ export default function ParamBoundingBoxSize() {
           size="sm"
           icon={<MdOutlineSwapVert />}
           fontSize={20}
-          onClick={() => {
-            dispatch(flipBoundingBoxAxes());
-            if (
-              ![null, 2 / 3, 16 / 9, 1 / 1].includes(
-                boundingBoxDimensions.height / boundingBoxDimensions.width
-              )
-            ) {
-              dispatch(setAspectRatio(null));
-            }
-          }}
+          onClick={handleToggleSize}
         />
         <IAIIconButton
           tooltip={t('ui.lockRatio')}
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamAspectRatio.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamAspectRatio.tsx
index bf2782b525..05e0b09cee 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamAspectRatio.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamAspectRatio.tsx
@@ -15,6 +15,8 @@ const aspectRatios = [
   { name: '1:1', value: 1 / 1 },
 ];
 
+export const mappedAspectRatios = aspectRatios.map((ar) => ar.value);
+
 export default function ParamAspectRatio() {
   const aspectRatio = useAppSelector(
     (state: RootState) => state.generation.aspectRatio
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamSize.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamSize.tsx
index 6ebcb0beb1..63700c4922 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamSize.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamSize.tsx
@@ -13,7 +13,7 @@ import { useTranslation } from 'react-i18next';
 import { FaLock } from 'react-icons/fa';
 import { MdOutlineSwapVert } from 'react-icons/md';
 import { activeTabNameSelector } from '../../../../ui/store/uiSelectors';
-import ParamAspectRatio from './ParamAspectRatio';
+import ParamAspectRatio, { mappedAspectRatios } from './ParamAspectRatio';
 import ParamHeight from './ParamHeight';
 import ParamWidth from './ParamWidth';
 
@@ -45,10 +45,27 @@ export default function ParamSize() {
   } = useAppSelector(sizeOptsSelector);
 
   const handleLockRatio = useCallback(() => {
-    dispatch(setAspectRatio(width / height));
-    dispatch(setShouldLockAspectRatio(!shouldLockAspectRatio));
+    if (shouldLockAspectRatio) {
+      dispatch(setShouldLockAspectRatio(false));
+      if (!mappedAspectRatios.includes(width / height)) {
+        dispatch(setAspectRatio(null));
+      } else {
+        dispatch(setAspectRatio(width / height));
+      }
+    } else {
+      dispatch(setShouldLockAspectRatio(true));
+      dispatch(setAspectRatio(width / height));
+    }
   }, [shouldLockAspectRatio, width, height, dispatch]);
 
+  const handleToggleSize = useCallback(() => {
+    dispatch(toggleSize());
+    dispatch(setAspectRatio(null));
+    if (shouldLockAspectRatio) {
+      dispatch(setAspectRatio(height / width));
+    }
+  }, [dispatch, shouldLockAspectRatio, width, height]);
+
   return (
     <Flex
       sx={{
@@ -87,7 +104,7 @@ export default function ParamSize() {
           isDisabled={
             activeTabName === 'img2img' ? !shouldFitToWidthHeight : false
           }
-          onClick={() => dispatch(toggleSize())}
+          onClick={handleToggleSize}
         />
         <IAIIconButton
           tooltip={t('ui.lockRatio')}
diff --git a/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts b/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
index c7b8d2b41f..5a830b6e95 100644
--- a/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
+++ b/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
@@ -152,9 +152,6 @@ export const generationSlice = createSlice({
     },
     toggleSize: (state) => {
       const [width, height] = [state.width, state.height];
-      if (![null, 2 / 3, 16 / 9, 1 / 1].includes(height / width)) {
-        state.aspectRatio = null;
-      }
       state.width = height;
       state.height = width;
     },
@@ -278,12 +275,6 @@ export const generationSlice = createSlice({
       }
     },
     setShouldLockAspectRatio: (state, action: PayloadAction<boolean>) => {
-      if (
-        action.payload === false &&
-        ![null, 2 / 3, 16 / 9, 1 / 1].includes(state.aspectRatio)
-      ) {
-        state.aspectRatio = null;
-      }
       state.shouldLockAspectRatio = action.payload;
     },
   },