Update app.py

app.py

@@ -15,13 +15,20 @@ app = Flask(__name__)
 
 # 环境变量配置
 HF_TOKEN = os.environ.get("HF_TOKEN")
-SPACE_ID = os.environ.get("SPACE_ID")
+SPACE_ID = os.environ.get("SPACE_ID", "username/spacename")  # 格式: username/spacename
 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"
+
 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")
@@ -29,30 +36,51 @@ def validate_auth():
     return header_key == API_KEY or param_key == API_KEY
 
 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)
+    if not ext:
+        ext = ".jpg"  # 默认扩展名
     return f"{timestamp}-{unique_id}{ext}"
 
 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}"
+    """保存文件到 Space"""
+    # 先保存到本地
+    local_path = os.path.join("images", filename)
+    with open(local_path, "wb") as f:
+        f.write(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)
 
 @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({
@@ -61,111 +89,135 @@ def upload():
             "status_txt": "Unauthorized"
         }), 401
 
-    # 获取文件数据
-    source = None
-    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")
-            original_filename = "uploaded_file"
-    else:
-        source = request.args.get("source")
+    try:
+        # 获取文件数据
+        source = None
         original_filename = "uploaded_file"
 
-    # 处理Base64、文件或URL
-    if isinstance(source, bytes):
-        content = source
-    elif source and source.startswith(("http://", "https://")):
-        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)
-        ext = mimetypes.guess_extension(mime_type)
-        original_filename = f"uploaded_file{ext}"
-    else:
-        return jsonify({
-            "status_code": 400,
-            "error": {"message": "Invalid source"},
-            "status_txt": "Bad Request"
-        }), 400
+        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
+        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 = {
+            "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,
+                "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,
+            },
+            "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)
 
-    # 验证大小限制
-    if len(content) > MAX_SIZE:
+    except Exception as e:
         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:
-        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(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(response_data)
+            "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)