Add Test Mode toggle

bahamutww · Apr 20, 2024 · 3a6a3cc · 3a6a3cc
1 parent a68dab8
commit 3a6a3cc
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 11 deletions.
diff --git a/src/components/ImageGeneration/GenerationForm/GenerationForm.tsx b/src/components/ImageGeneration/GenerationForm/GenerationForm.tsx
@@ -714,6 +714,9 @@ const GenerationFormInner = ({ onSuccess }: { onSuccess?: () => void }) => {
                       resources: [{ type: ModelType.VAE, baseModelSet: baseModel }],
                     }}
                   />
+                  {currentUser?.isModerator && (
+                    <InputSwitch name="staging" label="Test Mode" labelPosition="left" />
+                  )}
                 </Stack>
               </Accordion.Panel>
             </Accordion.Item>

diff --git a/src/components/ImageGeneration/GenerationForm/generation.utils.ts b/src/components/ImageGeneration/GenerationForm/generation.utils.ts
@@ -132,7 +132,7 @@ export const useEstimateTextToImageJobCost = () => {
   const input = useGenerationFormStore(
     useCallback(
       (state) => {
-        const { aspectRatio, steps, quantity, sampler, draft } = state;
+        const { aspectRatio, steps, quantity, sampler, draft, staging } = state;
         if (!status.charge || !baseModel) return null;
 
         return {
@@ -141,6 +141,7 @@ export const useEstimateTextToImageJobCost = () => {
           steps: steps ?? generation.defaultValues.steps,
           quantity: quantity ?? generation.defaultValues.quantity,
           sampler: sampler ?? generation.defaultValues.sampler,
+          staging,
           draft,
         };
       },

diff --git a/src/components/ImageGeneration/utils/generationRequestHooks.ts b/src/components/ImageGeneration/utils/generationRequestHooks.ts
@@ -44,7 +44,7 @@ export const useGetGenerationRequests = (
     if (!isLoading) updateFromEvents();
   }, [isLoading]);
 
-  return { data, requests, images, isLoading, ...rest };
+  return { data, requests, images, isLoading: !currentUser ? false : isLoading, ...rest };
 };
 
 export const updateGenerationRequest = (

diff --git a/src/server/schema/generation.schema.ts b/src/server/schema/generation.schema.ts
@@ -135,6 +135,7 @@ const sharedGenerationParamsSchema = z.object({
   quantity: z.coerce.number().min(1).max(20),
   nsfw: z.boolean().optional(),
   draft: z.boolean().optional(),
+  staging: z.boolean().optional(),
   baseModel: z.string().optional(),
   aspectRatio: z.string(),
 });

diff --git a/src/server/services/generation/generation.service.ts b/src/server/services/generation/generation.service.ts
@@ -785,15 +785,15 @@ export const createGenerationRequest = async ({
   // console.log(JSON.stringify(generationRequest));
   // console.log('________');
 
-  const charge = status.charge;
-  const response = await fetch(
-    `${env.SCHEDULER_ENDPOINT}/requests${charge ? '?charge=true' : ''}`,
-    {
-      method: 'POST',
-      headers: { 'Content-Type': 'application/json' },
-      body: JSON.stringify(generationRequest),
-    }
-  );
+  const schedulerUrl = new URL(`${env.SCHEDULER_ENDPOINT}/requests`);
+  if (status.charge) schedulerUrl.searchParams.set('charge', 'true');
+  if (params.staging) schedulerUrl.searchParams.set('staging', 'true');
+
+  const response = await fetch(schedulerUrl.toString(), {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify(generationRequest),
+  });
 
   // console.log('________');
   // console.log(response);