Update app.py

app.py

@@ -10,7 +10,6 @@ from app.logger import setup_logger
 import itertools
 import utils.extra_config
 import logging
-from flask import Flask, request, send_from_directory, render_template_string
 import threading
 
 if __name__ == "__main__":
@@ -57,35 +56,6 @@ if __name__ == "__main__":
 
 setup_logger(log_level=args.verbose, use_stdout=args.log_stdout)
 
-def start_file_browser():
-    app = Flask(__name__)
-    UPLOAD_FOLDER = "/data"
-
-    @app.route("/files")
-    def index():
-        files = os.listdir(UPLOAD_FOLDER)
-        file_list = ''.join(f'<li><a href="/files/download/{f}">{f}</a></li>' for f in files)
-        return render_template_string('''
-            <h2>File Browser (/data)</h2>
-            <ul>{{ file_list | safe }}</ul>
-            <h3>Upload File</h3>
-            <form method=post enctype=multipart/form-data action="/files/upload">
-                <input type=file name=file>
-                <input type=submit value=Upload>
-            </form>
-        ''', file_list=file_list)
-
-    @app.route("/files/download/<path:filename>")
-    def download_file(filename):
-        return send_from_directory(UPLOAD_FOLDER, filename, as_attachment=True)
-
-    @app.route("/files/upload", methods=["POST"])
-    def upload_file():
-        f = request.files["file"]
-        f.save(os.path.join(UPLOAD_FOLDER, f.filename))
-        return 'Uploaded! <a href="/files">Go back</a>'
-
-    app.run(host="0.0.0.0", port=7861, debug=False)
 
 
 def apply_custom_paths():
@@ -335,7 +305,86 @@ def start_comfyui(asyncio_loop=None):
 
     cuda_malloc_warning()
 
+    # Inject file browser routes into the FastAPI app
+    # In your app.py (inside start_comfyui)
+    from fastapi import UploadFile, Request, Form
+    from fastapi.responses import HTMLResponse, FileResponse, RedirectResponse
+    from fastapi.routing import APIRouter
+    import mimetypes
+    
+    UPLOAD_FOLDER = "/data"
+    router = APIRouter()
+    
+    # Helper to get safe paths
+    def get_safe_path(subpath: str):
+        full_path = os.path.abspath(os.path.join(UPLOAD_FOLDER, subpath))
+        if not full_path.startswith(os.path.abspath(UPLOAD_FOLDER)):
+            raise ValueError("Invalid path")
+        return full_path
+    
+    @router.get("/files/{path:path}", response_class=HTMLResponse)
+    async def browse_files(path: str):
+        full_path = get_safe_path(path)
+        if not os.path.exists(full_path):
+            return HTMLResponse("Not found", status_code=404)
+    
+        if os.path.isfile(full_path):
+            mime_type, _ = mimetypes.guess_type(full_path)
+            return FileResponse(full_path, media_type=mime_type)
+    
+        # Show folder listing
+        files = os.listdir(full_path)
+        files.sort()
+        parent_path = os.path.dirname(path)
+        list_html = f'<li><a href="/files/{parent_path}">.. (Up)</a></li>' if path else ''
+        for f in files:
+            sub = os.path.join(path, f)
+            encoded = sub.replace(" ", "%20")
+            entry_path = os.path.join(full_path, f)
+            preview = ''
+            if os.path.isfile(entry_path):
+                mime, _ = mimetypes.guess_type(entry_path)
+                if mime and mime.startswith("image"):
+                    preview = f'<br><img src="/files/{encoded}" style="max-height:100px;">'
+                elif mime and mime.startswith("video"):
+                    preview = f'<br><video src="/files/{encoded}" style="max-height:100px;" controls></video>'
+            list_html += f'<li><a href="/files/{encoded}">{f}</a> '
+            list_html += f'<a href="/files/delete/{encoded}" style="color:red">[Delete]</a>{preview}</li>'
+    
+        return f'''
+            <h2>Browsing: /{path}</h2>
+            <ul>{list_html}</ul>
+            <h3>Upload File</h3>
+            <form method=post enctype=multipart/form-data action="/files/upload/{path}">
+                <input type=file name=file>
+                <input type=submit value=Upload>
+            </form>
+        '''
+    
+    @router.post("/files/upload/{path:path}")
+    async def upload_file(path: str, file: UploadFile):
+        full_path = get_safe_path(path)
+        if not os.path.exists(full_path):
+            os.makedirs(full_path)
+        file_path = os.path.join(full_path, file.filename)
+        with open(file_path, "wb") as f:
+            f.write(await file.read())
+        return RedirectResponse(url=f"/files/{path}", status_code=303)
+    
+    @router.get("/files/delete/{path:path}")
+    async def delete_file(path: str):
+        try:
+            file_path = get_safe_path(path)
+            if os.path.isfile(file_path):
+                os.remove(file_path)
+                return RedirectResponse(url="/files/" + os.path.dirname(path), status_code=303)
+            return HTMLResponse("Not a file", status_code=400)
+        except Exception as e:
+            return HTMLResponse(f"Error: {str(e)}", status_code=400)
+
+
     prompt_server.add_routes()
+    prompt_server.app.include_router(router)
     hijack_progress(prompt_server)
 
     threading.Thread(target=prompt_worker, daemon=True, args=(q, prompt_server,)).start()
@@ -364,8 +413,6 @@ def start_comfyui(asyncio_loop=None):
 
 
 if __name__ == "__main__":
-    #start file browser
-    threading.Thread(target=start_file_browser, daemon=True).start()
     
     # Running directly, just start ComfyUI.
     logging.info("ComfyUI version: {}".format(comfyui_version.__version__))