api

app.py

@@ -1,20 +1,74 @@
 import hashlib
 import os
+import base64
 from io import BytesIO
+from typing import Optional
 
-import gradio as gr
 import grpc
+import uvicorn
 from PIL import Image
 from cachetools import LRUCache
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 
 from inference_pb2 import HairSwapRequest, HairSwapResponse
 from inference_pb2_grpc import HairSwapServiceStub
 from utils.shape_predictor import align_face
 
+app = FastAPI(
+    title="HairFastGAN API",
+    description="API for HairFastGAN: Realistic and Robust Hair Transfer with a Fast Encoder-Based Approach",
+    version="1.0.0"
+)
+
+# Global cache for aligned faces
+align_cache = LRUCache(maxsize=10)
+
+
+class HairSwapRequest(BaseModel):
+    face: str  # Base64 encoded image
+    shape: Optional[str] = None  # Base64 encoded image
+    color: Optional[str] = None  # Base64 encoded image
+    blending: str = "Article"
+    poisson_iters: int = 0
+    poisson_erosion: int = 15
+    align_face_img: bool = True
+    align_shape_img: bool = True
+    align_color_img: bool = True
+
+
+class HairSwapResponse(BaseModel):
+    image: str  # Base64 encoded image
+
+
+def base64_to_image(base64_str: str) -> Image.Image:
+    """Convert base64 string to PIL Image"""
+    if not base64_str:
+        return None
+    
+    # Remove header if present
+    if "base64," in base64_str:
+        base64_str = base64_str.split("base64,")[1]
+    
+    image_bytes = base64.b64decode(base64_str)
+    image = Image.open(BytesIO(image_bytes))
+    return image
+
+
+def image_to_base64(img: Image.Image, format="JPEG") -> str:
+    """Convert PIL Image to base64 string"""
+    if img is None:
+        return None
+    
+    buffered = BytesIO()
+    img.save(buffered, format=format)
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    return f"data:image/{format.lower()};base64,{img_str}"
+
 
 def get_bytes(img):
     if img is None:
-        return img
+        return None
 
     buffered = BytesIO()
     img.save(buffered, format="JPEG")
@@ -39,113 +93,102 @@ def center_crop(img):
     return img
 
 
