import AlertInfo from '@/components/alert-info'; import IconFont from '@/components/icon-font'; import FieldComponent from '@/components/seal-form/field-component'; 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, 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'; import classNames from 'classnames'; import _ from 'lodash'; import 'overlayscrollbars/overlayscrollbars.css'; import React, { forwardRef, memo, useCallback, useEffect, useImperativeHandle, useMemo, useRef, useState } from 'react'; import { CREAT_IMAGE_API } from '../apis'; import { OpenAIViewCode, promptList } from '../config'; import { ImageAdvancedParamsConfig, ImageconstExtraConfig, ImageParamsConfig as paramsConfig } from '../config/params-config'; import { MessageItem, ParamsSchema } from '../config/types'; import '../style/ground-left.less'; import '../style/system-message-wrap.less'; import DynamicParams from './dynamic-params'; import MessageInput from './message-input'; import ViewCodeModal from './view-code-modal'; interface MessageProps { modelList: Global.BaseOption[]; loaded?: boolean; ref?: any; } const initialValues = { n: 1, size: '512x512', seed: null, sampler: 'euler_a', cfg_scale: 4.5, sample_steps: 10, negative_prompt: null, schedule: 'discrete' }; const GroundImages: React.FC = forwardRef((props, ref) => { const { modelList } = props; const messageId = useRef(0); const [isOpenaiCompatible, setIsOpenaiCompatible] = useState(false); const [imageList, setImageList] = useState< { dataUrl: string; height: number | string; width: string | number; maxHeight: string | number; maxWidth: string | number; uid: number; span?: number; loading?: boolean; progress?: number; }[] >([ // { // dataUrl: // 'https://zos.alipayobjects.com/rmsportal/jkjgkEfvpUPVyRjUImniVslZfWPnJuuZ.png', // height: 'auto', // width: 'auto', // uid: 0, // span: 12, // loading: false, // progress: 60 // }, // { // dataUrl: // 'https://gw.alipayobjects.com/zos/antfincdn/LlvErxo8H9/photo-1503185912284-5271ff81b9a8.webp', // height: 'auto', // width: 'auto', // uid: 1, // span: 12, // progress: 15 // }, // { // dataUrl: // 'https://zos.alipayobjects.com/rmsportal/jkjgkEfvpUPVyRjUImniVslZfWPnJuuZ.png', // height: 'auto', // width: 'auto', // uid: 3, // span: 12, // progress: 10 // }, // { // dataUrl: // 'https://gw.alipayobjects.com/zos/antfincdn/LlvErxo8H9/photo-1503185912284-5271ff81b9a8.webp', // height: 'auto', // width: 'auto', // uid: 4, // span: 12, // progress: 15 // } ]); const intl = useIntl(); const [searchParams] = useSearchParams(); const selectModel = searchParams.get('model') || ''; const [parameters, setParams] = useState({}); const [show, setShow] = useState(false); const [loading, setLoading] = useState(false); const [tokenResult, setTokenResult] = useState(null); const [collapse, setCollapse] = useState(false); const scroller = useRef(null); const paramsRef = useRef(null); const messageListLengthCache = useRef(0); const requestToken = useRef(null); const [currentPrompt, setCurrentPrompt] = useState(''); const form = useRef(null); const inputRef = useRef(null); const size = Form.useWatch('size', form.current?.form); const { initialize, updateScrollerPosition } = useOverlayScroller(); const { initialize: innitializeParams } = useOverlayScroller(); useImperativeHandle(ref, () => { return { viewCode() { setShow(true); }, setCollapse() { setCollapse(!collapse); }, collapse: collapse }; }); const generateNumber = (min: number, max: number) => { return Math.floor(Math.random() * (max - min + 1) + min); }; const handleRandomPrompt = useCallback(() => { const randomIndex = generateNumber(0, promptList.length - 1); const randomPrompt = promptList[randomIndex]; inputRef.current?.handleInputChange({ target: { value: randomPrompt } }); }, []); const setImageSize = useCallback(() => { let size: Record = { 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; }, [parameters.n]); const finalParameters = useMemo(() => { if (parameters.size === 'custom') { return { ..._.omit(parameters, ['width', 'height']), size: parameters.width && parameters.height ? `${parameters.width}x${parameters.height}` : '' }; } return { ..._.omit(parameters, ['width', 'height']) }; }, [parameters]); const setMessageId = () => { messageId.current = messageId.current + 1; return messageId.current; }; const handleStopConversation = () => { requestToken.current?.abort?.(); setLoading(false); }; const submitMessage = async (current?: { content: string }) => { try { await form.current?.form?.validateFields(); if (!parameters.model) return; const size: any = setImageSize(); setLoading(true); setMessageId(); setCurrentPrompt(current?.content || ''); let newImageList = Array(parameters.n) .fill({}) .map((item, index: number) => { return { dataUrl: 'data:image/png;base64,', ...size, progress: 0, height: '100%', width: '100%', loading: true, uid: index }; }); setImageList(newImageList); requestToken.current?.abort?.(); requestToken.current = new AbortController(); const params = { stream: true, stream_options: {}, prompt: current?.content || currentPrompt || '', ..._.omitBy(finalParameters, (value: string) => !value) }; const result: any = await fetchChunkedData({ data: params, url: `${CREAT_IMAGE_API}?t=${Date.now()}`, signal: requestToken.current.signal }); if (result.error) { setTokenResult({ error: true, errorMessage: result?.data?.error?.message || result?.data?.error || '' }); setImageList([]); return; } const { reader, decoder } = result; const imgSize = _.split(finalParameters.size, 'x'); await readStreamData(reader, decoder, (chunk: any) => { if (chunk?.error) { setTokenResult({ error: true, errorMessage: chunk?.error?.message || chunk?.message || '' }); return; } chunk?.data?.forEach((item: any) => { const imgItem = newImageList[item.index]; if (item.b64_json) { imgItem.dataUrl += item.b64_json; } const progress = _.round(item.progress, 0); newImageList[item.index] = { dataUrl: imgItem.dataUrl, height: '100%', width: '100%', maxHeight: `${imgSize[1]}px`, maxWidth: `${imgSize[0]}px`, uid: imgItem.uid, span: imgItem.span, loading: progress < 100, progress: progress }; }); setImageList([...newImageList]); }); } catch (error) { console.log('error:', error); requestToken.current?.abort?.(); setImageList([]); } finally { setLoading(false); } }; const handleClear = () => { if (!imageList.length) { return; } setMessageId(); setImageList([]); setTokenResult(null); }; const handleSendMessage = (message: Omit) => { const currentMessage = message.content ? message : undefined; submitMessage(currentMessage); }; const handleCloseViewCode = () => { setShow(false); }; const handleToggleParamsStyle = () => { if (isOpenaiCompatible) { form.current?.form?.setFieldsValue({ seed: null, sampler: 'euler_a', cfg_scale: 4.5, sample_steps: 10, negative_prompt: null, schedule: 'discrete' }); setParams((pre: object) => { return { ..._.omit(pre, ['quality', 'style']), seed: null, sampler: 'euler_a', cfg_scale: 4.5, sample_steps: 10, negative_prompt: null, schedule: 'discrete' }; }); } else { setParams((pre: object) => { return { quality: 'standard', style: null, ..._.omit(pre, [ 'seed', 'sampler', 'cfg_scale', 'sample_steps', 'negative_prompt', 'schedule' ]) }; }); } setIsOpenaiCompatible(!isOpenaiCompatible); }; const renderExtra = useMemo(() => { if (!isOpenaiCompatible) { return []; } return ImageconstExtraConfig.map((item: ParamsSchema) => { return ( ); }); }, [ImageconstExtraConfig, isOpenaiCompatible, intl]); const renderAdvanced = useMemo(() => { if (isOpenaiCompatible) { return []; } return ImageAdvancedParamsConfig.map((item: ParamsSchema) => { return ( ); }); }, [ImageAdvancedParamsConfig, isOpenaiCompatible, intl]); const renderCustomSize = useMemo(() => { if (size === 'custom') { return (
); } return null; }, [size, intl]); useEffect(() => { if (scroller.current) { initialize(scroller.current); } }, [scroller.current, initialize]); useEffect(() => { if (paramsRef.current) { innitializeParams(paramsRef.current); } }, [paramsRef.current, innitializeParams]); useEffect(() => { if (loading) { updateScrollerPosition(); } }, [imageList, loading]); useEffect(() => { if (imageList.length > messageListLengthCache.current) { updateScrollerPosition(); } messageListLengthCache.current = imageList.length; }, [imageList.length]); return (
<>
{!imageList.length && (
{intl.formatMessage({ id: 'playground.params.empty.tips' })}
)}
{tokenResult && (
)}
{intl.formatMessage({ id: 'playground.image.prompt' })} } />
{intl.formatMessage({ id: 'playground.parameters' })}
} setParams={setParams} paramsConfig={paramsConfig} initialValues={initialValues} params={parameters} selectedModel={selectModel} modelList={modelList} extra={[renderCustomSize, ...renderExtra, ...renderAdvanced]} />
); }); export default memo(GroundImages);