- 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
85 lines
2.4 KiB
Python
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}")
|