Update app.py

app.py

@@ -1,151 +1,394 @@
 import hashlib
 import os
 from io import BytesIO
+import base64
 
 import gradio as gr
-import grpc
 from PIL import Image
 from cachetools import LRUCache
+import torch
+import numpy as np
+
+# Direct HairFast imports (no gRPC needed!)
+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
 
-from inference_pb2 import HairSwapRequest, HairSwapResponse
-from inference_pb2_grpc import HairSwapServiceStub
 from utils.shape_predictor import align_face
 
+# Global variables
+hair_fast_model = None
+align_cache = LRUCache(maxsize=10)
+
+
+def initialize_hairfast():
+    """Initialize HairFast model"""
+    global hair_fast_model
+    
+    if not HAIRFAST_AVAILABLE:
+        print("❌ HairFast not available")
+        return False
+        
+    try:
+        print("🔄 Initializing HairFast model...")
+        
+        # Get default arguments
+        parser = get_parser()
+        args = parser.parse_args([])  # Use default arguments
+        
+        # Override some settings for HF Spaces
+        args.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        args.batch_size = 1  # Keep small for HF Spaces
+        
+        # Initialize HairFast
+        hair_fast_model = HairFast(args)
+        
+        print(f"✅ HairFast initialized successfully on {args.device}!")
+        return True
+        
+    except Exception as e:
+        print(f"❌ HairFast initialization failed: {e}")
+        hair_fast_model = None
+        return False
+
 
 def get_bytes(img):
+    """Convert PIL Image to bytes"""
     if img is None:
         return img
-
     buffered = BytesIO()
     img.save(buffered, format="JPEG")
     return buffered.getvalue()
 
 
 def bytes_to_image(image: bytes) -> Image.Image:
+    """Convert bytes to PIL Image"""
     image = Image.open(BytesIO(image))
     return image
 
 
