Updated docker file

5cb75e6 over 1 year ago

5.09 kB

from fastapi import FastAPI, Request, Depends, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import os
from pathlib import Path

from seed_plans import seed_plans

from router.image import background_remove
from router.voice import clone
from router.try_on import try_clothes
from router.lip_sync import lip_sync
import torch
from transformers import AutoModelForImageSegmentation
from torchvision import transforms

from starlette.middleware.base import BaseHTTPMiddleware
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create FastAPI app
app = FastAPI(
    title="AI Image Processing API",
    description="API for image processing with AI models",
    version="1.0.0",
    debug=True,
)

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
import traceback


class CustomCORSMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if request.method == "OPTIONS":
            response = Response()
        else:
            try:
                response = await call_next(request)
            except Exception as e:
                response = Response(
                    content=f"Internal Server Error: {str(e)}", status_code=500
                )

        # Get the origin from the request
        origin = request.headers.get("origin")

        # Allow specific origins (not wildcard) when using credentials
        allowed_origins = [
            "http://localhost:3000",
            "http://localhost:3001",
            "http://localhost:5173",
            "http://localhost:8080",
        ]

        if origin in allowed_origins:
            response.headers["Access-Control-Allow-Origin"] = (
                origin  # Specific origin, not "*"
            )

        response.headers["Access-Control-Allow-Credentials"] = (
            "true"  # Must be string "true"
        )
        response.headers["Access-Control-Allow-Methods"] = "*"
        response.headers["Access-Control-Allow-Headers"] = "*"
        response.headers["Access-Control-Max-Age"] = "86400"

        return response


app.add_middleware(CustomCORSMiddleware)

