Download backend/server.py from Rizul2159/Elun: direct link, hf CLI and curl.
- Browser
- Download file 16.1 kB
-
https://huggingface.co/spaces/Rizul2159/Elun/resolve/main/backend/server.py
- Command line
-
hf download hf://spaces/Rizul2159/Elun/backend/server.py
-
curl -L -o server.py https://huggingface.co/spaces/Rizul2159/Elun/resolve/main/backend/server.py
16.1 kB
| 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 | |
| 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 --- | |
| def health_check(): | |
| return {"status": "healthy"} | |
| # --- Serve React Frontend (Static Files) --- | |
| # Serve index.html for root and SPA routing | |
| 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") | |