Adds new page for image generation (its ugly)

.env-example

@@ -124,8 +124,8 @@ NEXT_PUBLIC_SEARCH_CLIENT_KEY=aSampleKey
 
 # Scheduler endpoint
 SCHEDULER_ENDPOINT=url
-GENERATION_ENDPOINT=url
-ORCHESTRATOR_TOKEN=cooltoken
+ORCHESTRATOR_ENDPOINT=url
+ORCHESTRATOR_ACCESS_TOKEN=cooltoken
 
 # Civitai Buzz
 BUZZ_ENDPOINT=https://localhost

src/components/ImageGeneration/GenerationForm/BaseModelProvider.tsx

@@ -41,7 +41,7 @@ export function BaseModelProvider<T extends FieldValues>({
 
   return (
     <BaseModelsContext.Provider value={{ baseModels, baseModel }}>
-      <InputText type="hidden" name="baseModel" />
+      <InputText type="hidden" name="baseModel" clearable={false} hidden />
       {children({ baseModel })}
     </BaseModelsContext.Provider>
   );

src/env/schema.mjs

@@ -78,9 +78,9 @@ export const serverSchema = z.object({
   TRPC_ORIGINS: commaDelimitedStringArray().optional(),
   CANNY_SECRET: z.string().optional(),
   SCHEDULER_ENDPOINT: z.string().url().optional(),
-  GENERATION_ENDPOINT: z.string().url().optional(),
+  ORCHESTRATOR_ENDPOINT: z.string().url().optional(),
   GENERATION_CALLBACK_HOST: z.string().url().optional(),
-  ORCHESTRATOR_TOKEN: z.string().optional(),
+  ORCHESTRATOR_ACCESS_TOKEN: z.string().optional(),
   AXIOM_TOKEN: z.string().optional(),
   AXIOM_ORG_ID: z.string().optional(),
   AXIOM_DATASTREAM: z.string().optional(),

src/pages/generate/index.tsx

@@ -0,0 +1,69 @@
+import { Center, Container, Grid, Stack, Tabs, Text, ThemeIcon } from '@mantine/core';
+import { IconLock } from '@tabler/icons-react';
+import { Feed } from '~/components/ImageGeneration/Feed';
+import { GenerateFormLogic } from '~/components/ImageGeneration/GenerationForm/GenerateFormLogic';
+import { Queue } from '~/components/ImageGeneration/Queue';
+import { useGetGenerationRequests } from '~/components/ImageGeneration/utils/generationRequestHooks';
+import { IsClient } from '~/components/IsClient/IsClient';
+import { useCurrentUser } from '~/hooks/useCurrentUser';
+import { createServerSideProps } from '~/server/utils/server-side-helpers';
+import { getLoginLink } from '~/utils/login-helpers';
+
+export const getServerSideProps = createServerSideProps({
+  useSession: true,
+  resolver: async ({ session, features, ctx }) => {
+    if (!session)
+      return {
+        redirect: {
+          destination: getLoginLink({ returnUrl: ctx.req.url }),
+          permanent: false,
+        },
+      };
+
+    if (!features?.imageGeneration) return { notFound: true };
+  },
+});
+
+export default function GeneratePage() {
+  const currentUser = useCurrentUser();
+
+  const result = useGetGenerationRequests({});
+
+  if (currentUser?.muted)
+    return (
+      <Center h="100%" w="75%" mx="auto">
+        <Stack spacing="xl" align="center">
+          <ThemeIcon size="xl" radius="xl" color="yellow">
+            <IconLock />
+          </ThemeIcon>
+          <Text align="center">You cannot create new generations because you have been muted</Text>
+        </Stack>
+      </Center>
+    );
+
+  return (
+    <Container size="lg">
+      <Grid gutter={48}>
+        <Grid.Col span={5} maw={400}>
+          <IsClient>
+            <GenerateFormLogic />
+          </IsClient>
+        </Grid.Col>
+        <Grid.Col span={7} sx={{ maxWidth: 'unset', flexGrow: 1 }}>
+          <Tabs variant="pills" defaultValue="queue" radius="xl" color="gray">
+            <Tabs.List>
+              <Tabs.Tab value="queue">Queue</Tabs.Tab>
+              <Tabs.Tab value="feed">Feed</Tabs.Tab>
+            </Tabs.List>
+            <Tabs.Panel value="queue">
+              <Queue {...result} />
+            </Tabs.Panel>
+            <Tabs.Panel value="feed">
+              <Feed {...result} />
+            </Tabs.Panel>
+          </Tabs>
+        </Grid.Col>
+      </Grid>
+    </Container>
+  );
+}

src/server/services/training.service.ts

@@ -45,8 +45,9 @@ const assetUrlRegex =
   /\/v\d\/consumer\/jobs\/(?<jobId>[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})\/assets\/(?<assetName>\S+)$/i;
 
 export const moveAsset = async ({ url, modelId }: MoveAssetInput) => {
-  if (!env.GENERATION_ENDPOINT) throw throwBadRequestError('Missing GENERATION_ENDPOINT env');
-  if (!env.ORCHESTRATOR_TOKEN) throw throwBadRequestError('Missing ORCHESTRATOR_TOKEN env');
+  if (!env.ORCHESTRATOR_ENDPOINT) throw throwBadRequestError('Missing ORCHESTRATOR_ENDPOINT env');
+  if (!env.ORCHESTRATOR_ACCESS_TOKEN)
+    throw throwBadRequestError('Missing ORCHESTRATOR_ACCESS_TOKEN env');
 
   const urlMatch = url.match(assetUrlRegex);
   if (!urlMatch || !urlMatch.groups) throw throwBadRequestError('Invalid URL');
@@ -61,11 +62,11 @@ export const moveAsset = async ({ url, modelId }: MoveAssetInput) => {
     destinationUri,
   };
 
-  const response = await fetch(`${env.GENERATION_ENDPOINT}/v1/consumer/jobs?wait=true`, {
+  const response = await fetch(`${env.ORCHESTRATOR_ENDPOINT}/v1/consumer/jobs?wait=true`, {
     method: 'POST',
     headers: {
       'Content-Type': 'application/json',
-      Authorization: `Bearer ${env.ORCHESTRATOR_TOKEN}`,
+      Authorization: `Bearer ${env.ORCHESTRATOR_ACCESS_TOKEN}`,
     },
     body: JSON.stringify(reqBody),
   });
@@ -93,19 +94,20 @@ export const moveAsset = async ({ url, modelId }: MoveAssetInput) => {
 };
 
 export const deleteAssets = async (jobId: string) => {
-  if (!env.GENERATION_ENDPOINT) throw throwBadRequestError('Missing GENERATION_ENDPOINT env');
-  if (!env.ORCHESTRATOR_TOKEN) throw throwBadRequestError('Missing ORCHESTRATOR_TOKEN env');
+  if (!env.ORCHESTRATOR_ENDPOINT) throw throwBadRequestError('Missing ORCHESTRATOR_ENDPOINT env');
+  if (!env.ORCHESTRATOR_ACCESS_TOKEN)
+    throw throwBadRequestError('Missing ORCHESTRATOR_ACCESS_TOKEN env');
 
   const reqBody = {
     $type: 'clearAssets',
     jobId,
   };
 
-  const response = await fetch(`${env.GENERATION_ENDPOINT}/v1/consumer/jobs?wait=true`, {
+  const response = await fetch(`${env.ORCHESTRATOR_ENDPOINT}/v1/consumer/jobs?wait=true`, {
     method: 'POST',
     headers: {
       'Content-Type': 'application/json',
-      Authorization: `Bearer ${env.ORCHESTRATOR_TOKEN}`,
+      Authorization: `Bearer ${env.ORCHESTRATOR_ACCESS_TOKEN}`,
     },
     body: JSON.stringify(reqBody),
   });
@@ -126,8 +128,9 @@ export const createTrainingRequest = async ({
   userId,
   modelVersionId,
 }: CreateTrainingRequestInput & { userId: number }) => {
-  if (!env.GENERATION_ENDPOINT) throw throwBadRequestError('Missing GENERATION_ENDPOINT env');
-  if (!env.ORCHESTRATOR_TOKEN) throw throwBadRequestError('Missing ORCHESTRATOR_TOKEN env');
+  if (!env.ORCHESTRATOR_ENDPOINT) throw throwBadRequestError('Missing ORCHESTRATOR_ENDPOINT env');
+  if (!env.ORCHESTRATOR_ACCESS_TOKEN)
+    throw throwBadRequestError('Missing ORCHESTRATOR_ACCESS_TOKEN env');
 
   const modelVersions = await dbWrite.$queryRaw<TrainingRequest[]>`
     SELECT mv."trainingDetails",
@@ -216,11 +219,11 @@ export const createTrainingRequest = async ({
 
   // console.log(JSON.stringify(generationRequest));
 
-  const response = await fetch(`${env.GENERATION_ENDPOINT}/v1/consumer/jobs`, {
+  const response = await fetch(`${env.ORCHESTRATOR_ENDPOINT}/v1/consumer/jobs`, {
     method: 'POST',
     headers: {
       'Content-Type': 'application/json',
-      Authorization: `Bearer ${env.ORCHESTRATOR_TOKEN}`,
+      Authorization: `Bearer ${env.ORCHESTRATOR_ACCESS_TOKEN}`,
     },
     body: JSON.stringify(generationRequest),
   });