feat: generate images
This commit is contained in:
jialin
@@ -0,0 +1,310 @@
|
||||
import useOverlayScroller from '@/hooks/use-overlay-scroller';
|
||||
import useRequestToken from '@/hooks/use-request-token';
|
||||
import { useIntl, useSearchParams } from '@umijs/max';
|
||||
import { Spin } from 'antd';
|
||||
import classNames from 'classnames';
|
||||
import _ from 'lodash';
|
||||
import 'overlayscrollbars/overlayscrollbars.css';
|
||||
import {
|
||||
forwardRef,
|
||||
memo,
|
||||
useEffect,
|
||||
useImperativeHandle,
|
||||
useMemo,
|
||||
useRef,
|
||||
useState
|
||||
} from 'react';
|
||||
import { createImages } from '../apis';
|
||||
import { Roles, generateMessages } from '../config';
|
||||
import { ImageParamsConfig as paramsConfig } from '../config/params-config';
|
||||
import { MessageItem } from '../config/types';
|
||||
import '../style/ground-left.less';
|
||||
import '../style/system-message-wrap.less';
|
||||
import MessageInput from './message-input';
|
||||
import MessageContent from './multiple-chat/message-content';
|
||||
import ReferenceParams from './reference-params';
|
||||
import RerankerParams from './reranker-params';
|
||||
import ViewCodeModal from './view-code-modal';
|
||||
|
||||
interface MessageProps {
|
||||
modelList: Global.BaseOption<string>[];
|
||||
loaded?: boolean;
|
||||
ref?: any;
|
||||
}
|
||||
|
||||
const initialValues = {
|
||||
n: 1,
|
||||
size: '512x512',
|
||||
quality: 'standard',
|
||||
style: 'vivid'
|
||||
};
|
||||
|
||||
const GroundImages: React.FC<MessageProps> = forwardRef((props, ref) => {
|
||||
const { modelList } = props;
|
||||
const messageId = useRef<number>(0);
|
||||
const [messageList, setMessageList] = useState<MessageItem[]>([]);
|
||||
|
||||
const intl = useIntl();
|
||||
const requestSource = useRequestToken();
|
||||
const [searchParams] = useSearchParams();
|
||||
const selectModel = searchParams.get('model') || '';
|
||||
const [parameters, setParams] = useState<any>({});
|
||||
const [systemMessage, setSystemMessage] = useState('');
|
||||
const [show, setShow] = useState(false);
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [tokenResult, setTokenResult] = useState<any>(null);
|
||||
const [collapse, setCollapse] = useState(false);
|
||||
const contentRef = useRef<any>('');
|
||||
const scroller = useRef<any>(null);
|
||||
const currentMessageRef = useRef<any>(null);
|
||||
const paramsRef = useRef<any>(null);
|
||||
const messageListLengthCache = useRef<number>(0);
|
||||
const requestToken = useRef<any>(null);
|
||||
|
||||
const { initialize, updateScrollerPosition } = useOverlayScroller();
|
||||
const { initialize: innitializeParams } = useOverlayScroller();
|
||||
|
||||
useImperativeHandle(ref, () => {
|
||||
return {
|
||||
viewCode() {
|
||||
setShow(true);
|
||||
},
|
||||
setCollapse() {
|
||||
setCollapse(!collapse);
|
||||
},
|
||||
collapse: collapse
|
||||
};
|
||||
});
|
||||
|
||||
const viewCodeMessage = useMemo(() => {
|
||||
return generateMessages([
|
||||
{ role: Roles.System, content: systemMessage },
|
||||
...messageList
|
||||
]);
|
||||
}, [messageList, systemMessage]);
|
||||
|
||||
const setMessageId = () => {
|
||||
messageId.current = messageId.current + 1;
|
||||
return messageId.current;
|
||||
};
|
||||
|
||||
const handleNewMessage = (message?: { role: string; content: string }) => {
|
||||
const newMessage = message || {
|
||||
role:
|
||||
_.last(messageList)?.role === Roles.User ? Roles.Assistant : Roles.User,
|
||||
content: ''
|
||||
};
|
||||
messageList.push({
|
||||
...newMessage,
|
||||
uid: messageId.current + 1
|
||||
});
|
||||
setMessageId();
|
||||
setMessageList([...messageList]);
|
||||
};
|
||||
|
||||
const handleStopConversation = () => {
|
||||
requestToken.current?.cancel?.();
|
||||
setLoading(false);
|
||||
};
|
||||
|
||||
const submitMessage = async (current?: { role: string; content: string }) => {
|
||||
if (!parameters.model) return;
|
||||
try {
|
||||
setLoading(true);
|
||||
setMessageId();
|
||||
|
||||
requestToken.current?.cancel?.();
|
||||
requestToken.current = requestSource();
|
||||
|
||||
currentMessageRef.current = current
|
||||
? [
|
||||
{
|
||||
...current,
|
||||
uid: messageId.current
|
||||
}
|
||||
]
|
||||
: [];
|
||||
|
||||
contentRef.current = '';
|
||||
setMessageList((pre) => {
|
||||
return [...pre, ...currentMessageRef.current];
|
||||
});
|
||||
|
||||
const params = {
|
||||
prompt: current?.content || '',
|
||||
...parameters
|
||||
};
|
||||
|
||||
const result = await createImages(params, {
|
||||
cancelToken: requestToken.current.token
|
||||
});
|
||||
|
||||
const imgList = _.map(result.data, (item: any, index: number) => {
|
||||
return {
|
||||
dataUrl: `data:image/png;base64,${item.b64_json}`,
|
||||
created: result.created,
|
||||
uid: index
|
||||
};
|
||||
});
|
||||
setMessageList((pre) => {
|
||||
return [
|
||||
...pre,
|
||||
{
|
||||
content: '',
|
||||
role: Roles.Assistant,
|
||||
imgs: imgList,
|
||||
uid: messageId.current
|
||||
}
|
||||
];
|
||||
});
|
||||
console.log('result:', imgList);
|
||||
|
||||
setMessageId();
|
||||
} catch (error) {
|
||||
// console.log('error:', error);
|
||||
} finally {
|
||||
setLoading(false);
|
||||
}
|
||||
};
|
||||
const handleClear = () => {
|
||||
if (!messageList.length) {
|
||||
return;
|
||||
}
|
||||
setMessageId();
|
||||
setMessageList([]);
|
||||
setTokenResult(null);
|
||||
};
|
||||
|
||||
const handleSendMessage = (message: Omit<MessageItem, 'uid'>) => {
|
||||
console.log('message:', message);
|
||||
const currentMessage =
|
||||
message.content || message.imgs?.length ? message : undefined;
|
||||
submitMessage(currentMessage);
|
||||
};
|
||||
|
||||
const handleCloseViewCode = () => {
|
||||
setShow(false);
|
||||
};
|
||||
|
||||
const handleSelectModel = () => {};
|
||||
|
||||
const handlePresetPrompt = (list: { role: string; content: string }[]) => {
|
||||
const sysMsg = list.filter((item) => item.role === 'system');
|
||||
const userMsg = list
|
||||
.filter((item) => item.role === 'user')
|
||||
.map((item) => {
|
||||
setMessageId();
|
||||
return {
|
||||
...item,
|
||||
uid: messageId.current
|
||||
};
|
||||
});
|
||||
setSystemMessage(sysMsg[0]?.content || '');
|
||||
setMessageList(userMsg);
|
||||
};
|
||||
|
||||
useEffect(() => {
|
||||
if (scroller.current) {
|
||||
initialize(scroller.current);
|
||||
}
|
||||
}, [scroller.current, initialize]);
|
||||
|
||||
useEffect(() => {
|
||||
if (paramsRef.current) {
|
||||
innitializeParams(paramsRef.current);
|
||||
}
|
||||
}, [paramsRef.current, innitializeParams]);
|
||||
|
||||
useEffect(() => {
|
||||
if (loading) {
|
||||
updateScrollerPosition();
|
||||
}
|
||||
}, [messageList, loading]);
|
||||
|
||||
useEffect(() => {
|
||||
if (messageList.length > messageListLengthCache.current) {
|
||||
updateScrollerPosition();
|
||||
}
|
||||
messageListLengthCache.current = messageList.length;
|
||||
}, [messageList.length]);
|
||||
|
||||
return (
|
||||
<div className="ground-left-wrapper">
|
||||
<div className="ground-left">
|
||||
<div className="message-list-wrap" ref={scroller}>
|
||||
<>
|
||||
<div className="content">
|
||||
<MessageContent
|
||||
spans={{
|
||||
span: 24,
|
||||
count: 1
|
||||
}}
|
||||
actions={[]}
|
||||
messageList={messageList}
|
||||
setMessageList={setMessageList}
|
||||
editable={false}
|
||||
loading={loading}
|
||||
/>
|
||||
{loading && (
|
||||
<Spin size="small">
|
||||
<div style={{ height: '46px' }}></div>
|
||||
</Spin>
|
||||
)}
|
||||
</div>
|
||||
</>
|
||||
</div>
|
||||
{tokenResult && (
|
||||
<div style={{ height: 40 }}>
|
||||
<ReferenceParams usage={tokenResult}></ReferenceParams>
|
||||
</div>
|
||||
)}
|
||||
<div className="ground-left-footer">
|
||||
<MessageInput
|
||||
scope="chat"
|
||||
placeholer="Type <kbd>/</kbd> to input prompt"
|
||||
actions={['clear']}
|
||||
loading={loading}
|
||||
disabled={!parameters.model}
|
||||
isEmpty={!messageList.length}
|
||||
handleSubmit={handleSendMessage}
|
||||
addMessage={handleNewMessage}
|
||||
handleAbortFetch={handleStopConversation}
|
||||
clearAll={handleClear}
|
||||
setModelSelections={handleSelectModel}
|
||||
presetPrompt={handlePresetPrompt}
|
||||
modelList={modelList}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
<div
|
||||
className={classNames('params-wrapper', {
|
||||
collapsed: collapse
|
||||
})}
|
||||
ref={paramsRef}
|
||||
>
|
||||
<div className="box">
|
||||
<RerankerParams
|
||||
setParams={setParams}
|
||||
paramsConfig={paramsConfig}
|
||||
initialValues={initialValues}
|
||||
params={parameters}
|
||||
selectedModel={selectModel}
|
||||
modelList={modelList}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<ViewCodeModal
|
||||
open={show}
|
||||
payLoad={{
|
||||
prompt: currentMessageRef.current?.[0]?.content
|
||||
}}
|
||||
parameters={parameters}
|
||||
onCancel={handleCloseViewCode}
|
||||
title={intl.formatMessage({ id: 'playground.viewcode' })}
|
||||
></ViewCodeModal>
|
||||
</div>
|
||||
);
|
||||
});
|
||||
|
||||
export default memo(GroundImages);
|
||||
Reference in New Issue
Block a user