feat: support upscale and variation in midjourney, support image zoom (…

…#64)
GPTGenius · May 24, 2023 · 9c280ba · 9c280ba
1 parent 57334e4
commit 9c280ba
Show file tree

Hide file tree

Showing 11 changed files with 434 additions and 65 deletions.
diff --git a/package.json b/package.json
@@ -39,7 +39,8 @@
     "markdown-it": "^13.0.1",
     "markdown-it-highlightjs": "^4.0.1",
     "markdown-it-kbd": "^2.2.2",
-    "midjourney-fetch": "0.1.4",
+    "medium-zoom": "^1.0.8",
+    "midjourney-fetch": "1.0.0",
     "react": "^18.0.0",
     "react-dom": "^18.0.0",
     "replicate-fetch": "^0.1.1",

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/src/components/MessageBox/index.tsx b/src/components/MessageBox/index.tsx
@@ -1,18 +1,27 @@
 import { FC, useCallback, useContext, useEffect } from 'react';
 import { throttle } from 'lodash-es';
+import type { MessageType } from 'midjourney-fetch';
 import GlobalContext from '@contexts/global';
 import { ConversationMode, Message } from '@interfaces';
 import markdown from '@utils/markdown';
 import { getRelativeTime } from '@utils/date';
 import SystemAvatar from '@components/Avatar/system';
 import useCopyCode from '@hooks/useCopyCode';
+import MidjourneyOperations from '@components/MidjourneyOperations';
+import { hasUpscaleOrVariation } from '@utils/midjourney';
 import './index.css';
 
 const MessageItem: FC<{
   message: Message;
+  onOperationClick?: (
+    type: MessageType,
+    customId: string,
+    messageId: string,
+    prompt: string
+  ) => void;
   mode?: ConversationMode;
   index?: number;
-}> = ({ message, mode, index }) => {
+}> = ({ message, onOperationClick, mode, index }) => {
   const { i18n } = useContext(GlobalContext);
   const isExpired = message.expiredAt && message.expiredAt <= Date.now();
   const createdAt = getRelativeTime(message.createdAt, true);
@@ -25,27 +34,43 @@ const MessageItem: FC<{
       {message.role === 'assistant' ? (
         <SystemAvatar className="mt-[14px] mr-2" role={message.imageModel} />
       ) : null}
-      <div
-        dangerouslySetInnerHTML={{
-          __html: isExpired
-            ? i18n.status_image_expired
-            : markdown.render(message.content),
-        }}
-        className={`prose message-box shadow-sm p-4 ${
-          message.role === 'user' ? 'bg-gradient text-white' : 'bg-[#ebeced]'
-        } ${
-          mode === 'image' ? 'img-no-margin' : ''
-        } break-words overflow-hidden rounded-[16px]`}
-      />
-      {createdAt ? (
+      <div>
         <div
-          className={`message-box-time hover:visible invisible text-[#a1a7a8] text-sm absolute top-[-20px] ${
-            message.role === 'user' ? 'right-0' : 'left-[calc(32px+0.5rem)]'
-          }`}
-        >
-          {createdAt}
-        </div>
-      ) : null}
+          dangerouslySetInnerHTML={{
+            __html: isExpired
+              ? i18n.status_image_expired
+              : markdown.render(message.content),
+          }}
+          className={`prose message-box shadow-sm p-4 ${
+            message.role === 'user' ? 'bg-gradient text-white' : 'bg-[#ebeced]'
+          } ${
+            mode === 'image' ? 'img-no-margin' : ''
+          } break-words overflow-hidden rounded-[16px]`}
+        />
+        {createdAt ? (
+          <div
+            className={`message-box-time hover:visible invisible text-[#a1a7a8] text-sm absolute top-[-20px] ${
+              message.role === 'user' ? 'right-0' : 'left-[calc(32px+0.5rem)]'
+            }`}
+          >
+            {createdAt}
+          </div>
+        ) : null}
+        {message.midjourneyMessage &&
+        hasUpscaleOrVariation(message.midjourneyMessage) ? (
+          <MidjourneyOperations
+            message={message.midjourneyMessage}
+            onClick={(type, customId) =>
+              onOperationClick(
+                type,
+                customId,
+                message.midjourneyMessage.id,
+                message.midjourneyMessage.prompt
+              )
+            }
+          />
+        ) : null}
+      </div>
     </div>
   );
 };
@@ -55,7 +80,13 @@ const MessageBox: FC<{
   messages: Message[];
   mode: ConversationMode;
   loading: boolean;
-}> = ({ streamMessage, messages, mode, loading }) => {
+  onOperationClick?: (
+    type: MessageType,
+    customId: string,
+    messageId: string,
+    prompt: string
+  ) => void;
+}> = ({ streamMessage, messages, mode, loading, onOperationClick }) => {
   const { i18n } = useContext(GlobalContext);
 
   useCopyCode(i18n.success_copy);
@@ -100,7 +131,13 @@ const MessageBox: FC<{
         />
       ) : null}
       {messages.map((message, index) => (
-        <MessageItem key={index} index={index} mode={mode} message={message} />
+        <MessageItem
+          key={index}
+          index={index}
+          mode={mode}
+          message={message}
+          onOperationClick={loading ? () => null : onOperationClick}
+        />
       ))}
       {streamMessage ? (
         <MessageItem message={{ role: 'assistant', content: streamMessage }} />

diff --git a/src/components/MidjourneyOperations/index.tsx b/src/components/MidjourneyOperations/index.tsx
@@ -0,0 +1,60 @@
+import { FC } from 'react';
+import { Tag } from 'antd';
+import type { MessageType } from 'midjourney-fetch';
+import { MidjourneyMessage } from '@interfaces';
+import {
+  filterUpscale,
+  filterVariation,
+  isComponentAvailable,
+} from '@utils/midjourney';
+
+const { CheckableTag } = Tag;
+
+const MidjourneyOperations: FC<{
+  message: MidjourneyMessage;
+  onClick?: (type: MessageType, id: string) => void;
+}> = ({ message, onClick }) => {
+  const upscale = filterUpscale(message.components[0]?.components ?? []);
+  const variation = filterVariation([
+    ...(message.components[0]?.components ?? []),
+    ...(message.components[1]?.components ?? []),
+  ]);
+  return (
+    <div>
+      {upscale.map((option) => {
+        const isAvailable = isComponentAvailable(option);
+        return (
+          <CheckableTag
+            className="mt-3"
+            key={option.custom_id || option.label}
+            checked={!isAvailable}
+            style={isAvailable ? { background: 'rgba(0, 0, 0, 0.06)' } : {}}
+            onClick={() =>
+              isAvailable && onClick?.('upscale', option.custom_id)
+            }
+          >
+            {option.label || option.emoji?.name}
+          </CheckableTag>
+        );
+      })}
+      {variation.map((option) => {
+        const isAvailable = isComponentAvailable(option);
+        return (
+          <CheckableTag
+            className="mt-3"
+            key={option.custom_id}
+            checked={!isAvailable}
+            style={isAvailable ? { background: 'rgba(0, 0, 0, 0.06)' } : {}}
+            onClick={() =>
+              isAvailable && onClick?.('variation', option.custom_id)
+            }
+          >
+            {option.label || option.emoji?.name}
+          </CheckableTag>
+        );
+      })}
+    </div>
+  );
+};
+
+export default MidjourneyOperations;
diff --git a/src/interfaces/index.ts b/src/interfaces/index.ts
@@ -1,14 +1,23 @@
+import type { MessageItem } from 'midjourney-fetch';
 import {
   LayoutConfig,
   SupportedImageModels,
   SupportedImgSize,
   SupportedModel,
 } from '@configs';
 
+export type MidjourneyMessage = Pick<
+  MessageItem,
+  'id' | 'components' | 'attachments'
+> & {
+  prompt: string;
+};
+
 export interface Message {
   content: string;
   role: 'assistant' | 'user';
   imageModel?: SupportedImageModels; // distinguish avator
+  midjourneyMessage?: MidjourneyMessage;
   createdAt?: number;
   expiredAt?: number; // for image mode
 }

diff --git a/src/layouts/Layout.astro b/src/layouts/Layout.astro
@@ -179,6 +179,13 @@ const { title } = Astro.props;
     overflow: hidden;
     text-overflow: ellipsis;
   }
+  /** medium zoom */
+  .medium-zoom-overlay {
+    z-index: 20;
+  }
+  .medium-zoom-image {
+    z-index: 21;
+  }
   /* mobile style */
   @media screen and (max-width: 768px) {
     html {