Update app.py

app.py

@@ -3,11 +3,10 @@ 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
@@ -53,6 +52,54 @@ 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"]
 
@@ -101,85 +148,42 @@ 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)