Update app.py

app.py

@@ -1,13 +1,13 @@
 import os
 import base64
+import hashlib
 import mimetypes
-from datetime import datetime, timedelta
+from datetime import datetime
 from io import BytesIO
 import requests
-import json
 
-from flask import Flask, request, jsonify, redirect, render_template, send_from_directory
-from huggingface_hub import HfApi
+from flask import Flask, request, jsonify, redirect, render_template
+from huggingface_hub import HfApi, HfFileSystem
 from PIL import Image
 import uuid
 
@@ -15,72 +15,46 @@ app = Flask(__name__)
 
 # 环境变量配置
 HF_TOKEN = os.environ.get("HF_TOKEN")
-SPACE_ID = os.environ.get("SPACE_ID", "username/spacename")  # 格式: username/spacename
+DATASET_ID = os.environ.get("DATASET_ID")
 MAX_SIZE = int(os.environ.get("MAX_SIZE", 10)) * 1024 * 1024  # 转字节
 API_KEY = os.environ.get("API_KEY")
 
-# 从 SPACE_ID 构建域名
-SPACE_DOMAIN = f"{SPACE_ID.replace('/', '-')}.hf.space"
-
+fs = HfFileSystem(token=HF_TOKEN)
 hf_api = HfApi(token=HF_TOKEN)
 
-# 确保上传目录存在
-os.makedirs("images", exist_ok=True)
-
 def validate_auth():
-    """验证 API 密钥"""
     if not API_KEY:
         return True
     header_key = request.headers.get("X-API-Key")
     param_key = request.args.get("key")
     return header_key == API_KEY or param_key == API_KEY
 
-def generate_filename(original_filename):
-    """生成唯一的文件名"""
+def generate_filename(data, mime_type=None):
+    ext = mimetypes.guess_extension(mime_type or "application/octet-stream") or ".bin"
     unique_id = uuid.uuid4().hex[:8]
     timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-    _, ext = os.path.splitext(original_filename)
-    if not ext:
-        ext = ".jpg"  # 默认扩展名
     return f"{timestamp}-{unique_id}{ext}"
 
-def save_to_space(content, filename):
-    """保存文件到 Space"""
-    # 先保存到本地
-    local_path = os.path.join("images", filename)
-    with open(local_path, "wb") as f:
-        f.write(content)
+def save_to_dataset(content, filename):
+    path = f"datasets/{DATASET_ID}/resolve/main/images/{filename}"
+    fs.mkdir(f"datasets/{DATASET_ID}/images", exist_ok=True)
+    fs.write_bytes(f"datasets/{DATASET_ID}/images/{filename}", content)
     
-    # 上传到 Space
-    try:
-        hf_api.upload_file(
-            repo_id=SPACE_ID,
-            path_in_repo=f"images/{filename}",
-            path_or_fileobj=local_path
-        )
-        # 返回可访问的 URL
-        return f"https://{SPACE_DOMAIN}/images/{filename}"
-    except Exception as e:
-        print(f"Upload error: {e}")
-        return None
-    finally:
-        # 清理本地文件
-        if os.path.exists(local_path):
-            os.remove(local_path)
+    # 获取正确的 URL
+    repo_id = f"datasets/{DATASET_ID}"
+    file_path = f"images/{filename}"
+    url = hf_api.hf_hub_url(repo_id, filename=file_path)
+    
+    # 获取实际的下载 URL
+    response = requests.head(url, allow_redirects=True)
+    return response.url
 
 @app.route("/", methods=["GET"])
 def index():
-    """渲染上传页面"""
     return render_template("index.html")
 
-@app.route("/images/<path:filename>")
-def serve_image(filename):
-    """提供图片访问服务"""
-    return send_from_directory("images", filename)
-
 @app.route("/api/1/upload", methods=["GET", "POST"])
 def upload():
-    """处理上传请求"""
     # 鉴权检查
     if not validate_auth():
         return jsonify({
@@ -89,135 +63,72 @@ def upload():
             "status_txt": "Unauthorized"
         }), 401
 
-    try:
-        # 获取文件数据
-        source = None
-        original_filename = "uploaded_file"
-
-        if request.method == "POST":
-            if "source" in request.files:
-                file = request.files["source"]
-                source = file.read()
-                original_filename = file.filename
-            else:
-                source = request.form.get("source")
-        else:
-            source = request.args.get("source")
-
-        if not source:
-            return jsonify({
-                "status_code": 400,
-                "error": {"message": "No source provided"},
-                "status_txt": "Bad Request"
-            }), 400
-
-        # 处理不同类型的输入
-        if isinstance(source, bytes):
-            content = source
-        elif source.startswith(("http://", "https://")):
-            response = requests.get(source)
-            content = response.content
-            original_filename = os.path.basename(source)
-        elif source.startswith("data:"):
-            try:
-                header, data = source.split(",", 1)
-                mime_type = header.split(":")[1].split(";")[0]
-                content = base64.b64decode(data)
-                ext = mimetypes.guess_extension(mime_type)
-                original_filename = f"uploaded_file{ext}"
-            except Exception as e:
-                return jsonify({
-                    "status_code": 400,
-                    "error": {"message": "Invalid base64 data"},
-                    "status_txt": "Bad Request"
-                }), 400
+    # 获取文件数据
+    source = None
+    if request.method == "POST":
+        if "source" in request.files:
+            file = request.files["source"]
+            source = file.read()
         else:
