Upload 2 files

app.py

@@ -0,0 +1,672 @@
+import hashlib
+import os
+from io import BytesIO
+import base64
+import requests
+from pathlib import Path
+import subprocess
+import shutil
+import gc
+import time
+import json
+import threading
+
+import gradio as gr
+from PIL import Image
+from cachetools import LRUCache
+import torch
+import numpy as np
+import torchvision.transforms.functional as F
+
+# FastAPI imports for enhanced frontend integration
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional
+import uvicorn
+
+# T4 Medium GPU Optimizations
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.max_split_size_mb = 512
+
+# API Models for FastAPI
+class ImageRequest(BaseModel):
+    source_image: str  # base64 string
+    shape_image: Optional[str] = None  # base64 string
+    color_image: Optional[str] = None  # base64 string
+    blending: str = "Article"
+    poisson_iters: int = 0
+    poisson_erosion: int = 15
+
+class ImageResponse(BaseModel):
+    success: bool
+    result_image: Optional[str] = None  # base64 string
+    message: str
+
+# Download working face landmarks
+def download_face_landmarks():
+    """Download working dlib face landmarks predictor"""
+    landmarks_path = 'pretrained_models/ShapeAdaptor/shape_predictor_68_face_landmarks.dat'
+    
+    if os.path.exists(landmarks_path) and os.path.getsize(landmarks_path) > 50000000:
+        print("Face landmarks already exists and appears valid")
+        return True
+    
+    try:
+        print("Downloading working face landmarks predictor...")
+        url = 'https://github.com/davisking/dlib-models/raw/master/shape_predictor_68_face_landmarks.dat.bz2'
+        
+        os.makedirs(os.path.dirname(landmarks_path), exist_ok=True)
+        
+        response = requests.get(url, stream=True, timeout=300)
+        response.raise_for_status()
+        
+        import bz2
+        compressed_data = response.content
+        
+        with open(landmarks_path, 'wb') as f:
+            f.write(bz2.decompress(compressed_data))
+        
+        print(f"Face landmarks downloaded successfully ({os.path.getsize(landmarks_path)/1024/1024:.1f}MB)")
+        return True
+        
+    except Exception as e:
+        print(f"Failed to download face landmarks: {e}")
+        return False
+
+# Comprehensive model download for full accuracy
+def download_all_missing_models():
+    """Download ALL required models for full accuracy"""
+    
+    all_required_models = {
+        'pretrained_models/ArcFace/backbone_ir50.pth': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/ArcFace/backbone_ir50.pth',
+        'pretrained_models/ArcFace/ir_se50.pth': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/ArcFace/ir_se50.pth',
+        'pretrained_models/BiSeNet/face_parsing_79999_iter.pth': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/BiSeNet/face_parsing_79999_iter.pth',
+        'pretrained_models/FeatureStyleEncoder/backbone.pth': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/FeatureStyleEncoder/backbone.pth',
+        'pretrained_models/FeatureStyleEncoder/psp_ffhq_encode.pt': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/FeatureStyleEncoder/psp_ffhq_encode.pt',
+        'pretrained_models/FeatureStyleEncoder/79999_iter.pth': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/FeatureStyleEncoder/79999_iter.pth',
+        'pretrained_models/FeatureStyleEncoder/143_enc.pth': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/FeatureStyleEncoder/143_enc.pth',
+        'pretrained_models/encoder4editing/e4e_ffhq_encode.pt': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/encoder4editing/e4e_ffhq_encode.pt',
+        'pretrained_models/sean_checkpoints/CelebA-HQ_pretrained/latest_net_G.pth': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/sean_checkpoints/CelebA-HQ_pretrained/latest_net_G.pth',
+        'pretrained_models/PostProcess/pp_model.pth': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/PostProcess/pp_model.pth',
+        'pretrained_models/PostProcess/latent_avg.pt': 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/PostProcess/latent_avg.pt',
+    }
+    
+    print("Checking for missing models for full accuracy...")
+    
+    missing_models = []
+    existing_models = []
+    
+    for model_path, url in all_required_models.items():
+        if os.path.exists(model_path):
+            existing_models.append(os.path.basename(model_path))
+        else:
+            missing_models.append((model_path, url))
+    
+    if existing_models:
+        print(f"Found existing models: {', '.join(existing_models)}")
+    
+    if missing_models:
+        print(f"Downloading {len(missing_models)} missing models for full accuracy...")
+        
+        for i, (model_path, url) in enumerate(missing_models, 1):
+            try:
+                model_name = os.path.basename(model_path)
+                print(f"[{i}/{len(missing_models)}] Downloading: {model_name}")
+                
+                os.makedirs(os.path.dirname(model_path), exist_ok=True)
+                
+                start_time = time.time()
+                response = requests.get(url, stream=True, timeout=600)
+                response.raise_for_status()
+                
+                total_size = int(response.headers.get('content-length', 0))
+                downloaded = 0
+                
+                with open(model_path, 'wb') as f:
+                    for chunk in response.iter_content(chunk_size=1024*1024):
+                        if chunk:
+                            f.write(chunk)
+                            downloaded += len(chunk)
+                            
+                            if total_size > 100*1024*1024 and downloaded % (50*1024*1024) == 0:
+                                progress = (downloaded / total_size) * 100 if total_size > 0 else 0
+                                print(f"   Progress: {progress:.1f}% ({downloaded/1024/1024:.1f}MB/{total_size/1024/1024:.1f}MB)")
+                
+                elapsed = time.time() - start_time
+                print(f"   Downloaded: {model_name} ({downloaded/1024/1024:.1f}MB in {elapsed:.1f}s)")
+                
+            except Exception as e:
+                print(f"   Failed to download {model_name}: {e}")
+                continue
+    else:
+        print("All models already present!")
+    
+    return True
+
+def download_mask_generator_separately():
+    """Download large mask_generator.pth file separately"""
+    mask_path = 'pretrained_models/ShapeAdaptor/mask_generator.pth'
+    
+    if os.path.exists(mask_path):
+        print("Mask generator already exists")
+        return True
+    
+    try:
+        print("Downloading mask_generator.pth (919MB)...")
+        url = 'https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/pretrained_models/ShapeAdaptor/mask_generator.pth'
+        
+        os.makedirs(os.path.dirname(mask_path), exist_ok=True)
+        
+        response = requests.get(url, stream=True, timeout=900)
+        response.raise_for_status()
+        
+        total_size = int(response.headers.get('content-length', 0))
+        downloaded = 0
+        
+        with open(mask_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=2*1024*1024):
+                if chunk:
+                    f.write(chunk)
+                    downloaded += len(chunk)
+                    
+                    if downloaded % (100*1024*1024) == 0:
+                        progress = (downloaded / total_size) * 100 if total_size > 0 else 0
+                        print(f"   Mask Generator Progress: {progress:.1f}% ({downloaded/1024/1024:.1f}MB/{total_size/1024/1024:.1f}MB)")
+        
+        print(f"Successfully downloaded mask_generator.pth ({downloaded/1024/1024:.1f}MB)")
+        return True
+        
+    except Exception as e:
+        print(f"Failed to download mask_generator.pth: {e}")
+        return False
+
+# Download all models
+download_all_missing_models()
+download_mask_generator_separately()
+download_face_landmarks()
+
+# Direct HairFast imports
+try:
+    from hair_swap import HairFast, get_parser
+    HAIRFAST_AVAILABLE = True
+    print("HairFast successfully imported!")
+except ImportError as e:
+    print(f"HairFast import failed: {e}")
+    HAIRFAST_AVAILABLE = False
+
+try:
+    from utils.shape_predictor import align_face
+    ALIGN_AVAILABLE = True
+    print("Face alignment available!")
+except ImportError as e:
+    print(f"Face alignment not available: {e}")
+    ALIGN_AVAILABLE = False
+
+# Global variables
+hair_fast_model = None
+align_cache = LRUCache(maxsize=10)
+
+def get_gpu_memory():
+    """Check GPU memory for optimization"""
+    if torch.cuda.is_available():
+        return torch.cuda.get_device_properties(0).total_memory / 1e9
+    return 0
+
+def optimize_for_t4():
+    """T4 GPU specific optimizations"""
+    if torch.cuda.is_available():
+        gpu_memory = get_gpu_memory()
+        print(f"GPU Memory: {gpu_memory:.1f}GB")
+        
+        if gpu_memory < 20:
+            torch.cuda.empty_cache()
+            os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
+            print("T4 optimizations applied")
+
+def check_model_completeness():
+    """Check which critical models are available"""
+    critical_models = {
+        'StyleGAN': 'pretrained_models/StyleGAN/ffhq.pt',
+        'Blending': 'pretrained_models/Blending/checkpoint.pth',
+        'Rotate': 'pretrained_models/Rotate/rotate_best.pth',
+        'PostProcess': 'pretrained_models/PostProcess/pp_model.pth',
+        'FeatureEncoder': 'pretrained_models/FeatureStyleEncoder/143_enc.pth',
+        'E4E': 'pretrained_models/encoder4editing/e4e_ffhq_encode.pt',
+        'SEAN': 'pretrained_models/sean_checkpoints/CelebA-HQ_pretrained/latest_net_G.pth',
+        'ShapeAdaptor': 'pretrained_models/ShapeAdaptor/mask_generator.pth',
+    }
+    
+    available_models = {}
+    for name, path in critical_models.items():
+        available_models[name] = os.path.exists(path)
+    
+    return available_models
+
+def initialize_hairfast_original():
+    """Initialize HairFast exactly like original - use pure defaults"""
+    global hair_fast_model
+    
+    if not HAIRFAST_AVAILABLE:
+        print("HairFast not available")
+        return False
+        
+    try:
+        print("Initializing HairFast with original default arguments...")
+        
+        optimize_for_t4()
+        
+        available_models = check_model_completeness()
+        print("Available models:", {k: v for k, v in available_models.items() if v})
+        
+        parser = get_parser()
+        args = parser.parse_args([])
+        
+        hair_fast_model = HairFast(args)
+        
+        available_count = sum(available_models.values())
+        total_count = len(available_models)
+        accuracy_percentage = (available_count / total_count) * 100
+        
+        print(f"HairFast initialized with original defaults ({accuracy_percentage:.1f}% accuracy)")
+        
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        return True
+        
+    except Exception as e:
+        print(f"HairFast initialization failed: {e}")
+        hair_fast_model = None
+        torch.cuda.empty_cache()
+        return False
+
+def get_bytes(img):
+    """EXACT copy of original get_bytes function"""
+    if img is None:
+        return img
+    buffered = BytesIO()
+    img.save(buffered, format="JPEG")
+    return buffered.getvalue()
+
+def bytes_to_image(image: bytes) -> Image.Image:
+    """EXACT copy of original bytes_to_image function"""
+    image = Image.open(BytesIO(image))
+    return image
+
+def base64_to_image(base64_string):
+    """Convert base64 string to PIL Image with error handling"""
+    try:
+        if base64_string.startswith('data:image'):
+            base64_string = base64_string.split(',')[1]
+        image_bytes = base64.b64decode(base64_string)
+        image = Image.open(BytesIO(image_bytes))
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        return image
+    except Exception as e:
+        print(f"Error converting base64 to image: {e}")
+        return None
+
+def image_to_base64(image):
+    """Convert PIL Image to base64 string with maximum quality"""
+    if image is None:
+        return None
+    buffered = BytesIO()
+    # Use PNG for lossless quality
+    image.save(buffered, format="PNG", optimize=False)
+    img_bytes = buffered.getvalue()
+    img_base64 = base64.b64encode(img_bytes).decode('utf-8')
+    return f"data:image/png;base64,{img_base64}"
+
+def center_crop(img):
+    """EXACT copy of original center_crop function"""
+    width, height = img.size
+    side = min(width, height)
+    left = (width - side) / 2
+    top = (height - side) / 2
+    right = (width + side) / 2
+    bottom = (height + side) / 2
+    img = img.crop((left, top, right, bottom))
+    return img
+
+def resize(name):
+    """Fixed resize function with proper size handling"""
+    def resize_inner(img, align):
+        global align_cache
+        if name in align and ALIGN_AVAILABLE:
+            img_hash = hashlib.md5(get_bytes(img)).hexdigest()
+            if img_hash not in align_cache:
+                try:
+                    aligned_imgs = align_face(img, return_tensors=False)
+                    if aligned_imgs and len(aligned_imgs) > 0:
+                        img = aligned_imgs[0]
+                        if img.size != (1024, 1024):
+                            img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
+                        align_cache[img_hash] = img
+                    else:
+                        img = center_crop(img)
+                        img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
+                        align_cache[img_hash] = img
+                except Exception as e:
+                    print(f"Face alignment failed for {name}, using center crop: {e}")
+                    img = center_crop(img)
+                    img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
+                    align_cache[img_hash] = img
+            else:
+                img = align_cache[img_hash]
+        else:
+            if img.size != (1024, 1024):
+                img = center_crop(img)
+                img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
+        return img
+    return resize_inner
+
+def swap_hair_selective(face, shape, color, blending, poisson_iters, poisson_erosion):
+    """
+    Enhanced swap logic with selective transfer:
+    - If only shape provided: change hairstyle only
+    - If only color provided: change hair color only  
+    - If both provided: change both hairstyle and color
+    """
+    global hair_fast_model
+    
+    if hair_fast_model is None:
+        if not initialize_hairfast_original():
+            return None, "HairFast model not available. Please check if all model files are uploaded."
+    
+    if not face and not shape and not color:
+        return None, "Need to upload a face and at least a shape or color ❗"
+    elif not face:
+        return None, "Need to upload a face ❗"
+    elif not shape and not color:
+        return None, "Need to upload at least a shape or color ❗"
+
+    try:
+        print("Starting selective hair transfer...")
+        
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        def validate_size(img, name):
+            if img is not None:
+                if img.size != (1024, 1024):
+                    print(f"Resizing {name} from {img.size} to (1024, 1024)")
+                    img = center_crop(img)
+                    img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
+            return img
+        
+        face = validate_size(face, "face")
+        shape = validate_size(shape, "shape") if shape else None
+        color = validate_size(color, "color") if color else None
+        
+        # Determine transfer mode
+        has_shape = shape is not None
+        has_color = color is not None
+        
+        if has_shape and has_color:
+            transfer_mode = "both"
+            print("Transfer mode: Both hairstyle and color")
+        elif has_shape and not has_color:
+            transfer_mode = "shape_only"
+            color = face  # Use original face for color reference
+            print("Transfer mode: Hairstyle only (preserving original color)")
+        elif has_color and not has_shape:
+            transfer_mode = "color_only"
+            shape = face  # Use original face for shape reference
+            print("Transfer mode: Color only (preserving original hairstyle)")
+        
+        print(f"Final sizes - Face: {face.size}, Shape: {shape.size if shape else 'None'}, Color: {color.size if color else 'None'}")
+        
+        with torch.no_grad():
+            start_time = time.time()
+            
+            # Use the HairFast model's swap method with proper parameters
+            final_image = hair_fast_model.swap(face, shape, color)
+            
+            inference_time = time.time() - start_time
+            print(f"Inference completed in {inference_time:.2f} seconds")
+        
+        result_image = F.to_pil_image(final_image)
+        
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        success_message = f"Hair transfer ({transfer_mode}) completed successfully in {inference_time:.2f}s"
+        print(success_message)
+        
+        return result_image, success_message
+
+    except Exception as e:
+        torch.cuda.empty_cache()
+        gc.collect()
+        error_msg = f"Hair transfer failed: {str(e)}"
+        print(f"Detailed error: {e}")
+        return None, error_msg
+
+def hair_transfer_api(source_image, shape_image=None, color_image=None, 
+                     blending="Article", poisson_iters=0, poisson_erosion=15):
+    """Enhanced API function for frontend integration with quality preservation"""
+    try:
+        print("API call received - processing images...")
+        
+        if isinstance(source_image, str):
+            source_image = base64_to_image(source_image)
+            print(f"Source image loaded: {source_image.size if source_image else 'Failed'}")
+        
+        if isinstance(shape_image, str) and shape_image:
+            shape_image = base64_to_image(shape_image)
+            print(f"Shape image loaded: {shape_image.size if shape_image else 'Failed'}")
+        
+        if isinstance(color_image, str) and color_image:
+            color_image = base64_to_image(color_image)
+            print(f"Color image loaded: {color_image.size if color_image else 'Failed'}")
+        
+        if not source_image:
+            return None, "Failed to process source image"
+        
+        result_image, status_message = swap_hair_selective(
+            source_image, shape_image, color_image, 
+            blending, poisson_iters, poisson_erosion
+        )
+        
+        if result_image is not None:
+            print(f"Result image generated: {result_image.size}")
+            result_base64 = image_to_base64(result_image)
+            print("Result converted to base64 successfully")
+            return result_base64, status_message
+        else:
+            print("Result image is None")
+            return None, status_message
+        
+    except Exception as e:
+        error_msg = f"API Error: {str(e)}"
+        print(f"API Error details: {e}")
+        return None, error_msg
+
+def get_demo():
+    """Enhanced Gradio interface with selective transfer options"""
+    with gr.Blocks(
+        title="HairFastGAN - Selective Transfer", 
+        theme=gr.themes.Soft(),
+        css="""
+        .error-message {
+            color: red !important;
+            background-color: #ffebee !important;
+            padding: 10px !important;
+            border-radius: 5px !important;
+        }
+        .transfer-info {
+            background-color: #e3f2fd !important;
+            padding: 10px !important;
+            border-radius: 5px !important;
+            margin: 10px 0 !important;
+        }
+        """
+    ) as demo:
+        
+        gr.Markdown("## HairFastGan - Selective Transfer")
+        gr.Markdown(
+            '<div style="display: flex; align-items: center; gap: 10px;">'
+            '<span>Enhanced HairFastGAN with selective transfer:</span>'
+            '<a href="https://arxiv.org/abs/2404.01094"><img src="https://img.shields.io/badge/arXiv-2404.01094-b31b1b.svg" height=22.5></a>'
+            '<a href="https://github.com/AIRI-Institute/HairFastGAN"><img src="https://img.shields.io/badge/github-%23121011.svg?style=for-the-badge&logo=github&logoColor=white" height=22.5></a>'
+            '<a href="https://huggingface.co/AIRI-Institute/HairFastGAN"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/model-on-hf-md.svg" height=22.5></a>'
+            '<a href="https://colab.research.google.com/#fileId=https://huggingface.co/AIRI-Institute/HairFastGAN/blob/main/notebooks/HairFast_inference.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>'
+            '</div>'
+        )
+        
+        gr.Markdown(
+            """
+            <div class="transfer-info">
+            <b>🎯 Selective Transfer Modes:</b><br>
+            • <b>Shape Only:</b> Upload only shape image → Changes hairstyle while preserving original color<br>
+            • <b>Color Only:</b> Upload only color image → Changes hair color while preserving original hairstyle<br>
+            • <b>Both:</b> Upload both images → Changes both hairstyle and color
+            </div>
+            """, 
+            elem_classes="transfer-info"
+        )
+        
+        with gr.Row():
+            with gr.Column():
+                source = gr.Image(label="Source photo to try on the hairstyle", type="pil")
+                
+                with gr.Row():
+                    shape = gr.Image(label="Shape photo with desired hairstyle (optional)", type="pil")
+                    color = gr.Image(label="Color photo with desired hair color (optional)", type="pil")
+                
+                # Transfer mode indicator
+                transfer_status = gr.Textbox(label="Transfer Mode", interactive=False, 
+                                           value="Upload images to see transfer mode")
+                
+                with gr.Accordion("Advanced Options", open=False):
+                    blending = gr.Radio(["Article", "Alternative_v1", "Alternative_v2"], value='Article',
+                                        label="Color Encoder version")
+                    poisson_iters = gr.Slider(0, 2500, value=0, step=1, label="Poisson iters")
+                    poisson_erosion = gr.Slider(1, 100, value=15, step=1, label="Poisson erosion")
+                    align = gr.CheckboxGroup(["Face", "Shape", "Color"], value=["Face", "Shape", "Color"],
+                                            label="Image cropping [recommended]")
+                
+                btn = gr.Button("Get the haircut", variant="primary")
+                
+            with gr.Column():
+                output = gr.Image(label="Your result")
+                error_message = gr.Textbox(label="⚠️ Error ⚠️", visible=False, elem_classes="error-message")
+
+        # Function to update transfer mode display
+        def update_transfer_mode(shape_img, color_img):
+            if shape_img is not None and color_img is not None:
+                return "🎯 Both hairstyle and color will be transferred"
+            elif shape_img is not None and color_img is None:
+                return "🎨 Only hairstyle will be transferred (color preserved)"
+            elif shape_img is None and color_img is not None:
+                return "🌈 Only hair color will be transferred (style preserved)"
+            else:
+                return "Upload shape and/or color reference images"
+
+        # Update transfer mode when images change
+        shape.change(fn=update_transfer_mode, inputs=[shape, color], outputs=transfer_status)
+        color.change(fn=update_transfer_mode, inputs=[shape, color], outputs=transfer_status)
+
+        source.upload(fn=resize('Face'), inputs=[source, align], outputs=source)
+        shape.upload(fn=resize('Shape'), inputs=[shape, align], outputs=shape)
+        color.upload(fn=resize('Color'), inputs=[color, align], outputs=color)
+
+        btn.click(
+            fn=swap_hair_selective,
+            inputs=[source, shape, color, blending, poisson_iters, poisson_erosion],
+            outputs=[output, error_message],
+            api_name="predict"
+        )
+
+        gr.Markdown('''
+        ### How to use:
+        1. **Upload your source photo** (the face you want to modify)
+        2. **Choose your transfer mode:**
+           - For **hairstyle only**: Upload only a shape reference image
+           - For **color only**: Upload only a color reference image  
+           - For **both**: Upload both shape and color reference images
+        3. **Click "Get the haircut"** and wait for results!
+        
+        
+        ''')
+
+    return demo
+
+def create_api_app():
+    """Create FastAPI app for better frontend integration"""
+    app = FastAPI(title="HairFastGAN Selective Transfer API", version="2.0")
+    
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+    
+    @app.post("/api/hair-transfer", response_model=ImageResponse)
+    async def transfer_hair(request: ImageRequest):
+        """Enhanced API endpoint for hair transfer"""
+        try:
+            print(f"API request received: {len(request.source_image)} chars source")
+            
+            result_base64, message = hair_transfer_api(
+                source_image=request.source_image,
+                shape_image=request.shape_image,
+                color_image=request.color_image,
+                blending=request.blending,
+                poisson_iters=request.poisson_iters,
+                poisson_erosion=request.poisson_erosion
+            )
+            
+            if result_base64:
+                return ImageResponse(
+                    success=True,
+                    result_image=result_base64,
+                    message=message or "Hair transfer completed successfully"
+                )
+            else:
+                return ImageResponse(
+                    success=False,
+                    message=message or "Hair transfer failed"
+                )
+                
+        except Exception as e:
+            print(f"API endpoint error: {e}")
+            raise HTTPException(status_code=500, detail=str(e))
+    
+    @app.get("/api/health")
+    async def health_check():
+        """Health check endpoint"""
+        return {"status": "healthy", "model_loaded": hair_fast_model is not None}
+    
+    return app
+
+if __name__ == '__main__':
+    optimize_for_t4()
+    
+    align_cache = LRUCache(maxsize=10)
+    
+    if not initialize_hairfast_original():
+        print("Failed to initialize HairFast model")
+        exit(1)
+    
+    gradio_demo = get_demo()
+    fastapi_app = create_api_app()
+    
+    def run_fastapi():
+        uvicorn.run(fastapi_app, host="0.0.0.0", port=8000, log_level="info")
+    
+    def run_gradio():
+        gradio_demo.queue(max_size=10, default_concurrency_limit=2)
+        gradio_demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    
+    print("Starting FastAPI server on port 8000...")
+    fastapi_thread = threading.Thread(target=run_fastapi)
+    fastapi_thread.daemon = True
+    fastapi_thread.start()
+    
+    print("Starting Gradio server on port 7860...")
+    run_gradio()

requirements.txt

@@ -0,0 +1,20 @@
+pillow==10.0.0
+face_alignment==1.3.4
+addict==2.4.0
+git+https://github.com/openai/CLIP.git
+gdown==4.7.1
+grpcio==1.63.0
+grpcio_tools==1.63.0
+gradio==4.31.5
+cachetools==5.3.3
+dlib-binary==19.24.1
+opencv-python==4.11.0.86
+torch>=2.0.0
+torchvision>=0.15.0
+numpy>=1.21.0
+scipy>=1.7.0
+scikit-image>=0.18.0
+scikit-learn>=1.0.0
+numba>=0.56.0
+ninja  
+cmake