Files
video-flow/modules/storage.py
Tony Zhang 33a165a615 feat: video-flow initial commit
- app.py: Streamlit UI for video generation workflow
- main_flow.py: CLI tool with argparse support
- modules/: Business logic modules (script_gen, image_gen, video_gen, composer, etc.)
- config.py: Configuration with API keys and paths
- requirements.txt: Python dependencies
- docs/: System prompt documentation
2025-12-12 19:18:27 +08:00

85 lines
2.4 KiB
Python

"""
MatchMe Studio - Storage Module (R2)
"""
import os
import logging
import time
import uuid
import boto3
from botocore.exceptions import NoCredentialsError
from pathlib import Path
from typing import Optional
import config
logger = logging.getLogger(__name__)
def get_s3_client():
try:
return boto3.client(
's3',
endpoint_url=config.R2_ENDPOINT,
aws_access_key_id=config.R2_ACCESS_KEY,
aws_secret_access_key=config.R2_SECRET_KEY,
region_name='auto'
)
except Exception as e:
logger.error(f"Failed to create R2 client: {e}")
raise
def upload_file(file_path: str) -> Optional[str]:
"""Upload file to R2 and return Public URL."""
if not os.path.exists(file_path):
logger.error(f"File not found: {file_path}")
return None
# 使用 UUID 作为文件名,避免中文/特殊字符导致的 URL 问题
original_name = Path(file_path).name
ext = Path(file_path).suffix.lower() or ".bin"
object_name = f"{uuid.uuid4().hex}{ext}"
s3 = get_s3_client()
try:
logger.info(f"Uploading {original_name} to R2 as {object_name}...")
# 根据后缀设置正确的 Content-Type
if ext == ".png":
content_type = "image/png"
elif ext in [".jpg", ".jpeg"]:
content_type = "image/jpeg"
elif ext == ".mp4":
content_type = "video/mp4"
elif ext == ".mp3":
content_type = "audio/mpeg"
else:
content_type = "application/octet-stream"
s3.upload_file(
file_path,
config.R2_BUCKET_NAME,
object_name,
ExtraArgs={'ContentType': content_type}
)
public_url = f"{config.R2_PUBLIC_URL}/{object_name}"
logger.info(f"Upload successful: {public_url}")
return public_url
except Exception as e:
logger.error(f"R2 Upload Failed: {e}")
return None
def cleanup_temp(max_age_seconds: int = 3600):
"""Delete old temp files."""
logger.info("Running cleanup_temp...")
now = time.time()
if not config.TEMP_DIR.exists(): return
for f in config.TEMP_DIR.iterdir():
try:
if f.is_file() and (now - f.stat().st_mtime) > max_age_seconds:
f.unlink()
except Exception as e:
logger.warning(f"Failed to delete {f}: {e}")