Update app.py

app.py

@@ -222,13 +222,13 @@ def tiled_flux_img2img(pipe, prompt, image, strength, steps, guidance, generator
     return output
 
 def download_png(image):
-    """Convert image to PNG and return as downloadable file"""
+    """Convert image to PNG and return base64 string for download"""
     if image is None:
         raise gr.Error("No upscaled image available to download")
     buffer = io.BytesIO()
     image.save(buffer, format="PNG")
-    buffer.seek(0)
-    return buffer
+    base64_data = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    return base64_data
 
 @spaces.GPU(duration=120)
 def enhance_image(
@@ -405,6 +405,9 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
     # State to store the upscaled image
     upscaled_image_state = gr.State()
 
+    # Hidden textbox for base64 data
+    download_data = gr.Textbox(visible=False, elem_id="download_data")
+
     # Event handlers
     enhance_btn.click(
         fn=enhance_image,
@@ -424,7 +427,7 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
     download_btn.click(
         fn=download_png,
         inputs=[upscaled_image_state],
-        outputs=gr.File(label="Download Upscaled Image as PNG")
+        outputs=download_data
     )
     
     gr.HTML("""
@@ -494,6 +497,31 @@ with gr.Blocks(css=css, title="🎨 Flux dev Creative Upscaler - Florence-2 + FL
                 sliderInput.value = 50;
                 sliderInput.dispatchEvent(new Event('input'));
             }
+
+            const downloadData = document.querySelector('#download_data textarea');
+            if (downloadData) {
+                const observer = new MutationObserver(() => {
+                    const base64 = downloadData.value;
+                    if (base64) {
+                        const byteCharacters = atob(base64);
+                        const byteNumbers = new Array(byteCharacters.length);
+                        for (let i = 0; i < byteCharacters.length; i++) {
+                            byteNumbers[i] = byteCharacters.charCodeAt(i);
+                        }
+                        const byteArray = new Uint8Array(byteNumbers);
+                        const blob = new Blob([byteArray], {type: 'image/png'});
+                        const url = URL.createObjectURL(blob);
+                        const a = document.createElement('a');
+                        a.href = url;
+                        a.download = 'upscaled_image.png';
+                        a.click();
+                        URL.revokeObjectURL(url);
+                        // Clear the textbox
+                        downloadData.value = '';
+                    }
+                });
+                observer.observe(downloadData, {childList: true, subtree: true, characterData: true});
+            }
         });
     </script>
     """)