# Add CORS middleware for Next.js frontend
app.add_middleware(
    CORSMiddleware,
    # allow_origins=["http://localhost:3000"],  # NextJS default port
    # allow_credentials=True,
    allow_origins=["*"],  # Allow all origins
    allow_credentials=True,  # Must be False when using "*"
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variables
birefnet = None
birefnet_lite = None
transform_image = None


# Include routers
# app.include_router(ai_image_router)
app.include_router(background_remove.router)
app.include_router(clone.router)
app.include_router(try_clothes.router)

# Create temp directory for file storage if it doesn't exist
temp_dir = Path("./temp")
temp_dir.mkdir(exist_ok=True)


@app.on_event("startup")
def startup_event():
    """Create database tables if they don't exist."""
    print("DB URL:", os.getenv("DATABASE_URL"))
    try_clothes.initialize_models()
    lip_sync.load_model("wav2lip_checkpoints/wav2lip_gan.pth")

    logger.info("Loading BiRefNet models...")
    torch.set_float32_matmul_precision("medium")
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Load both models and store in app state
    app.state.birefnet = AutoModelForImageSegmentation.from_pretrained(
        "ZhengPeng7/BiRefNet", trust_remote_code=True
    )
    app.state.birefnet.to(device)

    app.state.birefnet_lite = AutoModelForImageSegmentation.from_pretrained(
        "ZhengPeng7/BiRefNet_lite", trust_remote_code=True
    )
    app.state.birefnet_lite.to(device)

    app.state.transform_image = transforms.Compose(
        [
            transforms.Resize((768, 768)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )

    logger.info("Models loaded successfully")


@app.get("/")
async def root():
    return {
        "message": "Voice Processing API",
        "version": "1.0.0",
        "endpoints": {
            "voice_cloning": "/voice/clone",
            "speech_enhancement": "/voice/speech-enhancement",
            "speech_separation": "/voice/speech-separation",
            "target_speaker_extraction": "/voice/target-speaker-extraction",
            "speech_super_resolution": "/voice/speech-super-resolution",
            "health_check": "/voice/health",
            "cleanup": "/voice/cleanup",
            "download": "/voice/download/{filename}",
        },
    }


@app.get("/health")
async def health_check():
    return {"status": "healthy"}


if __name__ == "__main__":
    # Start server with hot reload during development
    # uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        reload_dirs=["./ai-api"],
        log_level="debug",  # Add this line for debug logs
    )

	from fastapi import FastAPI, Request, Depends, HTTPException, status
	from fastapi.middleware.cors import CORSMiddleware
	import uvicorn
	import os
	from pathlib import Path

	from seed_plans import seed_plans

	from router.image import background_remove
	from router.voice import clone
	from router.try_on import try_clothes
	from router.lip_sync import lip_sync
	import torch
	from transformers import AutoModelForImageSegmentation
	from torchvision import transforms

	from starlette.middleware.base import BaseHTTPMiddleware
	import logging

	# Configure logging
	logging.basicConfig(level=logging.INFO)
	logger = logging.getLogger(__name__)

	# Create FastAPI app
	app = FastAPI(
	title="AI Image Processing API",
	description="API for image processing with AI models",
	version="1.0.0",
	debug=True,
	)

	from starlette.middleware.base import BaseHTTPMiddleware
	from starlette.responses import Response
	import traceback


	class CustomCORSMiddleware(BaseHTTPMiddleware):
	async def dispatch(self, request: Request, call_next):
	if request.method == "OPTIONS":
	response = Response()
	else:
	try:
	response = await call_next(request)
	except Exception as e:
	response = Response(
	content=f"Internal Server Error: {str(e)}", status_code=500
	)

	# Get the origin from the request
	origin = request.headers.get("origin")

	# Allow specific origins (not wildcard) when using credentials
	allowed_origins = [
	"http://localhost:3000",
	"http://localhost:3001",
	"http://localhost:5173",
	"http://localhost:8080",
	]

	if origin in allowed_origins:
	response.headers["Access-Control-Allow-Origin"] = (
	origin # Specific origin, not "*"
	)

	response.headers["Access-Control-Allow-Credentials"] = (
	"true" # Must be string "true"
	)
	response.headers["Access-Control-Allow-Methods"] = "*"
	response.headers["Access-Control-Allow-Headers"] = "*"
	response.headers["Access-Control-Max-Age"] = "86400"

	return response


	app.add_middleware(CustomCORSMiddleware)

	# Add CORS middleware for Next.js frontend
	app.add_middleware(
	CORSMiddleware,
	# allow_origins=["http://localhost:3000"], # NextJS default port
	# allow_credentials=True,
	allow_origins=["*"], # Allow all origins
	allow_credentials=True, # Must be False when using "*"
	allow_methods=["*"],
	allow_headers=["*"],
	)

	# Global variables
	birefnet = None
	birefnet_lite = None
	transform_image = None


	# Include routers
	# app.include_router(ai_image_router)
	app.include_router(background_remove.router)
	app.include_router(clone.router)
	app.include_router(try_clothes.router)

	# Create temp directory for file storage if it doesn't exist
	temp_dir = Path("./temp")
	temp_dir.mkdir(exist_ok=True)


	@app.on_event("startup")
	def startup_event():
	"""Create database tables if they don't exist."""
	print("DB URL:", os.getenv("DATABASE_URL"))
	try_clothes.initialize_models()
	lip_sync.load_model("wav2lip_checkpoints/wav2lip_gan.pth")

	logger.info("Loading BiRefNet models...")
	torch.set_float32_matmul_precision("medium")
	device = "cuda" if torch.cuda.is_available() else "cpu"

	# Load both models and store in app state
	app.state.birefnet = AutoModelForImageSegmentation.from_pretrained(
	"ZhengPeng7/BiRefNet", trust_remote_code=True
	)
	app.state.birefnet.to(device)

	app.state.birefnet_lite = AutoModelForImageSegmentation.from_pretrained(
	"ZhengPeng7/BiRefNet_lite", trust_remote_code=True
	)
	app.state.birefnet_lite.to(device)

	app.state.transform_image = transforms.Compose(
	[
	transforms.Resize((768, 768)),
	transforms.ToTensor(),
	transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
	]
	)

	logger.info("Models loaded successfully")


	@app.get("/")
	async def root():
	return {
	"message": "Voice Processing API",
	"version": "1.0.0",
	"endpoints": {
	"voice_cloning": "/voice/clone",
	"speech_enhancement": "/voice/speech-enhancement",
	"speech_separation": "/voice/speech-separation",
	"target_speaker_extraction": "/voice/target-speaker-extraction",
	"speech_super_resolution": "/voice/speech-super-resolution",
	"health_check": "/voice/health",
	"cleanup": "/voice/cleanup",
	"download": "/voice/download/{filename}",
	},
	}


	@app.get("/health")
	async def health_check():
	return {"status": "healthy"}


	if __name__ == "__main__":
	# Start server with hot reload during development
	# uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
	uvicorn.run(
	"main:app",
	host="0.0.0.0",
	port=8000,
	reload=True,
	reload_dirs=["./ai-api"],
	log_level="debug", # Add this line for debug logs
	)