-            try:
-                content = base64.b64decode(source)
-            except Exception as e:
-                return jsonify({
-                    "status_code": 400,
-                    "error": {"message": "Invalid source format"},
-                    "status_txt": "Bad Request"
-                }), 400
-
-        # 验证文件大小
-        if len(content) > MAX_SIZE:
-            return jsonify({
-                "status_code": 413,
-                "error": {"message": f"File exceeds {MAX_SIZE//1024//1024}MB limit"},
-                "status_txt": "Payload Too Large"
-            }), 413
-
-        # 获取图片信息
-        try:
-            image = Image.open(BytesIO(content))
-            width, height = image.size
-            mime_type = Image.MIME[image.format]
-        except Exception as e:
-            width = height = 0
-            mime_type = mimetypes.guess_type(original_filename)[0] or "application/octet-stream"
-
-        # 生成文件名并保存
-        filename = generate_filename(original_filename)
-        file_url = save_to_space(content, filename)
-
-        if not file_url:
-            return jsonify({
-                "status_code": 500,
-                "error": {"message": "Failed to save file"},
-                "status_txt": "Internal Server Error"
-            }), 500
-
-        # 构造响应数据
-        response_data = {
+            source = request.form.get("source")
+    else:
+        source = request.args.get("source")
+
+    # 处理Base64、文件或URL
+    if isinstance(source, bytes):
+        content = source
+    elif source.startswith(("http://", "https://")):
+        # 需要实现URL下载逻辑（示例代码省略）
+        return jsonify({"error": "URL download not implemented"}), 501
+    elif source.startswith("data:"):
+        header, data = source.split(",", 1)
+        mime_type = header.split(":")[1].split(";")[0]
+        content = base64.b64decode(data)
+    else:
+        content = base64.b64decode(source)
+
+    # 验证大小限制
+    if len(content) > MAX_SIZE:
+        return jsonify({
+            "status_code": 413,
+            "error": {"message": f"File exceeds {MAX_SIZE//1024//1024}MB limit"},
+            "status_txt": "Payload Too Large"
+        }), 413
+
+    # 处理图片属性
+    try:
+        image = Image.open(BytesIO(content))
+        width, height = image.size
+        mime_type = image.get_format_mimetype()
+    except:
+        width = height = 0
+        mime_type = None
+
+    # 生成存储信息
+    filename = generate_filename(content, mime_type)
+    file_url = save_to_dataset(content, filename)
+
+    # 构造响应
+    response_format = request.args.get("format", "json")
+    if response_format == "txt":
+        return file_url
+    elif response_format == "redirect":
+        return redirect(file_url, code=302)
+    else:
+        return jsonify({
             "status_code": 200,
-            "success": {
-                "message": "file uploaded",
-                "code": 200
-            },
+            "success": {"message": "file uploaded", "code": 200},
             "image": {
-                "name": os.path.splitext(filename)[0],
-                "extension": os.path.splitext(filename)[1][1:],
+                "filename": filename,
+                "url": file_url,
                 "size": len(content),
                 "width": width,
                 "height": height,
-                "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-                "date_gmt": (datetime.now() + timedelta(hours=3)).strftime("%Y-%m-%d %H:%M:%S"),
-                "title": request.form.get("title") or os.path.splitext(original_filename)[0],
-                "description": request.form.get("description"),
-                "nsfw": int(request.form.get("nsfw", 0)),
-                "md5": None,
-                "storage_mode": "space",
-                "original_filename": original_filename,
-                "views": 0,
-                "likes": 0,
-                "is_animated": int(getattr(image, "is_animated", False)),
-                "filename": filename,
                 "mime": mime_type,
-                "url": file_url,
-                "size_formatted": f"{len(content)/1024/1024:.1f} MB",
-                "display_url": file_url,
-                "url_viewer": file_url,
+                "size_formatted": f"{len(content)/1024/1024:.1f} MB"
             },
             "status_txt": "OK"
-        }
-
-        # 根据请求的格式返回响应
-        response_format = request.args.get("format", "json")
-        if response_format == "txt":
-            return file_url
-        elif response_format == "redirect":
-            return redirect(file_url, code=302)
-        else:
-            return jsonify(response_data)
-
-    except Exception as e:
-        return jsonify({
-            "status_code": 500,
-            "error": {"message": str(e)},
-            "status_txt": "Internal Server Error"
-        }), 500
+        })
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860, debug=True)