t2c / main.py
99i's picture
Update main.py
3e55f8d verified
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import os
import time
import uuid
from generate_image import generate_image
from fastapi.staticfiles import StaticFiles
app = FastAPI()
# 挂载 tmp 文件夹下的静态文件
app.mount("/tmp", StaticFiles(directory="tmp"), name="static")
host=os.environ.get("HOST", "https://99i-t2c.hf.space")
# 添加跨域支持
origins = [
"*"
]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 创建一个模型,用于接收 markdown 字符串
class Markdown(BaseModel):
markdown: str
@app.get("/")
def hello():
return "hello world"
# 创建一个路由,用于接收 markdown 字符串并生成图片
@app.post("/t2c")
async def generate_image_endpoint(markdown: Markdown):
# 生成随机文件名
image_filename = 'tmp/'+str(uuid.uuid4()) + ".png"
# 调用 generate_image 函数来生成图片
generate_image(markdown.markdown, image_filename)
# 返回图片的路径
return {"url": f"{host}/{image_filename}"}
# 创建一个函数,用于清理过期的图片
def clean_expired_images():
expiration_time = int(os.environ.get("EXPIRATION_TIME", "3600")) # 默认过期时间为1小时
current_time = time.time()
for filename in os.listdir("tmp"):
file_path = os.path.join("tmp", filename)
if current_time - os.path.getmtime(file_path) > expiration_time:
os.remove(file_path)
# 创建一个路由,用于触发清理过期的图片
@app.get("/ci")
async def clean_images_endpoint():
# 调用 clean_expired_images 函数来清理过期的图片
clean_expired_images()
return {"message": "Images cleaned"}