Update app.py

app.py

@@ -1,12 +1,13 @@
 import os
 import base64
-import hashlib
 import mimetypes
-from datetime import datetime
+from datetime import datetime, timedelta
 from io import BytesIO
+import requests
+import json
 
-from flask import Flask, request, jsonify, redirect, render_template
-from huggingface_hub import HfApi, HfFileSystem
+from flask import Flask, request, jsonify, redirect, render_template, send_from_directory
+from huggingface_hub import HfApi
 from PIL import Image
 import uuid
 
@@ -14,12 +15,10 @@ app = Flask(__name__)
 
 # 环境变量配置
 HF_TOKEN = os.environ.get("HF_TOKEN")
-DATASET_ID = os.environ.get("DATASET_ID")
+SPACE_ID = os.environ.get("SPACE_ID")
 MAX_SIZE = int(os.environ.get("MAX_SIZE", 10)) * 1024 * 1024  # 转字节
 API_KEY = os.environ.get("API_KEY")
-SPACE_DOMAIN = os.environ.get("SPACE_DOMAIN", "chb2026-image.hf.space")  # Space域名
 
-fs = HfFileSystem(token=HF_TOKEN)
 hf_api = HfApi(token=HF_TOKEN)
 
 def validate_auth():
@@ -29,22 +28,29 @@ def validate_auth():
     param_key = request.args.get("key")
     return header_key == API_KEY or param_key == API_KEY
 
-def generate_filename(data, mime_type=None):
-    ext = mimetypes.guess_extension(mime_type or "application/octet-stream") or ".bin"
+def generate_filename(original_filename):
     unique_id = uuid.uuid4().hex[:8]
     timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    _, ext = os.path.splitext(original_filename)
     return f"{timestamp}-{unique_id}{ext}"
 
-def save_to_dataset(content, filename):
-    fs.mkdir(f"datasets/{DATASET_ID}/images", exist_ok=True)
-    fs.write_bytes(f"datasets/{DATASET_ID}/images/{filename}", content)
-    # 返回代理URL
-    return f"https://{SPACE_DOMAIN}/proxy/datasets/{DATASET_ID}/resolve/main/images/{filename}"
+def save_to_space(content, filename):
+    path = f"images/{filename}"
+    hf_api.upload_file(
+        repo_id=SPACE_ID,
+        path_in_repo=path,
+        path_or_fileobj=content
+    )
+    return f"https://{SPACE_ID.split('/')[-1]}.hf.space/{path}"
 
 @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():
     # 鉴权检查
@@ -61,26 +67,31 @@ def upload():
         if "source" in request.files:
             file = request.files["source"]
             source = file.read()
+            original_filename = file.filename
         else:
             source = request.form.get("source")
+            original_filename = "uploaded_file"
     else:
         source = request.args.get("source")
+        original_filename = "uploaded_file"
 
     # 处理Base64、文件或URL
     if isinstance(source, bytes):
         content = source
     elif source and source.startswith(("http://", "https://")):
-        return jsonify({"error": "URL download not implemented"}), 501
+        response = requests.get(source)
+        content = response.content
+        original_filename = os.path.basename(source)
     elif source and source.startswith("data:"):
         header, data = source.split(",", 1)
         mime_type = header.split(":")[1].split(";")[0]
         content = base64.b64decode(data)
-    elif source:
-        content = base64.b64decode(source)
+        ext = mimetypes.guess_extension(mime_type)
+        original_filename = f"uploaded_file{ext}"
     else:
         return jsonify({
             "status_code": 400,
-            "error": {"message": "No file provided"},
+            "error": {"message": "Invalid source"},
             "status_txt": "Bad Request"
         }), 400
 
@@ -96,54 +107,65 @@ def upload():
     try:
         image = Image.open(BytesIO(content))
         width, height = image.size
-        mime_type = image.get_format_mimetype()
-        extension = mimetypes.guess_extension(mime_type) or ".jpg"
-        is_animated = getattr(image, "is_animated", False)
+        mime_type = Image.MIME[image.format]
     except:
         width = height = 0
-        mime_type = None
-        extension = ".bin"
-        is_animated = False
+        mime_type = mimetypes.guess_type(original_filename)[0] or "application/octet-stream"
 
     # 生成存储信息
-    filename = generate_filename(content, mime_type)
-    file_url = save_to_dataset(content, filename)
-
-    # 构造符合PicGo API的响应
-    response_format = request.args.get("format", "json")
-    
-    image_info = {
-        "name": os.path.splitext(filename)[0],
-        "extension": extension.lstrip("."),
-        "size": len(content),
-        "width": width,
-        "height": height,
-        "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-        "date_gmt": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S"),
-        "title": request.form.get("title", filename),
-        "description": request.form.get("description"),
-        "nsfw": int(request.form.get("nsfw", "0")),
-        "md5": hashlib.md5(content).hexdigest(),
-        "mime": mime_type,
-        "url": file_url,
-        "is_animated": int(is_animated),
-        "size_formatted": f"{len(content)/1024/1024:.1f} MB"
+    filename = generate_filename(original_filename)
+    file_url = save_to_space(BytesIO(content), filename)
+
+    # 构造响应
+    response_data = {
+        "status_code": 200,
+        "success": {"message": "file uploaded", "code": 200},
+        "image": {
+            "name": os.path.splitext(filename)[0],
+            "extension": os.path.splitext(filename)[1][1:],
+            "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,  # 可以实现MD5计算如果需要
+            "storage_mode": "space",
+            "source_md5": None,
+            "original_filename": original_filename,
+            "views": 0,
+            "likes": 0,
+            "is_animated": int(getattr(image, "is_animated", False)),
+            "is_360": 0,
+            "category_id": request.form.get("category_id"),
+            "id_encoded": uuid.uuid4().hex[:8],
+            "filename": filename,
+            "mime": mime_type,
+            "url": file_url,
+            "ratio": width / height if height else 0,
+            "size_formatted": f"{len(content)/1024/1024:.1f} MB",
+            "url_viewer": file_url,
+            "url_short": file_url,
+            "display_url": file_url,
+            "display_width": width,
+            "display_height": height,
+            "views_label": "views",
+            "likes_label": "likes",
+            "how_long_ago": "moments ago",
+            "delete_url": None  # 可以实现删除URL如果需要
+        },
+        "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({
-            "status_code": 200,
-            "success": {
-                "message": "file uploaded",
-                "code": 200
-            },
-            "image": image_info,
-            "status_txt": "OK"
-        })
+        return jsonify(response_data)
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860, debug=True)