backend/server.py

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

# Version 1.1 - High mode optimization in progressimport os
import sys
from pathlib import Path
import glob
import subprocess
import shutil
from PIL import Image

app = FastAPI()

# Allow all origins for Hugging Face deployment
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- Path Resolution for Cloud (Relative Paths) ---
BACKEND_DIR = Path(__file__).resolve().parent

# Key directories (all relative to backend folder)
TEST_IMG_DIR = BACKEND_DIR / "test_img"
ESRGAN_ROOT = BACKEND_DIR / "Real-ESRGAN"
ESRGAN_SCRIPT = ESRGAN_ROOT / "inference_realesrgan.py"
ESRGAN_INPUT_DIR = ESRGAN_ROOT / "inputs"
ESRGAN_OUTPUT_DIR = ESRGAN_ROOT / "results"
FACE_PARSING_DIR = BACKEND_DIR / "face-parsing.PyTorch"
FACE_PARSING_SCRIPT = FACE_PARSING_DIR / "test.py"
DIVIDED_REGIONS_DIR = BACKEND_DIR / "Divided Regions"

# Ensure directories exist on startup
TEST_IMG_DIR.mkdir(parents=True, exist_ok=True)
ESRGAN_INPUT_DIR.mkdir(parents=True, exist_ok=True)
ESRGAN_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
DIVIDED_REGIONS_DIR.mkdir(parents=True, exist_ok=True)


def run_esrgan_upscale(script_dir: Path, input_output_dir: Path, input_file_name: str, step: int):
    """
    Runs the Real-ESRGAN inference script with memory-optimized arguments.
    script_dir: Directory where inference_realesrgan.py is located
    input_output_dir: Directory containing inputs/ and results/ folders
    """
    input_path = input_output_dir / "inputs" / input_file_name
    output_path = input_output_dir / "results/"
    
    # Verify input file exists
    if not input_path.exists():
        raise HTTPException(status_code=500, detail=f"ESRGAN input file not found: {input_path}")
    
    # Optimization: Adaptive tile size based on image dimensions
    # Tile sizes MUST be even to avoid assertion errors
    try:
        img = Image.open(input_path)
        img_area = img.size[0] * img.size[1]
        
        if img_area > 200000:  # Large image
            tile_size = "128"  # Even tile size
            tile_pad = "2"
        elif img_area > 100000:  # Medium image
            tile_size = "256"  # Even tile size
            tile_pad = "2"
        else:  # Small image
            tile_size = "512"  # Even tile size
            tile_pad = "2"
        
        sys.stderr.write(f"DEBUG: ESRGAN adaptive settings - Image: {img.size}, Area: {img_area}, Tile: {tile_size}\n")
    except Exception as e:
        tile_size = "256"
        tile_pad = "2"
        sys.stderr.write(f"DEBUG: Using default tile settings: {e}\n")
    
    sys.stderr.write(f"DEBUG: ESRGAN Step {step} - Input: {input_path}\n")
    sys.stderr.write(f"DEBUG: ESRGAN Step {step} - Output dir: {output_path}\n")
    
    command = [
        "python", 
        "inference_realesrgan.py", 
        "-n", "RealESRGAN_x2plus",  # Faster 2x upscaling instead of 4x
        "-i", str(input_path), 
        "-o", str(output_path),
        "-t", tile_size,  # Adaptive tile size
        "--tile_pad", tile_pad,  # Minimal padding
        "--fp32"  # Force full precision on CPU
    ]
    
    sys.stderr.write(f"DEBUG: Running ESRGAN Step {step} with command: {' '.join(command)}\n")

    completed = None
    try:
        completed = subprocess.run(
            command,
            cwd=str(script_dir),
            check=True,
            capture_output=True,
            text=True,
            encoding='utf-8',
            errors='replace',
        )
    except subprocess.CalledProcessError as e:
        error_detail = e.stderr or e.stdout or str(e)
        sys.stderr.write(f"\n--- ESRGAN Step {step} FAILED ---\n{error_detail}\n")
        raise HTTPException(status_code=500, detail=f"ESRGAN upscale Step {step} failed: {error_detail}")

    if completed:
        sys.stderr.write(f"ESRGAN Step {step} stdout:\n{completed.stdout}\n")
        sys.stderr.write(f"ESRGAN Step {step} stderr:\n{completed.stderr}\n")
        sys.stderr.write(f"--- ESRGAN Step {step} SUCCESS ---\n")

    return completed


