Update app.py

app.py

@@ -6,6 +6,8 @@ import torch
 from PIL import Image
 from tqdm import tqdm
 import gradio as gr
+import base64
+from flask import Flask, send_file, request, Response, jsonify
 
 from safetensors.torch import save_file
 from src.pipeline import FluxPipeline
@@ -105,6 +107,79 @@ with gr.Blocks() as demo:
         inputs=[prompt, spatial_img, height, width, seed, control_type],
         outputs=single_output_image
     )
+    
+    # Add proxy endpoints to access temporary files
+    with gr.Tab("API Endpoints"):
+        gr.Markdown("## Image Proxy Endpoints")
+        gr.Markdown("The following endpoints allow external access to generated images:")
+        gr.Markdown("- `/proxy-image?path=PATH` - Access a generated image directly")
+        gr.Markdown("- `/proxy-image-base64?path=PATH` - Get a generated image as base64")
+        gr.Markdown("These endpoints should only be used programmatically.")
+
+# Create Flask app for proxy endpoints
+# This allows Gradio to handle its routes while we add custom Flask routes
+app = gr.flask_app
+
+@app.route('/proxy-image')
+def proxy_image():
+    """
+    Proxies an image from a local file path
+    Example usage: /proxy-image?path=/tmp/gradio/4f461b7833afb048edcda4fd8968db08aeb352e871a5445721d63b11f9f489c7/image.webp
+    """
+    file_path = request.args.get('path')
+    
+    if not file_path:
+        return "No path specified", 400
+    
+    # Security check to prevent directory traversal attacks
+    if '..' in file_path:
+        return "Invalid path", 400
+    
+    # Ensure path starts with a slash
+    if not file_path.startswith('/'):
+        file_path = '/' + file_path
+    
+    try:
+        # Read the file
+        with open(file_path, 'rb') as f:
+            file_data = f.read()
+        
+        # Return the file
+        return Response(file_data, mimetype='image/webp')
+    except Exception as e:
+        return f"Error accessing file: {str(e)}", 404
+
+@app.route('/proxy-image-base64')
+def proxy_image_base64():
+    """
+    Proxies an image from a local file path and returns it as base64
+    Example usage: /proxy-image-base64?path=/tmp/gradio/4f461b7833afb048edcda4fd8968db08aeb352e871a5445721d63b11f9f489c7/image.webp
+    """
+    file_path = request.args.get('path')
+    
+    if not file_path:
+        return jsonify({"error": "No path specified"}), 400
+    
+    # Security check to prevent directory traversal attacks
+    if '..' in file_path:
+        return jsonify({"error": "Invalid path"}), 400
+    
+    # Ensure path starts with a slash
+    if not file_path.startswith('/'):
+        file_path = '/' + file_path
+    
+    try:
+        # Read the file
+        with open(file_path, 'rb') as f:
+            file_data = f.read()
+        
+        # Convert to base64
+        base64_data = base64.b64encode(file_data).decode('utf-8')
+        
+        # Return the base64 data
+        return jsonify({"base64": base64_data, "mime_type": "image/webp"})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 404
 
 # Launch the Gradio app
 demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)