Update public/index.html

public/index.html

@@ -89,23 +89,23 @@
                     </div>
                     <div class="mb-6">
                         <button id="generateBtn" onclick="generateImage()" class="w-full px-8 py-3 bg-gradient-to-r from-neon-blue to-neon-purple rounded-full font-semibold hover:opacity-90 transition-opacity animate-glow">
-                            <i class="fas fa-magic mr-2"></i>Generate
-                        </button>
-                    </div>
+                                    <i class="fas fa-magic mr-2"></i>Generate
+                                </button>
+                            </div>
                     <div id="loadingState" class="hidden text-center py-6">
                         <div class="inline-flex items-center space-x-4">
                             <div class="w-8 h-8 border-4 border-neon-blue border-t-transparent rounded-full animate-spin"></div>
                             <span class="text-neon-blue text-lg">Generating image...</span>
+                            </div>
+                            <div id="queueInfo" class="mt-4 text-gray-400"></div>
                         </div>
-                        <div id="queueInfo" class="mt-4 text-gray-400"></div>
-                    </div>
                     <div id="resultSection" class="hidden mt-8">
-                        <div class="glass-effect rounded-2xl p-6 neon-border">
-                            <img id="generatedImage" class="w-full rounded-xl shadow-2xl" alt="Generated image">
+                            <div class="glass-effect rounded-2xl p-6 neon-border">
+                                <img id="generatedImage" class="w-full rounded-xl shadow-2xl" alt="Generated image">
                             <div class="mt-6 flex justify-center">
-                                <button onclick="downloadImage()" class="px-6 py-3 glass-effect rounded-full hover:neon-border transition-all">
-                                    <i class="fas fa-download mr-2"></i>Download
-                                </button>
+                                    <button onclick="downloadImage()" class="px-6 py-3 glass-effect rounded-full hover:neon-border transition-all">
+                                        <i class="fas fa-download mr-2"></i>Download
+                                    </button>
                             </div>
                         </div>
                     </div>
@@ -156,13 +156,20 @@
             }
             hideResult();
             showLoading(true);
-            if (model === 'heartsync') {
-                // SSE logic
+            if (model === 'heartsync' || model === 'flux') {
+                // SSE logic for both models
+                if (model === 'flux') {
+                    const res = getFluxResolution();
+                    width = res.width;
+                    height = res.height;
+                }
                 try {
-                    const response = await fetch('/api/generate', {
+                    const url = model === 'heartsync' ? '/api/generate' : '/api/generate/flux';
+                    const body = model === 'heartsync' ? { prompt } : { prompt, width, height };
+                    const response = await fetch(url, {
                         method: 'POST',
                         headers: { 'Content-Type': 'application/json' },
-                        body: JSON.stringify({ prompt })
+                        body: JSON.stringify(body)
                     });
                     const reader = response.body.getReader();
                     const decoder = new TextDecoder();
@@ -181,7 +188,11 @@
                                         showLoading(true, 'Processing your image...');
                                     } else if (data.type === 'success') {
                                         showLoading(false);
-                                        showResult(data.originalUrl);
+                                        if (model === 'flux') {
+                                            showResult(data.imageUrl);
+                                        } else {
+                                            showResult(data.originalUrl);
+                                        }
                                         return;
                                     } else if (data.type === 'error') {
                                         throw new Error(data.message);
@@ -194,33 +205,6 @@
                     showLoading(false);
                     alert('Generation failed: ' + error.message);
                 }
-            } else if (model === 'flux') {
-                // FLUX.1-dev JSON logic
-                const res = getFluxResolution();
-                width = res.width;
-                height = res.height;
-                try {
-                    const response = await fetch('/api/generate/flux', {
-                        method: 'POST',
-                        headers: { 'Content-Type': 'application/json' },
-                        body: JSON.stringify({ prompt, width, height })
-                    });
-                    const data = await response.json();
-                    if (data.success && data.image && data.image.url) {
-                        showLoading(false);
-                        showResult(data.image.url);
-                    } else if (data.success && typeof data.image === 'string') {
-                        // fallback for string URL
-                        showLoading(false);
-                        showResult(data.image);
-                    } else {
-                        showLoading(false);
-                        alert('Generation failed: ' + (data.error || 'Unknown error'));
-                    }
-                } catch (error) {
-                    showLoading(false);
-                    alert('Generation failed: ' + error.message);
-                }
             }
         }
         function downloadImage() {