Update app.py

app.py

@@ -3,10 +3,11 @@ import os
 import json
 import time
 import torch
-import base64
 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
@@ -52,54 +53,6 @@ def single_condition_generate_image(prompt, spatial_img, height, width, seed, co
     clear_cache(pipe.transformer)
     return image
 
-# Proxy endpoint function to serve images directly from file path
-def serve_file_from_path(file_path):
-    """
-    Serve a file directly from a path
-    """
-    if not file_path:
-        return None, "No path specified"
-    
-    # Security check to prevent directory traversal attacks
-    if '..' in file_path:
-        return None, "Invalid path"
-    
-    # Ensure path starts with a slash
-    if not file_path.startswith('/'):
-        file_path = '/' + file_path
-    
-    try:
-        # Open and return the image
-        return Image.open(file_path), None
-    except Exception as e:
-        return None, f"Error accessing file: {str(e)}"
-
-# Function to encode an image to base64
-def get_base64_from_path(file_path):
-    """
-    Get base64 encoded data from a file path
-    """
-    if not file_path:
-        return {"error": "No path specified"}
-    
-    # Security check to prevent directory traversal attacks
-    if '..' in file_path:
-        return {"error": "Invalid path"}
-    
-    # Ensure path starts with a slash
-    if not file_path.startswith('/'):
-        file_path = '/' + file_path
-    
-    try:
-        # Read the file and convert to base64
-        with open(file_path, 'rb') as f:
-            file_data = f.read()
-        
-        base64_data = base64.b64encode(file_data).decode('utf-8')
-        return {"base64": base64_data, "mime_type": "image/webp"}
-    except Exception as e:
-        return {"error": str(e)}
-
 # Define the Gradio interface components
 control_types = ["Ghibli"]
 
@@ -148,42 +101,85 @@ with gr.Blocks() as demo:
             label="Single Condition Examples"
         )
 
-    # Add proxy functionality as Gradio endpoints
-    with gr.Tab("Proxy Service"):
-        gr.Markdown("## File Proxy Service")
-        gr.Markdown("This tab provides a way to access files from local paths.")
-        
-        file_path_input = gr.Textbox(label="File Path", placeholder="Enter file path (e.g., /tmp/gradio/...)")
-        proxy_result = gr.Image(label="Retrieved Image", type="pil")
-        retrieve_btn = gr.Button("Retrieve Image")
-        
-        # Connect the button to the function
-        retrieve_btn.click(
-            serve_file_from_path,
-            inputs=[file_path_input],
-            outputs=[proxy_result, gr.Textbox(label="Error")]
-        )
-        
-        gr.Markdown("## Base64 Service")
-        gr.Markdown("Get a file as base64 encoded data.")
-        
-        base64_path_input = gr.Textbox(label="File Path", placeholder="Enter file path (e.g., /tmp/gradio/...)")
-        base64_result = gr.JSON(label="Base64 Result")
-        base64_btn = gr.Button("Get Base64")
-        
-        # Connect the button to the function
-        base64_btn.click(
-            get_base64_from_path,
-            inputs=[base64_path_input],
-            outputs=[base64_result]
-        )
-
     # Link the buttons to the functions
     single_generate_btn.click(
         single_condition_generate_image,
         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)