fix: enhance URL validation for Z.AI image generation (#5482) * fix: enhance URL validation for Z.AI image generation and handle potential delays Fixes #5480 * fix: retry 404 in a loop * fix: handle Z.AI image and video unavailability with appropriate warnings and responses

b8f2b1cfa66ae9dee462d329b8c114f747bc21f6

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed
Parents: b44b12b52734568cae63080f3d57be6a0f4fe28c
1 files changed, +23 -3Showing whitespace changes
src/endpoints/stable-diffusion.js+23 -3
@@ -1906,23 +1906,41 @@ zai.post('/generate', async (request, response) => {
1906 const data = await generateResponse.json();1906 const data = await generateResponse.json();
1907 console.debug('Z.AI image response:', data);1907 console.debug('Z.AI image response:', data);
19081908
1909 const url = data?.data?.[0]?.url;1909 const urlString = String(data?.data?.[0]?.url ?? '');
1910 if (!url || !isValidUrl(url) || !new URL(url).hostname.endsWith('.z.ai')) {1910 if (!urlString || !isValidUrl(urlString)) {
1911 console.warn('Z.AI returned an invalid image URL.');1911 console.warn('Z.AI returned an invalid image URL.');
1912 return response.sendStatus(500);1912 return response.sendStatus(500);
1913 }1913 }
19141914
1915 const url = new URL(urlString);
1916 if (!url.hostname.endsWith('.z.ai') && !url.hostname.endsWith('.ufileos.com')) {
1917 console.warn('Z.AI returned a URL with an unrecognized hostname.');
1918 return response.sendStatus(500);
1919 }
1920
1921 for (let attempt = 0; attempt < 5; attempt++) {
1915 const imageResponse = await fetch(url);1922 const imageResponse = await fetch(url);
1916 if (!imageResponse.ok) {1923 if (!imageResponse.ok) {
1924 // Sometimes the URL is valid but the image isn't immediately available
1925 if (imageResponse.status === 404) {
1926 console.info('Z.AI image not found yet, retrying...', { attempt: attempt + 1 });
1927 await delay(1000);
1928 continue;
1929 }
1930
1917 console.warn('Z.AI image fetch returned an error. Status:', imageResponse.status, imageResponse.statusText);1931 console.warn('Z.AI image fetch returned an error. Status:', imageResponse.status, imageResponse.statusText);
1918 return response.sendStatus(500);1932 return response.sendStatus(500);
1919 }1933 }
19201934
1921 const buffer = await imageResponse.arrayBuffer();1935 const buffer = await imageResponse.arrayBuffer();
1922 const image = Buffer.from(buffer).toString('base64');1936 const image = Buffer.from(buffer).toString('base64');
1923 const format = path.extname(url).substring(1).toLowerCase() || 'png';1937 const format = path.extname(url.pathname).substring(1).toLowerCase() || 'png';
19241938
1925 return response.send({ image, format });1939 return response.send({ image, format });
1940 }
1941
1942 console.warn('Z.AI image was not available after multiple attempts.');
1943 return response.sendStatus(500);
1926 } catch (error) {1944 } catch (error) {
1927 console.error(error);1945 console.error(error);
1928 return response.sendStatus(500);1946 return response.sendStatus(500);
@@ -2024,6 +2042,8 @@ zai.post('/generate-video', async (request, response) => {
2024 return response.send({ format: 'mp4', video: Buffer.from(contentBuffer).toString('base64') });2042 return response.send({ format: 'mp4', video: Buffer.from(contentBuffer).toString('base64') });
2025 }2043 }
2026 }2044 }
2045 console.warn('Z.AI video was not available after multiple attempts.');
2046 return response.sendStatus(500);
2027 } catch (error) {2047 } catch (error) {
2028 console.error(error);2048 console.error(error);
2029 return response.sendStatus(500);2049 return response.sendStatus(500);