From 324bd57e3dcabba8819c679472444144859f7612 Mon Sep 17 00:00:00 2001
From: jialin <jialinkuang@126.com>
Date: Wed, 27 Nov 2024 21:12:49 +0800
Subject: [PATCH] fix: image chunk data broken

---
 src/components/speech-content/speech-item.tsx | 13 ++++++
 .../components/ground-embedding.tsx           |  2 +-
 .../playground/components/ground-images.tsx   | 31 ++++++++++---
 .../playground/components/ground-reranker.tsx |  2 +-
 src/pages/playground/config/params-config.ts  |  4 +-
 src/utils/fetch-chunk-data.ts                 | 46 ++++++++++++++++++-
 6 files changed, 86 insertions(+), 12 deletions(-)

diff --git a/src/components/speech-content/speech-item.tsx b/src/components/speech-content/speech-item.tsx
index 247a99cf..3b278044 100644
--- a/src/components/speech-content/speech-item.tsx
+++ b/src/components/speech-content/speech-item.tsx
@@ -28,6 +28,18 @@ const SpeechItem: React.FC<SpeechContentProps> = (props) => {
     setCollapsed(!collapsed);
   };
 
+  const onDownload = () => {
+    const url = props.audioUrl || '';
+    const filename = Date.now() + '';
+
+    const link = document.createElement('a');
+    link.href = url;
+    link.download = filename;
+    document.body.appendChild(link);
+    link.click();
+    link.remove();
+  };
+
   return (
     <div>
       <div className="speech-item">
@@ -60,6 +72,7 @@ const SpeechItem: React.FC<SpeechContentProps> = (props) => {
           </Tooltip>
           <Tooltip title="Download">
             <Button
+              onClick={onDownload}
               icon={<DownloadOutlined />}
               type="text"
               size="small"
diff --git a/src/pages/playground/components/ground-embedding.tsx b/src/pages/playground/components/ground-embedding.tsx
index 90082601..0e688bdc 100644
--- a/src/pages/playground/components/ground-embedding.tsx
+++ b/src/pages/playground/components/ground-embedding.tsx
@@ -282,7 +282,7 @@ const GroundEmbedding: React.FC<MessageProps> = forwardRef((props, ref) => {
           .split('\n')
           .map((item: string) => {
             return {
-              text: item,
+              text: item?.trim(),
               uid: inputListRef.current?.setMessageId(),
               name: ''
             };
diff --git a/src/pages/playground/components/ground-images.tsx b/src/pages/playground/components/ground-images.tsx
index 0029ba1a..826f2811 100644
--- a/src/pages/playground/components/ground-images.tsx
+++ b/src/pages/playground/components/ground-images.tsx
@@ -4,7 +4,10 @@ import SealInput from '@/components/seal-form/seal-input';
 import SealSelect from '@/components/seal-form/seal-select';
 import useOverlayScroller from '@/hooks/use-overlay-scroller';
 import ThumbImg from '@/pages/playground/components/thumb-img';
-import { fetchChunkedData, readStreamData } from '@/utils/fetch-chunk-data';
+import {
+  fetchChunkedData,
+  readLargeStreamData as readStreamData
+} from '@/utils/fetch-chunk-data';
 import { FileImageOutlined, SwapOutlined } from '@ant-design/icons';
 import { useIntl, useSearchParams } from '@umijs/max';
 import { Button, Form, Tooltip } from 'antd';
@@ -138,18 +141,28 @@ const GroundImages: React.FC<MessageProps> = forwardRef((props, ref) => {
 
   const setImageSize = useCallback(() => {
     let size: Record<string, string | number> = {
+      with: 256,
+      height: 256,
       span: 12
     };
     if (parameters.n === 1) {
+      size.width = '100%';
+      size.height = '100%';
       size.span = 24;
     }
     if (parameters.n === 2) {
+      size.width = '50%';
+      size.height = 256;
       size.span = 12;
     }
     if (parameters.n === 3) {
+      size.width = '33%';
+      size.height = 256;
       size.span = 12;
     }
     if (parameters.n === 4) {
+      size.width = '25%';
+      size.height = 256;
       size.span = 12;
     }
     return size;
@@ -210,8 +223,7 @@ const GroundImages: React.FC<MessageProps> = forwardRef((props, ref) => {
       const params = {
         stream: true,
         stream_options: {
-          chunk_result: true,
-          chunk_size: 16 * 1024
+          chunk_result: true
         },
         prompt: current?.content || currentPrompt || '',
         ..._.omitBy(finalParameters, (value: string) => !value)
@@ -220,7 +232,6 @@ const GroundImages: React.FC<MessageProps> = forwardRef((props, ref) => {
       const result: any = await fetchChunkedData({
         data: params,
         url: CREAT_IMAGE_API,
-        // url: 'http://192.168.50.27:9090/v1/images/generations',
         signal: requestToken.current.signal
       });
 
@@ -235,12 +246,13 @@ const GroundImages: React.FC<MessageProps> = forwardRef((props, ref) => {
           });
           return;
         }
-
+        console.log('imgItem.dataUrl:', chunk.data);
         chunk?.data?.forEach((item: any) => {
           const imgItem = newImageList[item.index];
           if (item.b64_json) {
             imgItem.dataUrl += item.b64_json;
             // imgItem.cache.push(item.b64_json);
+            console.log('imgItem.dataUrl:', imgItem.dataUrl);
           }
           const progress = _.round(item.progress, 0);
           newImageList[item.index] = {
@@ -295,7 +307,7 @@ const GroundImages: React.FC<MessageProps> = forwardRef((props, ref) => {
       });
       setParams((pre: object) => {
         return {
-          ...pre,
+          ..._.omit(pre, ['quality', 'style']),
           seed: null,
           sampler: 'euler_a',
           cfg_scale: 1,
@@ -306,6 +318,8 @@ const GroundImages: React.FC<MessageProps> = forwardRef((props, ref) => {
     } else {
       setParams((pre: object) => {
         return {
+          quality: 'standard',
+          style: null,
           ..._.omit(pre, [
             'seed',
             'sampler',
@@ -320,6 +334,9 @@ const GroundImages: React.FC<MessageProps> = forwardRef((props, ref) => {
   };
 
   const renderExtra = useMemo(() => {
+    if (!isOpenaiCompatible) {
+      return [];
+    }
     return ImageconstExtraConfig.map((item: ParamsSchema) => {
       return (
         <Form.Item name={item.name} rules={item.rules} key={item.name}>
@@ -335,7 +352,7 @@ const GroundImages: React.FC<MessageProps> = forwardRef((props, ref) => {
         </Form.Item>
       );
     });
-  }, [ImageconstExtraConfig, intl]);
+  }, [ImageconstExtraConfig, isOpenaiCompatible, intl]);
 
   const renderAdvanced = useMemo(() => {
     if (isOpenaiCompatible) {
diff --git a/src/pages/playground/components/ground-reranker.tsx b/src/pages/playground/components/ground-reranker.tsx
index bb69ab3f..dd87308c 100644
--- a/src/pages/playground/components/ground-reranker.tsx
+++ b/src/pages/playground/components/ground-reranker.tsx
@@ -354,7 +354,7 @@ const GroundReranker: React.FC<MessageProps> = forwardRef((props, ref) => {
           .split('\n')
           .map((item: string) => {
             return {
-              text: item,
+              text: item?.trim(),
               uid: inputListRef.current?.setMessageId(),
               name: ''
             };
diff --git a/src/pages/playground/config/params-config.ts b/src/pages/playground/config/params-config.ts
index f0844fb6..6b0d9f9d 100644
--- a/src/pages/playground/config/params-config.ts
+++ b/src/pages/playground/config/params-config.ts
@@ -112,12 +112,12 @@ export const ImageParamsConfig: ParamsSchema[] = [
     type: 'Select',
     name: 'size',
     options: [
+      { label: 'playground.params.custom', value: 'custom', locale: true },
       { label: '256x256', value: '256x256' },
       { label: '512x512', value: '512x512' },
       { label: '1024x1024', value: '1024x1024' },
       { label: '1792x1024', value: '1792x1024' },
-      { label: '1024x1792', value: '1024x1792' },
-      { label: 'playground.params.custom', value: 'custom', locale: true }
+      { label: '1024x1792', value: '1024x1792' }
     ],
     label: {
       text: 'playground.params.size',
diff --git a/src/utils/fetch-chunk-data.ts b/src/utils/fetch-chunk-data.ts
index 008158cf..f4d15a2d 100644
--- a/src/utils/fetch-chunk-data.ts
+++ b/src/utils/fetch-chunk-data.ts
@@ -59,7 +59,9 @@ export const fetchChunkedData = async (params: {
     };
   }
   const reader = response?.body?.getReader();
-  const decoder = new TextDecoder('utf-8');
+  const decoder = new TextDecoder('utf-8', {
+    fatal: true
+  });
   return {
     reader,
     decoder
@@ -77,12 +79,54 @@ export const readStreamData = async (
   }
 
   let chunk = decoder.decode(value, { stream: true });
+
   extractJSON(chunk).forEach((data) => {
     callback?.(data);
   });
   await readStreamData(reader, decoder, callback);
 };
 
+export const readLargeStreamData = async (
+  reader: any,
+  decoder: TextDecoder,
+  callback: (data: any) => void
+) => {
+  let buffer = '';
+
+  const processStream = async () => {
+    const { done, value } = await reader.read();
+    if (done) {
+      if (buffer) {
+        try {
+          extractJSON(buffer).forEach((data) => {
+            callback?.(data);
+          });
+        } catch (e) {
+          console.error('parse buffer failed:', buffer);
+        }
+      }
+      return;
+    }
+
+    // cache each chunk
+    buffer += decoder.decode(value, { stream: true });
+
+    const extractedData = extractJSON(buffer);
+
+    extractedData.forEach((data) => {
+      callback?.(data);
+    });
+
+    const lastIndex = buffer.lastIndexOf('}');
+    buffer = lastIndex !== -1 ? buffer.slice(lastIndex + 1) : buffer;
+
+    // next chunk
+    await processStream();
+  };
+
+  await processStream();
+};
+
 export const readTextEventStreamData = async (
   reader: any,
   decoder: TextDecoder,