@app.post("/process")
async def process_image(
    file: UploadFile = File(...),
    quality: str = Form("high"),
    density: int = Form(50),
    palette: str = Form("math"),
):
    """
    Save upload, run optional ESRGAN upscaling, run face parsing script, return image.
    """
    quality = (quality or "high").strip().lower()
    if quality not in {"low", "medium", "high"}:
        raise HTTPException(status_code=400, detail="Invalid quality value. Use 'low', 'medium', or 'high'.")

    uploaded_file_name = "test.jpg"
    final_input_for_face_parsing = TEST_IMG_DIR / uploaded_file_name

    # --- Directory Setup and Cleanup ---
    try:
        TEST_IMG_DIR.mkdir(parents=True, exist_ok=True)
        ESRGAN_INPUT_DIR.mkdir(parents=True, exist_ok=True)
        ESRGAN_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        for f in ESRGAN_INPUT_DIR.glob("*.*"):
            os.remove(f)
        for f in ESRGAN_OUTPUT_DIR.glob("*.*"):
            os.remove(f)
        if final_input_for_face_parsing.exists():
            os.remove(final_input_for_face_parsing)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Failed to set up directories: {exc}")

    # --- Save Uploaded File ---
    esrgan_initial_input_path = ESRGAN_INPUT_DIR / uploaded_file_name
    try:
        data = await file.read()
        if not data:
            raise HTTPException(status_code=400, detail="Uploaded file is empty")
        
        with open(esrgan_initial_input_path, "wb") as out:
            out.write(data)
        
        # Verify file was saved
        if not esrgan_initial_input_path.exists():
            raise HTTPException(status_code=500, detail=f"File save failed: {esrgan_initial_input_path}")
        
        sys.stderr.write(f"DEBUG: Uploaded file saved: {esrgan_initial_input_path} ({len(data)} bytes)\n")
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Failed to save uploaded file: {exc}")

    # --- Optimization: Pre-process input image for faster upscaling ---
    try:
        img = Image.open(esrgan_initial_input_path)
        sys.stderr.write(f"DEBUG: Image opened successfully: {img.size}, mode: {img.mode}\n")
        
        # Convert to RGB first if needed
        if img.mode != "RGB":
            sys.stderr.write(f"DEBUG: Converting {img.mode} to RGB\n")
            if img.mode == "RGBA":
                img_rgb = Image.new("RGB", img.size, (255, 255, 255))
                img_rgb.paste(img, mask=img.split()[3])
                img = img_rgb
            else:
                img = img.convert("RGB")
        
        # Resize if too large (max 512 for better performance)
        max_size = 512
        if img.size[0] > max_size or img.size[1] > max_size:
            img.thumbnail((max_size, max_size), Image.LANCZOS)
            sys.stderr.write(f"DEBUG: Resized to {img.size}\n")
        
        # CRITICAL: Ensure dimensions are EVEN (divisible by 2) for ESRGAN
        width, height = img.size
        new_width = width - (width % 2)  # Make even
        new_height = height - (height % 2)  # Make even
        
        if new_width != width or new_height != height:
            img = img.crop((0, 0, new_width, new_height))
            sys.stderr.write(f"DEBUG: Cropped to even dimensions: {img.size}\n")
        
        # Verify dimensions are even
        assert img.size[0] % 2 == 0 and img.size[1] % 2 == 0, f"Dimensions still odd: {img.size}"
        
        # Save as JPEG
        img.save(esrgan_initial_input_path, format="JPEG", quality=85)
        sys.stderr.write(f"DEBUG: Preprocessed image saved with dimensions {img.size}\n")
    except Exception as e:
        sys.stderr.write(f"ERROR: Pre-processing failed: {e}\n")
        raise HTTPException(status_code=500, detail=f"Image preprocessing failed: {e}")

    # --- Handle ESRGAN Quality Options ---
    if quality == "low":
        shutil.copy(esrgan_initial_input_path, final_input_for_face_parsing)

    elif quality in {"medium", "high"}:
        input_file_name = uploaded_file_name
        
        # Clean outputs before ESRGAN step 1
        for f in ESRGAN_OUTPUT_DIR.glob("*.*"):
            try:
                os.remove(f)
            except Exception:
                pass
        
        # ESRGAN Step 1
        try:
            run_esrgan_upscale(ESRGAN_SCRIPT.parent, ESRGAN_ROOT, input_file_name, 1)
        except Exception as e:
            sys.stderr.write(f"ESRGAN Step 1 failed: {str(e)}\n")
            raise
        
        # Search for output files - try multiple patterns
        output_files_step1 = []
        for ext in ["*.png", "*.jpg", "*.jpeg"]:
            output_files_step1.extend(list(ESRGAN_OUTPUT_DIR.glob(ext)))
        if not output_files_step1:
            for ext in ["*.png", "*.jpg", "*.jpeg"]:
                output_files_step1.extend(list(ESRGAN_OUTPUT_DIR.glob(f"*/{ext}")))
        if not output_files_step1:
            for ext in ["*.png", "*.jpg", "*.jpeg"]:
                output_files_step1.extend(list(ESRGAN_OUTPUT_DIR.glob(f"**/{ext}")))
        
        if not output_files_step1:
            output_contents = list(ESRGAN_OUTPUT_DIR.rglob("*"))
            content_names = [str(p.relative_to(ESRGAN_OUTPUT_DIR)) for p in output_contents[:20]]
            sys.stderr.write(f"ESRGAN Step 1 output files not found. Directory: {ESRGAN_OUTPUT_DIR}\n")
            sys.stderr.write(f"Contents: {content_names}\n")
            raise HTTPException(status_code=500, 
                detail=f"ESRGAN Step 1 output not found. Directory: {ESRGAN_OUTPUT_DIR}. Contents: {content_names}")
        
        
        first_output_image = max(output_files_step1, key=lambda p: p.stat().st_mtime)
        sys.stderr.write(f"DEBUG: ESRGAN Step 1 output found: {first_output_image}\n")

        if quality == "high":
            # ESRGAN Step 2 - Load and fix dimensions FIRST
            step1_img = Image.open(first_output_image)
            width, height = step1_img.size
            sys.stderr.write(f"DEBUG: Step 1 output size: {step1_img.size}\n")
            
            # Ensure EVEN dimensions (divisible by 2)
            new_width = width - (width % 2)
            new_height = height - (height % 2)
            
            if new_width != width or new_height != height:
                step1_img = step1_img.crop((0, 0, new_width, new_height))
                sys.stderr.write(f"DEBUG: Step 2 input cropped to even dimensions: {step1_img.size}\n")
            
            # Verify dimensions are even
            assert step1_img.size[0] % 2 == 0 and step1_img.size[1] % 2 == 0, f"Step 2 input dimensions not even: {step1_img.size}"
            
            # Now clean up directories
            for f in ESRGAN_INPUT_DIR.glob("*.*"):
                try:
                    os.remove(f)
                except Exception:
                    pass
            
            for f in ESRGAN_OUTPUT_DIR.glob("*.*"):
                try:
                    os.remove(f)
                except Exception:
                    pass
            
            # Save the processed image for Step 2 as PNG to preserve exact dimensions
            step1_img.save(ESRGAN_INPUT_DIR / input_file_name, format="JPEG", quality=95)
            
            # Verify saved file has even dimensions
            verify_img = Image.open(ESRGAN_INPUT_DIR / input_file_name)
            if verify_img.size[0] % 2 != 0 or verify_img.size[1] % 2 != 0:
                sys.stderr.write(f"ERROR: Saved image has odd dimensions: {verify_img.size}\n")
                raise HTTPException(status_code=500, detail=f"Step 2 input has odd dimensions after save: {verify_img.size}")
            sys.stderr.write(f"DEBUG: Step 2 input saved and verified with size {verify_img.size}\n")
            
            try:
                run_esrgan_upscale(ESRGAN_SCRIPT.parent, ESRGAN_ROOT, input_file_name, 2)
            except Exception as e:
                sys.stderr.write(f"ESRGAN Step 2 failed: {str(e)}\n")
                raise

            # Search for output files from Step 2
            output_files_step2 = []
            for ext in ["*.png", "*.jpg", "*.jpeg"]:
                output_files_step2.extend(list(ESRGAN_OUTPUT_DIR.glob(ext)))
            if not output_files_step2:
                for ext in ["*.png", "*.jpg", "*.jpeg"]:
                    output_files_step2.extend(list(ESRGAN_OUTPUT_DIR.glob(f"*/{ext}")))
            if not output_files_step2:
                for ext in ["*.png", "*.jpg", "*.jpeg"]:
                    output_files_step2.extend(list(ESRGAN_OUTPUT_DIR.glob(f"**/{ext}")))
            
            if not output_files_step2:
                output_contents = list(ESRGAN_OUTPUT_DIR.rglob("*"))
                content_names = [str(p.relative_to(ESRGAN_OUTPUT_DIR)) for p in output_contents[:20]]
                sys.stderr.write(f"ESRGAN Step 2 output files not found. Contents: {content_names}\n")
                raise HTTPException(status_code=500, detail=f"ESRGAN Step 2 output not found. Contents: {content_names}")
            
            final_output_image = max(output_files_step2, key=lambda p: p.stat().st_mtime)
            sys.stderr.write(f"DEBUG: ESRGAN Step 2 output found: {final_output_image}\n")
        else:
            final_output_image = first_output_image

        sys.stderr.write(f"DEBUG: Copying output to parser: {final_output_image} -> {final_input_for_face_parsing}\n")
        shutil.copy(final_output_image, final_input_for_face_parsing)

    # --- Run Face Parsing Script ---
    if not FACE_PARSING_SCRIPT.exists():
        raise HTTPException(status_code=500, detail=f"Processing script not found at {FACE_PARSING_SCRIPT}")

    try:
        completed = subprocess.run(
            ["python", str(FACE_PARSING_SCRIPT), str(final_input_for_face_parsing), quality, palette],
            cwd=str(FACE_PARSING_DIR),
            capture_output=True,
            text=True,
            encoding='utf-8',
            errors='replace',
            check=False,
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Failed to execute processing script: {exc}")

    if completed.returncode != 0:
        raise HTTPException(status_code=500, detail=f"Script failed: {completed.stderr or completed.stdout}")

    # --- Locate and Return Output ---
    candidates = []
    for search_dir in [DIVIDED_REGIONS_DIR, FACE_PARSING_DIR, BACKEND_DIR]:
        for pattern in ["gift_worthy_mathematical_face*", "gift-worthy*"]:
            for ext in ("png", "jpg", "jpeg", "webp"):
                candidates.extend(
                    glob.glob(str(search_dir / f"**/{pattern}.{ext}"), recursive=True)
                )

    if not candidates:
        raise HTTPException(status_code=500, detail="Processed output not found.")

    output_path = max((Path(p) for p in candidates), key=lambda p: p.stat().st_mtime)

    # Cleanup input file
    try:
        if final_input_for_face_parsing.exists():
            os.remove(final_input_for_face_parsing)
    except Exception:
        pass

    return FileResponse(path=str(output_path), media_type="image/png")


# --- Health check endpoint for Hugging Face ---
@app.get("/health")
def health_check():
    return {"status": "healthy"}


# --- Serve React Frontend (Static Files) ---
# Serve index.html for root and SPA routing
@app.get("/", response_class=FileResponse)
async def serve_index():
    dist_path = BACKEND_DIR / "dist" / "index.html"
    if dist_path.exists():
        return dist_path
    return {"error": "Frontend not built"}

# Mount other static assets
dist_path = BACKEND_DIR / "dist"
if dist_path.exists() and (dist_path / "assets").exists():
    app.mount("/assets", StaticFiles(directory=str(dist_path / "assets")), name="assets")