From 5a3b0497c1f4a2c3eced8a9eaeb4833c4911a6e8 Mon Sep 17 00:00:00 2001
From: Jong Eun Lee <jongeun@lablup.com>
Date: Mon, 6 Jan 2025 17:09:44 +0800
Subject: [PATCH] improve image generation model endpoint generation and
 realted UX

---
 react/src/components/ChatContent.tsx          | 10 +++++++-
 react/src/components/ModelCardChat.tsx        |  1 +
 .../lablupTalkativotUI/ChatUIModal.tsx        | 10 +++++++-
 .../lablupTalkativotUI/LLMChatCard.tsx        | 24 +++++++++++++++----
 4 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/react/src/components/ChatContent.tsx b/react/src/components/ChatContent.tsx
index ff6f8f159c..b66d97b367 100644
--- a/react/src/components/ChatContent.tsx
+++ b/react/src/components/ChatContent.tsx
@@ -13,12 +13,14 @@ import { useLazyLoadQuery } from 'react-relay/hooks';
 interface ChatContentProps {
   endpointId: string;
   endpointUrl: string;
+  endpointName: string;
   basePath: string;
 }
 
 const ChatContent: React.FC<ChatContentProps> = ({
   endpointId,
   endpointUrl,
+  endpointName,
   basePath,
 }) => {
   const { t } = useTranslation();
@@ -56,6 +58,7 @@ const ChatContent: React.FC<ChatContentProps> = ({
         fetchPolicy: 'network-only',
       },
     );
+  const isTextToImageModel = _.includes(endpointName, 'stable-diffusion');
 
   const newestValidToken =
     _.orderBy(endpoint_token_list?.items, ['valid_until'], ['desc'])[0]
@@ -85,7 +88,12 @@ const ChatContent: React.FC<ChatContentProps> = ({
   return (
     <LLMChatCard
       endpointId={endpointId || ''}
-      baseURL={new URL(basePath, endpointUrl).toString()}
+      baseURL={
+        isTextToImageModel
+          ? new URL('/generate-image', endpointUrl || '').toString()
+          : new URL(basePath, endpointUrl || '').toString()
+      }
+      isImageGeneration={isTextToImageModel}
       models={_.map(modelsResult?.data, (m) => ({
         id: m.id,
         name: m.id,
diff --git a/react/src/components/ModelCardChat.tsx b/react/src/components/ModelCardChat.tsx
index cd4728e443..5b1d6c6600 100644
--- a/react/src/components/ModelCardChat.tsx
+++ b/react/src/components/ModelCardChat.tsx
@@ -67,6 +67,7 @@ const ModelCardChat: React.FC<ModelCardChatProps> = ({
     <ChatContent
       endpointId={healthyEndpoint[0]?.endpoint_id as string}
       endpointUrl={healthyEndpoint[0]?.url as string}
+      endpointName={healthyEndpoint[0]?.name as string}
       basePath={basePath}
     />
   ) : (
diff --git a/react/src/components/lablupTalkativotUI/ChatUIModal.tsx b/react/src/components/lablupTalkativotUI/ChatUIModal.tsx
index 48d897b601..154aad9c6c 100644
--- a/react/src/components/lablupTalkativotUI/ChatUIModal.tsx
+++ b/react/src/components/lablupTalkativotUI/ChatUIModal.tsx
@@ -71,6 +71,7 @@ const EndpointChatContent: React.FC<ChatUIBasicProps> = ({
     graphql`
       fragment ChatUIModalFragment on Endpoint {
         endpoint_id
+        name
         url
         status
       }
@@ -90,6 +91,7 @@ const EndpointChatContent: React.FC<ChatUIBasicProps> = ({
     `,
     endpointTokenFrgmt,
   );
+  const isTextToImageModel = _.includes(endpoint?.name, 'stable-diffusion');
 
   const newestToken = _.maxBy(
     endpointTokenList?.items,
@@ -124,7 +126,12 @@ const EndpointChatContent: React.FC<ChatUIBasicProps> = ({
   ) : (
     <LLMChatCard
       endpointId={endpoint?.endpoint_id || ''}
-      baseURL={new URL(basePath, endpoint?.url || '').toString()}
+      baseURL={
+        isTextToImageModel
+          ? new URL('/generate-image', endpoint?.url || '').toString()
+          : new URL(basePath, endpoint?.url || '').toString()
+      }
+      isImageGeneration={isTextToImageModel}
       models={_.map(modelsResult?.data, (m) => ({
         id: m.id,
         name: m.id,
@@ -133,6 +140,7 @@ const EndpointChatContent: React.FC<ChatUIBasicProps> = ({
       style={{ flex: 1 }}
       allowCustomModel={_.isEmpty(modelsResult?.data)}
       alert={
+        !isTextToImageModel &&
         _.isEmpty(modelsResult?.data) && (
           <Alert
             type="warning"
diff --git a/react/src/components/lablupTalkativotUI/LLMChatCard.tsx b/react/src/components/lablupTalkativotUI/LLMChatCard.tsx
index a8a208b91a..46d55e8f0e 100644
--- a/react/src/components/lablupTalkativotUI/LLMChatCard.tsx
+++ b/react/src/components/lablupTalkativotUI/LLMChatCard.tsx
@@ -230,7 +230,7 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
     setLoadingImageGeneration(true);
     try {
       const response = await fetch(
-        'https://stable-diffusion-3m.asia03.app.backend.ai/generate-image',
+        customModelFormRef.current?.getFieldValue('baseURL'),
         {
           method: 'POST',
           headers: {
@@ -244,7 +244,9 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
       );
       if (response.ok) {
         const responseData = await response.json();
-        return 'data:image/png;base64,' + responseData.image_base64;
+        return _.startsWith(responseData.image_base64, 'data:image/png;base64,')
+          ? responseData.image_base64
+          : 'data:image/png;base64,' + responseData.image_base64;
       } else {
         throw new Error('Error generating image');
       }
@@ -414,8 +416,8 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
               });
 
               if (isImageGeneration) {
+                const generationId = _.uniqueId();
                 try {
-                  const imageBase64 = await generateImage(input, 'accessKey');
                   setMessages((prevMessages) => [
                     ...prevMessages,
                     {
@@ -424,7 +426,20 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
                       content: input,
                     },
                     {
-                      id: _.uniqueId(),
+                      id: generationId,
+                      role: 'assistant',
+                      content: 'Processing...',
+                    },
+                  ]);
+                  setInput('');
+                  const imageBase64 = await generateImage(input, 'accessKey');
+                  setMessages((prevMessages) => [
+                    ..._.filter(
+                      prevMessages,
+                      (message) => message.id !== generationId,
+                    ),
+                    {
+                      id: generationId,
                       role: 'assistant',
                       content: '',
                       experimental_attachments: [
@@ -510,6 +525,7 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
                 required: true,
               },
             ]}
+            hidden={isImageGeneration}
           >
             <Input placeholder="llm-model" />
           </Form.Item>