-def resize(name):
-    def resize_inner(img, align):
-        global align_cache
-
-        if name in align:
-            img_hash = hashlib.md5(get_bytes(img)).hexdigest()
-
-            if img_hash not in align_cache:
-                img = align_face(img, return_tensors=False)[0]
-                align_cache[img_hash] = img
-            else:
-                img = align_cache[img_hash]
-
-        elif img.size != (1024, 1024):
-            img = center_crop(img)
-            img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
-
-        return img
-
-    return resize_inner
-
-
-def swap_hair(face, shape, color, blending, poisson_iters, poisson_erosion):
-    if not face and not shape and not color:
-        return gr.update(visible=False), gr.update(value="Need to upload a face and at least a shape or color ❗", visible=True)
-    elif not face:
-        return gr.update(visible=False), gr.update(value="Need to upload a face ❗", visible=True)
-    elif not shape and not color:
-        return gr.update(visible=False), gr.update(value="Need to upload at least a shape or color ❗", visible=True)
-
-    face_bytes, shape_bytes, color_bytes = map(lambda item: get_bytes(item), (face, shape, color))
-
-    if shape_bytes is None:
-        shape_bytes = b'face'
-    if color_bytes is None:
-        color_bytes = b'shape'
-
-    with grpc.insecure_channel(os.environ['SERVER']) as channel:
-        stub = HairSwapServiceStub(channel)
-
-        output: HairSwapResponse = stub.swap(
-            HairSwapRequest(face=face_bytes, shape=shape_bytes, color=color_bytes, blending=blending,
-                            poisson_iters=poisson_iters, poisson_erosion=poisson_erosion, use_cache=True)
-        )
-
-    output = bytes_to_image(output.image)
-    return gr.update(value=output, visible=True), gr.update(visible=False)
-
-
-def get_demo():
-    with gr.Blocks() as demo:
-        gr.Markdown("## HairFastGan")
-        gr.Markdown(
-            '<div style="display: flex; align-items: center; gap: 10px;">'
-            '<span>Official HairFastGAN Gradio demo:</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>'
-        )
-        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")
-                with gr.Accordion("Advanced Options", open=False):
-                    blending = gr.Radio(["Article", "Alternative_v1", "Alternative_v2"], value='Article',
-                                        label="Color Encoder version", info="Selects a model for hair color transfer.")
-                    poisson_iters = gr.Slider(0, 2500, value=0, step=1, label="Poisson iters",
-                                              info="The power of blending with the original image, helps to recover more details. Not included in the article, disabled by default.")
-                    poisson_erosion = gr.Slider(1, 100, value=15, step=1, label="Poisson erosion",
-                                                info="Smooths out the blending area.")
-                    align = gr.CheckboxGroup(["Face", "Shape", "Color"], value=["Face", "Shape", "Color"],
-                                             label="Image cropping [recommended]",
-                                             info="Selects which images to crop by face")
-                btn = gr.Button("Get the haircut")
-            with gr.Column():
-                output = gr.Image(label="Your result")
-                error_message = gr.Textbox(label="⚠️ Error ⚠️", visible=False, elem_classes="error-message")
-
-        gr.Examples(examples=[["input/0.png", "input/1.png", "input/2.png"], ["input/6.png", "input/7.png", None],
-                              ["input/10.jpg", None, "input/11.jpg"]],
-                    inputs=[source, shape, color], outputs=output)
-
-        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, inputs=[source, shape, color, blending, poisson_iters, poisson_erosion],
-                  outputs=[output, error_message])
-
-        gr.Markdown('''To cite the paper by the authors
-    ```
-        @article{nikolaev2024hairfastgan,
-          title={HairFastGAN: Realistic and Robust Hair Transfer with a Fast Encoder-Based Approach},
-          author={Nikolaev, Maxim and Kuznetsov, Mikhail and Vetrov, Dmitry and Alanov, Aibek},
-          journal={arXiv preprint arXiv:2404.01094},
-          year={2024}
-        }
-    ```
-        ''')
-    return demo
-
-
-if __name__ == '__main__':
-    align_cache = LRUCache(maxsize=10)
-    demo = get_demo()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+def process_image(img, should_align=True):
+    global align_cache
+    
+    if should_align:
+        img_bytes = get_bytes(img)
+        img_hash = hashlib.md5(img_bytes).hexdigest()
+
+        if img_hash not in align_cache:
+            img = align_face(img, return_tensors=False)[0]
+            align_cache[img_hash] = img
+        else:
+            img = align_cache[img_hash]
+    elif img.size != (1024, 1024):
+        img = center_crop(img)
+        img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
+
+    return img
+
+
+@app.post("/swap-hair", response_model=HairSwapResponse)
+async def swap_hair(request: HairSwapRequest):
+    """
+    Swap hair in the source face image with the shape and/or color from provided images.
+    
+    - face: Source image as base64 string (required)
+    - shape: Image with desired hairstyle shape as base64 string (optional, but either shape or color is required)
+    - color: Image with desired hair color as base64 string (optional, but either shape or color is required)
+    - blending: Color Encoder version ("Article", "Alternative_v1", or "Alternative_v2")
+    - poisson_iters: Power of blending with original image (0-2500)
+    - poisson_erosion: Smooths out blending area (1-100)
+    - align_face_img: Whether to align the face image
+    - align_shape_img: Whether to align the shape image
+    - align_color_img: Whether to align the color image
+    
+    Returns the processed image as a base64-encoded JPEG.
+    """
+    # Validate inputs
+    if not request.face:
+        raise HTTPException(status_code=400, detail="Need to provide a face image")
+    if not request.shape and not request.color:
+        raise HTTPException(status_code=400, detail="Need to provide at least a shape or color image")
+    
+    # Convert base64 to images
+    try:
+        face_img = base64_to_image(request.face)
+        
+        shape_img = None
+        if request.shape:
+            shape_img = base64_to_image(request.shape)
+            shape_img = process_image(shape_img, request.align_shape_img)
+            
+        color_img = None
+        if request.color:
+            color_img = base64_to_image(request.color)
+            color_img = process_image(color_img, request.align_color_img)
+            
+        # Process face image (always required)
+        face_img = process_image(face_img, request.align_face_img)
+        
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error processing images: {str(e)}")
+    
+    # Convert images to bytes
+    face_bytes = get_bytes(face_img)
+    shape_bytes = get_bytes(shape_img) if shape_img else b'face'
+    color_bytes = get_bytes(color_img) if color_img else b'shape'
+    
+    # Call gRPC service
+    try:
+        with grpc.insecure_channel(os.environ['SERVER']) as channel:
+            stub = HairSwapServiceStub(channel)
+            
+            output: HairSwapResponse = stub.swap(
+                HairSwapRequest(
+                    face=face_bytes, 
+                    shape=shape_bytes, 
+                    color=color_bytes, 
+                    blending=request.blending,
+                    poisson_iters=request.poisson_iters, 
+                    poisson_erosion=request.poisson_erosion, 
+                    use_cache=True
+                )
+            )
+            
+        # Convert result to image
+        output_img = bytes_to_image(output.image)
+        
+        # Convert image to base64
+        base64_img = image_to_base64(output_img)
+        
+        return HairSwapResponse(image=base64_img)
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during hair swapping: {str(e)}")
+
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 8000))
+    uvicorn.run(app, host="0.0.0.0", port=port)

requirements.txt

@@ -5,6 +5,7 @@ git+https://github.com/openai/CLIP.git
 gdown==3.12.2
 grpcio==1.63.0
 grpcio_tools==1.63.0
-gradio==4.31.5
+fastapi==0.104.1
+uvicorn==0.23.2
 cachetools==5.3.3
 dlib==19.24.1