Update server.js

server.js

@@ -570,14 +570,13 @@ app.post('/api/generate/anime-uc', authenticateToken, async (req, res) => {
     }
 });
 
-// Flux Generation endpoint (now GET, returns only image URL)
+// Flux Generation endpoint (now uses UNO-FLUX model)
 app.get('/api/generate/flux', authenticateToken, async (req, res) => {
     try {
         const { prompt, aspect } = req.query;
         if (!prompt) {
             return res.status(400).json({ success: false, error: 'Prompt is required' });
         }
-        
         // Aspect ratio logic
         let width = 1280, height = 720;
         if (aspect === '9:16') {
@@ -586,106 +585,39 @@ app.get('/api/generate/flux', authenticateToken, async (req, res) => {
             width = 1024; height = 1024;
         }
 
-        const sessionHash = Math.random().toString(36).substring(2, 15);
-
-        console.log(`[Flux] Starting generation for prompt: "${prompt}" with session: ${sessionHash}`);
-
-        // Join the queue using the correct API structure
-        const joinResponse = await axios.post('https://yanze-pulid-flux.hf.space/queue/join', {
-            data: [
-                prompt,
-                null, // id_image (null for no ID image)
-                0,    // start_step
-                4,    // guidance
-                "-1", // seed
-                1,    // true_cfg
-                width,
-                height,
-                20,   // num_steps
-                1,    // id_weight
-                "bad quality, worst quality, text, signature, watermark, extra limbs", // neg_prompt
-                1,    // timestep_to_start_cfg
-                128   // max_sequence_length
-            ],
-            event_data: null,
-            fn_index: 4,
-            trigger_id: 21,
-            session_hash: sessionHash
-        }, {
-            headers: {
-                'Content-Type': 'application/json',
-                'Referer': 'https://yanze-pulid-flux.hf.space/'
-            }
-        });
-
-        console.log(`[Flux] Queue joined, event_id: ${joinResponse.data?.event_id}`);
-
-        // Poll for results
-        let attempts = 0;
-        const maxAttempts = 60;
-        
-        while (attempts < maxAttempts) {
-            try {
-                const dataResponse = await axios.get(
-                    `https://yanze-pulid-flux.hf.space/queue/data?session_hash=${encodeURIComponent(sessionHash)}`,
-                    {
-                        headers: {
-                            'Accept': 'text/event-stream',
-                            'Referer': 'https://yanze-pulid-flux.hf.space/'
-                        },
-                        timeout: 10000
-                    }
-                );
-
-                const data = dataResponse.data;
-                console.log(`[Flux] Poll attempt ${attempts + 1}, data length: ${data.length}`);
-                
-                // Parse the event stream data
-                const lines = data.split('\n');
-                let eventData = '';
-                
-                for (const line of lines) {
-                    if (line.startsWith('data: ')) {
-                        eventData += line.substring(6);
-                    } else if (line === '') {
-                        if (eventData) {
-                            try {
-                                const json = JSON.parse(eventData);
-                                console.log(`[Flux] Parsed event: ${json.msg}`);
-                                
-                                if (json.msg === 'process_completed' && json.output?.data?.[0]?.url) {
-                                    const imageUrl = json.output.data[0].url;
-                                    console.log(`[Flux] Found image:`, imageUrl);
-                                    
-                                    res.json({ success: true, imageUrl: imageUrl });
-                                    return;
-                                }
-
-                            } catch (error) {
-                                console.error('[Flux] Error parsing event data:', error, 'Raw data:', eventData);
-                            }
-                            eventData = '';
-                        }
-                    } else if (line.startsWith(':')) {
-                        continue;
-                    } else if (line !== '') {
-                         eventData += line;
-                    }
-                }
-
-            } catch (pollError) {
-                console.error(`[Flux] Poll error on attempt ${attempts + 1}:`, pollError.message);
-            }
-
-            await new Promise(resolve => setTimeout(resolve, 2000)); // Wait 2 seconds between polls
-            attempts++;
+        // Use a default image for all required image_prompt1-4
+        const defaultImageUrl = 'https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png';
+        async function fetchImageBlob(url) {
+            const response = await fetch(url);
+            if (!response.ok) throw new Error('Failed to fetch image: ' + url);
+            return await response.blob();
         }
+        const [img1, img2, img3, img4] = await Promise.all([
+            fetchImageBlob(defaultImageUrl),
+            fetchImageBlob(defaultImageUrl),
+            fetchImageBlob(defaultImageUrl),
+            fetchImageBlob(defaultImageUrl)
+        ]);
+
+        // Connect to UNO-FLUX Gradio client
+        const client = await Client.connect("bytedance-research/UNO-FLUX");
+        const result = await client.predict("/gradio_generate", {
+            prompt,
+            width,
+            height,
+            guidance: 4,
+            num_steps: 25,
+            seed: -1,
+            image_prompt1: img1,
+            image_prompt2: img2,
+            image_prompt3: img3,
+            image_prompt4: img4
+        });
 
-        console.log(`[Flux] Generation timed out after ${maxAttempts} attempts`);
-        res.status(500).json({ success: false, error: 'Generation timed out' });
-
+        // result.data[0] is the generated image URL
+        res.json({ success: true, imageUrl: result.data[0] });
     } catch (error) {
-        console.error('[Flux] Generation endpoint error:', error);
+        console.error('[Flux] UNO-FLUX endpoint error:', error);
         res.status(500).json({ success: false, error: error.message || 'Failed to generate image' });
     }
 });
@@ -922,14 +854,15 @@ app.post('/api/admin/process-old', async (req, res) => {
   }
 });
 
-// Proxy route for HuggingFace/Gradio images
+// Proxy route for HuggingFace/Gradio images (update to allow UNO-FLUX output URLs)
 app.get('/proxy-image', async (req, res) => {
     const { url, proxied } = req.query;
     if (!url) {
         return res.status(400).send('Missing url parameter');
     }
-    if (!url.startsWith('https://yanze-pulid-flux.hf.space/')) {
-        return res.status(403).send('Forbidden: Only allowed to proxy from yanze-pulid-flux.hf.space');
+    // Allow proxying from yanze-pulid-flux.hf.space and UNO-FLUX output URLs
+    if (!url.startsWith('https://yanze-pulid-flux.hf.space/') && !url.startsWith('https://hf.space/')) {
+        return res.status(403).send('Forbidden: Only allowed to proxy from yanze-pulid-flux.hf.space or hf.space');
     }
     try {
         const response = await fetch(url, {
@@ -960,6 +893,58 @@ app.get('/proxy-image', async (req, res) => {
     }
 });
 
+// UNO-FLUX model generation endpoint
+app.post('/api/generate/uno-flux', async (req, res) => {
+    try {
+        const {
+            prompt = "handsome woman in the city",
+            width = 512,
+            height = 512,
+            guidance = 4,
+            num_steps = 25,
+            seed = -1,
+            image_prompt1,
+            image_prompt2,
+            image_prompt3,
+            image_prompt4
+        } = req.body;
+
+        // Helper to fetch image as blob
+        async function fetchImageBlob(url) {
+            const response = await fetch(url);
+            if (!response.ok) throw new Error('Failed to fetch image: ' + url);
+            return await response.blob();
+        }
+
+        // Fetch all image prompts as blobs
+        const [img1, img2, img3, img4] = await Promise.all([
+            fetchImageBlob(image_prompt1),
+            fetchImageBlob(image_prompt2),
+            fetchImageBlob(image_prompt3),
+            fetchImageBlob(image_prompt4)
+        ]);
+
+        // Connect to UNO-FLUX Gradio client
+        const client = await Client.connect("bytedance-research/UNO-FLUX");
+        const result = await client.predict("/gradio_generate", {
+            prompt,
+            width,
+            height,
+            guidance,
+            num_steps,
+            seed,
+            image_prompt1: img1,
+            image_prompt2: img2,
+            image_prompt3: img3,
+            image_prompt4: img4
+        });
+
+        res.json({ success: true, data: result.data });
+    } catch (err) {
+        res.status(500).json({ success: false, error: err.message });
+    }
+});
+
 app.listen(port, () => {
   console.log(`Server running on port ${port}`);
 });