+def base64_to_image(base64_string):
+    """Convert base64 string to PIL Image"""
+    try:
+        if base64_string.startswith('data:image'):
+            base64_string = base64_string.split(',')[1]
+        image_bytes = base64.b64decode(base64_string)
+        return Image.open(BytesIO(image_bytes))
+    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"""
+    if image is None:
+        return None
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    img_bytes = buffered.getvalue()
+    img_base64 = base64.b64encode(img_bytes).decode('utf-8')
+    return f"data:image/jpeg;base64,{img_base64}"
+
+
+def pil_to_tensor(image):
+    """Convert PIL to tensor for HairFast"""
+    if isinstance(image, Image.Image):
+        # Convert to tensor format expected by HairFast
+        image_array = np.array(image)
+        if image_array.max() > 1:
+            image_array = image_array / 255.0
+        tensor = torch.from_numpy(image_array).permute(2, 0, 1).float()
+        return tensor
+    return image
+
+
+def tensor_to_pil(tensor):
+    """Convert tensor to PIL Image"""
+    if isinstance(tensor, torch.Tensor):
+        if tensor.dim() == 4:
+            tensor = tensor.squeeze(0)
+        if tensor.dim() == 3:
+            tensor = tensor.permute(1, 2, 0)
+        tensor = tensor.detach().cpu().numpy()
+        if tensor.max() <= 1:
+            tensor = (tensor * 255).astype(np.uint8)
+        return Image.fromarray(tensor)
+    return tensor
+
+
 def center_crop(img):
+    """Center crop image to square"""
     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):
+    """Image resize function with face alignment"""
     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
+                try:
+                    img = align_face(img, return_tensors=False)[0]
+                    align_cache[img_hash] = img
+                except Exception as e:
+                    print(f"Face alignment failed: {e}")
+                    img = center_crop(img)
+                    img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
             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):
+def swap_hair_direct(face, shape, color, blending, poisson_iters, poisson_erosion):
+    """Direct hair swapping using HairFast (no gRPC)"""
+    global hair_fast_model
+    
+    # Initialize model if needed
+    if hair_fast_model is None:
+        if not initialize_hairfast():
+            return gr.update(visible=False), gr.update(
+                value="❌ HairFast model not available. Please check if all model files are uploaded.", 
+                visible=True
+            )
+    
+    # Validation
     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)
+        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)
+        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)
+        return gr.update(visible=False), gr.update(
+            value="Need to upload at least a shape or color ❗", 
+            visible=True
         )
 
-    output = bytes_to_image(output.image)
-    return gr.update(value=output, visible=True), gr.update(visible=False)
+    try:
+        print("🔄 Starting hair transfer...")
+        
+        # Use shape as color if color is not provided
+        if color is None:
+            color = shape
+        if shape is None:
+            shape = color
+        
+        # Direct HairFast inference
+        result_tensor = hair_fast_model.swap(
+            face_img=face,
+            shape_img=shape, 
+            color_img=color,
+            benchmark=False,
+            align=True,  # Use face alignment
+            seed=3407
+        )
+        
+        # Convert result tensor to PIL Image
+        result_image = tensor_to_pil(result_tensor)
+        
+        print("✅ Hair transfer completed successfully!")
+        return gr.update(value=result_image, visible=True), gr.update(visible=False)
+
+    except Exception as e:
+        error_msg = f"❌ Hair transfer failed: {str(e)}"
+        print(error_msg)
+        return gr.update(visible=False), gr.update(value=error_msg, visible=True)
+
+
+def hair_transfer_api(source_image, shape_image=None, color_image=None, 
+                     blending="Article", poisson_iters=0, poisson_erosion=15):
+    """API function for React integration"""
+    global hair_fast_model
+    
+    try:
+        # Handle base64 inputs
+        if isinstance(source_image, str):
+            source_image = base64_to_image(source_image)
+        if isinstance(shape_image, str):
+            shape_image = base64_to_image(shape_image)
+        if isinstance(color_image, str):
+            color_image = base64_to_image(color_image)
+        
+        # Initialize model if needed
+        if hair_fast_model is None:
+            if not initialize_hairfast():
+                return None, "❌ HairFast model not available"
+        
+        # Validation
+        if source_image is None:
+            return None, "❌ Source image is required"
+        
+        # Use source as reference if no references provided
+        if shape_image is None and color_image is None:
+            return None, "❌ At least shape or color reference image is required"
+        
+        if color_image is None:
+            color_image = shape_image
+        if shape_image is None:
+            shape_image = color_image
+        
+        # Direct HairFast inference
+        result_tensor = hair_fast_model.swap(
+            face_img=source_image,
+            shape_img=shape_image,
+            color_img=color_image,
+            benchmark=False,
+            align=True,
+            seed=3407
+        )
+        
+        # Convert to PIL and then base64
+        result_image = tensor_to_pil(result_tensor)
+        result_base64 = image_to_base64(result_image)
+        
+        return result_base64, "✅ Hair transfer completed successfully!"
+        
+    except Exception as e:
+        error_msg = f"❌ API Error: {str(e)}"
+        print(error_msg)
+        return None, error_msg
 
 
 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>'
-        )
+    """Create Gradio interface"""
+    with gr.Blocks(
+        title="HairFastGAN Direct API", 
+        theme=gr.themes.Soft()
+    ) as demo:
+        
+        gr.HTML("""
+        <div style="text-align: center; padding: 20px;">
+            <h1>🎨 HairFastGAN - Direct Model Inference</h1>
+            <p>High-quality hair transfer without gRPC dependency</p>
+        </div>
+        """)
+        
         with gr.Row():
             with gr.Column():
-                source = gr.Image(label="Source photo to try on the hairstyle", type="pil")
+                gr.HTML("<h3>📤 Input Images</h3>")
+                
+                source = gr.Image(
+                    label="Source Photo (Person's Face)", 
+                    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")
+                    shape = gr.Image(
+                        label="Hair Shape Reference (Optional)", 
+                        type="pil"
+                    )
+                    color = gr.Image(
+                        label="Hair Color Reference (Optional)", 
+                        type="pil"
+                    )
+                
+                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 Iterations",
+                        info="Detail recovery strength"
+                    )
+                    poisson_erosion = gr.Slider(
+                        1, 100, value=15, step=1, 
+                        label="Poisson Erosion",
+                        info="Blending smoothness"
+                    )
+                    align = gr.CheckboxGroup(
+                        ["Face", "Shape", "Color"], 
+                        value=["Face", "Shape", "Color"],
+                        label="Face Alignment [Recommended]"
+                    )
+                
+                btn = gr.Button("🎨 Transfer Hair Style", variant="primary", size="lg")
+                
             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)
+                gr.HTML("<h3>📥 Result</h3>")
+                
+                output = gr.Image(label="Result Image", type="pil")
+                error_message = gr.Textbox(
+                    label="⚠️ Status", 
+                    visible=False,
+                    elem_classes="error-message"
+                )
+
+        # Example gallery
+        gr.HTML("<h3>💡 Examples</h3>")
+        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
+        )
 
+        # Event handlers  
         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])
+        btn.click(
+            fn=swap_hair_direct,
+            inputs=[source, shape, color, blending, poisson_iters, poisson_erosion],
+            outputs=[output, error_message],
+            api_name="predict"  # For React integration
+        )
 
-        gr.Markdown('''To cite the paper by the authors
-    ```
+        # Citation
+        gr.Markdown('''
+        ### 📖 Citation
+        ```bibtex
         @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__':
+    # Initialize cache
     align_cache = LRUCache(maxsize=10)
+    
+    # Create demo
     demo = get_demo()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    
+    # Launch with API enabled
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_api=True,
+        share=False
+    )