Complete FaceLift deployment with all dependencies

app.py

@@ -1,5 +1,25 @@
+"""
+FaceLift + Gemini - HuggingFace Space Entry Point
+This is the main file that HuggingFace Space will run.
+"""
+
+import os
+# Set OMP_NUM_THREADS to 1 to avoid libgomp crash in HF Spaces
+os.environ["OMP_NUM_THREADS"] = "1"
+
+import gradio as gr
 from gradio_app import create_demo
 
-demo = create_demo()
-demo.queue(max_size=10)
-demo.launch()
+# HuggingFace Spaces automatically provides GPU
+# The app will use environment variables for API keys
+
+if __name__ == "__main__":
+    # Create the Gradio interface
+    demo = create_demo()
+    
+    # Launch with HuggingFace Space settings
+    demo.queue().launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )

gradio_app.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 """
-FaceLift: Single Image 3D Face Reconstruction
-Generates 3D head models from single images using multi-view diffusion and GS-LRM.
+FaceLift: Single Image 3D Face Reconstruction (Gemini Edition)
+Generates 3D head models from single images using Gemini 2.0 Flash and GS-LRM.
 """
 
 import json
@@ -30,45 +30,39 @@ from einops import rearrange
 from PIL import Image
 from huggingface_hub import snapshot_download
 
-from gslrm.model.gaussians_renderer import render_turntable, imageseq2video
-from mvdiffusion.pipelines.pipeline_mvdiffusion_unclip import StableUnCLIPImg2ImgPipeline
 from utils_folder.face_utils import preprocess_image, preprocess_image_without_cropping
+from gemini_generator import GeminiGenerator
 
 # HuggingFace repository configuration
 HF_REPO_ID = "wlyu/OpenFaceLift"
 
 def download_weights_from_hf() -> Path:
-    """Download model weights from HuggingFace if not already present.
-    
-    Returns:
-        Path to the downloaded repository
-    """
+    """Download model weights from HuggingFace if not already present."""
     workspace_dir = Path(__file__).parent
     
     # Check if weights already exist locally
-    mvdiffusion_path = workspace_dir / "checkpoints/mvdiffusion/pipeckpts"
     gslrm_path = workspace_dir / "checkpoints/gslrm/ckpt_0000000000021125.pt"
-    prompt_embeds_path = workspace_dir / "mvdiffusion/data/fixed_prompt_embeds_6view/clr_embeds.pt"
     
-    if mvdiffusion_path.exists() and gslrm_path.exists() and prompt_embeds_path.exists():
+    if gslrm_path.exists():
         print("Using local model weights")
         return workspace_dir
     
     print(f"Downloading model weights from HuggingFace: {HF_REPO_ID}")
-    print("This may take a few minutes on first run...")
     
     # Download to checkpoints directory
+    # Repo structure is 'gslrm/ckpt...', so we download to 'checkpoints' folder to get 'checkpoints/gslrm/ckpt...'
     snapshot_download(
         repo_id=HF_REPO_ID,
         local_dir=str(workspace_dir / "checkpoints"),
         local_dir_use_symlinks=False,
+        allow_patterns=["gslrm/*"] # Only download GS-LRM
     )
     
     print("Model weights downloaded successfully!")
     return workspace_dir
 
 class FaceLiftPipeline:
-    """Pipeline for FaceLift 3D head generation from single images."""
+    """Pipeline for FaceLift 3D head generation (Gemini Only)."""
     
     def __init__(self):
         # Download weights from HuggingFace if needed
@@ -76,23 +70,18 @@ class FaceLiftPipeline:
         
         # Setup paths
         self.output_dir = workspace_dir / "outputs"
-        self.examples_dir = workspace_dir / "examples"
         self.output_dir.mkdir(exist_ok=True)
         
         # Parameters
         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         self.image_size = 512
-        self.camera_indices = [2, 1, 0, 5, 4, 3]
+        self.camera_indices = [2, 1, 0, 5, 4, 3] # Front, Back, Left, Right, Top, Bottom
         
-        # Load models
-        print("Loading models...")
-        self.mvdiffusion_pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained(
-            str(workspace_dir / "checkpoints/mvdiffusion/pipeckpts"),
-            torch_dtype=torch.float16,
-        )
-        self.mvdiffusion_pipeline.unet.enable_xformers_memory_efficient_attention()
-        self.mvdiffusion_pipeline.to(self.device)
+        # Initialize Gemini Generator
+        self.gemini_generator = GeminiGenerator()
         
+        # Load GS-LRM model (Reconstruction only)
+        print("Loading GS-LRM model...")
         with open(workspace_dir / "configs/gslrm.yaml", "r") as f:
             config = edict(yaml.safe_load(f))
         
@@ -108,84 +97,191 @@ class FaceLiftPipeline:
         self.gs_lrm_model.load_state_dict(checkpoint["model"])
         self.gs_lrm_model.to(self.device)
         
-        self.color_prompt_embedding = torch.load(
-            workspace_dir / "mvdiffusion/data/fixed_prompt_embeds_6view/clr_embeds.pt",
-            map_location=self.device
-        )
-        
         with open(workspace_dir / "utils_folder/opencv_cameras.json", 'r') as f:
             self.cameras_data = json.load(f)["frames"]
         
         print("Models loaded successfully!")
     
-    def generate_3d_head(self, image_path, mode="Single Image", auto_crop=True, guidance_scale=3.0, 
+    def generate_3d_head(self, image_path, api_key, model_type="Gemini", auto_crop=True, guidance_scale=3.0, 
                          random_seed=4, num_steps=50):
-        """Generate 3D head from single image or 6-view strip."""
+        """Generate 3D head from single image."""
         try:
+            # Update API Key if provided
+            if api_key:
+                self.gemini_generator.configure_key(api_key)
+
             # Setup output directory
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
             output_dir = self.output_dir / timestamp
             output_dir.mkdir(exist_ok=True)
             
-            # Load input
-            original_img = Image.open(image_path)
-            input_path = output_dir / "input.png"
-            original_img.save(input_path)
-
+            # Preprocess input
+            original_img = np.array(Image.open(image_path))
+            
+            # Check for pre-generated multiview image (Grid or Strip) BEFORE cropping
+            h, w, _ = original_img.shape
+            aspect_ratio = w / h
+            
+            print(f"[DEBUG] Image dimensions: {w}x{h}, Aspect ratio: {aspect_ratio:.3f}")
+            
+            is_strip = 5.5 < aspect_ratio < 6.5
+            is_grid = 1.1 < aspect_ratio < 2.0  # Widened range to catch cropped/resized grids
+            
+            print(f"[DEBUG] is_strip: {is_strip}, is_grid: {is_grid}")
+            
+            
             selected_views = []
+            original_views = []  # Keep original aspect ratios for multiview composite
             
-            if mode == "6-View Strip":
-                print("Processing 6-View Strip...")
-                # Expect 6x1 strip (3072x512)
-                if original_img.width != self.image_size * 6 or original_img.height != self.image_size:
-                    # Try to resize if aspect ratio is correct
-                    if abs(original_img.width / original_img.height - 6.0) < 0.1:
-                         original_img = original_img.resize((self.image_size * 6, self.image_size), Image.LANCZOS)
-                    else:
-                         raise ValueError(f"Input must be 6x1 strip (e.g. 3072x512). Got {original_img.size}")
-                
-                # Split views
+            if is_strip:
+                print("Detected pre-generated multiview image (6x1). Skipping generation & cropping.")
+                input_image = Image.fromarray(original_img)
+                single_view_width = w // 6
                 for i in range(6):
-                    view = original_img.crop((self.image_size * i, 0, self.image_size * (i + 1), self.image_size))
-                    selected_views.append(view)
-                
-                # Visualization of input
-                multiview_image = original_img
-                multiview_path = output_dir / "multiview.png"
-                multiview_image.save(multiview_path)
-                
+                    left = i * single_view_width
+                    right = (i + 1) * single_view_width
+                    view = input_image.crop((left, 0, right, h))
+                    
+                    # Pad to square (add white borders)
+                    view_w, view_h = view.size
+                    target_size = max(view_w, view_h)
+                    
+                    # Create white canvas
+                    view_square = Image.new('RGB', (target_size, target_size), (255, 255, 255))
+                    # Paste view centered
+                    paste_x = (target_size - view_w) // 2
+                    paste_y = (target_size - view_h) // 2
+                    view_square.paste(view, (paste_x, paste_y))
+                    
+                    original_views.append(view_square.copy())  # Save square original
+                    if view_square.size != (self.image_size, self.image_size):
+                        view_square = view_square.resize((self.image_size, self.image_size))
+                    selected_views.append(view_square)
+                    
+            elif is_grid:
+                # Grid layout detected - could be 3x2 or 2x3
+                # Determine which based on dimensions
+                if w > h:
+                    # Landscape: 3x2 (3 columns, 2 rows)
+                    print(f"Detected 3x2 grid layout. Aspect Ratio: {aspect_ratio}. Skipping generation & cropping.")
+                    input_image = Image.fromarray(original_img)
+                    single_view_width = w // 3
+                    single_view_height = h // 2
+                    
+                    # Row 1: Top 3 views
+                    for i in range(3):
+                        left = i * single_view_width
+                        right = (i + 1) * single_view_width
+                        view = input_image.crop((left, 0, right, single_view_height))
+                        
+                        # Pad to square (add white borders)
+                        view_w, view_h = view.size
+                        target_size = max(view_w, view_h)
+                        
+                        # Create white canvas
+                        view_square = Image.new('RGB', (target_size, target_size), (255, 255, 255))
+                        # Paste view centered
+                        paste_x = (target_size - view_w) // 2
+                        paste_y = (target_size - view_h) // 2
+                        view_square.paste(view, (paste_x, paste_y))
+                        
+                        original_views.append(view_square.copy())  # Save square original
+                        if view_square.size != (self.image_size, self.image_size):
+                            view_square = view_square.resize((self.image_size, self.image_size))
+                        selected_views.append(view_square)
+                    
+                    # Row 2: Bottom 3 views
+                    for i in range(3):
+                        left = i * single_view_width
+                        right = (i + 1) * single_view_width
+                        view = input_image.crop((left, single_view_height, right, h))
+                        
+                        # Pad to square (add white borders)
+                        view_w, view_h = view.size
+                        target_size = max(view_w, view_h)
+                        
+                        # Create white canvas
+                        view_square = Image.new('RGB', (target_size, target_size), (255, 255, 255))
+                        # Paste view centered
+                        paste_x = (target_size - view_w) // 2
+                        paste_y = (target_size - view_h) // 2
+                        view_square.paste(view, (paste_x, paste_y))
+                        
+                        original_views.append(view_square.copy())  # Save square original
+                        if view_square.size != (self.image_size, self.image_size):
+                            view_square = view_square.resize((self.image_size, self.image_size))
+                        selected_views.append(view_square)
+                else:
+                    # Portrait: 2x3 (2 columns, 3 rows)
+                    print(f"Detected 2x3 grid layout. Aspect Ratio: {aspect_ratio}. Skipping generation & cropping.")
+                    input_image = Image.fromarray(original_img)
+                    single_view_width = w // 2
+                    single_view_height = h // 3
+                    
+                    # Process all 6 views row by row
+                    for row in range(3):
+                        for col in range(2):
+                            left = col * single_view_width
+                            right = (col + 1) * single_view_width
+                            top = row * single_view_height
+                            bottom = (row + 1) * single_view_height
+                            view = input_image.crop((left, top, right, bottom))
+                            
+                            # Pad to square (add white borders)
+                            view_w, view_h = view.size
+                            target_size = max(view_w, view_h)
+                            
+                            # Create white canvas
+                            view_square = Image.new('RGB', (target_size, target_size), (255, 255, 255))
+                            # Paste view centered
+                            paste_x = (target_size - view_w) // 2
+                            paste_y = (target_size - view_h) // 2
+                            view_square.paste(view, (paste_x, paste_y))
+                            
+                            original_views.append(view_square.copy())  # Save original
+                            if view_square.size != (self.image_size, self.image_size):
+                                view_square = view_square.resize((self.image_size, self.image_size))
+                            selected_views.append(view_square)
+            
+            
             else:
-                # Single Image Mode
-                input_image_arr = np.array(original_img)
-                input_image = preprocess_image(input_image_arr) if auto_crop else \
-                             preprocess_image_without_cropping(input_image_arr)
+                # Normal flow: Preprocess -> Generate
+                input_image = preprocess_image(original_img) if auto_crop else \
+                             preprocess_image_without_cropping(original_img)
                 
+                # Gemini generation requires API key
+                if not api_key:
+                    raise gr.Error("API Key is required for generating new views. Please provide a Gemini API Key or upload a pre-generated 2x3 or 6x1 grid image.")
+                
+                print("Generating multi-view images with Gemini...")
                 if input_image.size != (self.image_size, self.image_size):
                     input_image = input_image.resize((self.image_size, self.image_size))
                 
-                # Generate multi-view images
-                generator = torch.Generator(device=self.mvdiffusion_pipeline.unet.device)
-                generator.manual_seed(random_seed)
-                
-                result = self.mvdiffusion_pipeline(
-                    input_image, None,
-                    prompt_embeds=self.color_prompt_embedding,
-                    guidance_scale=guidance_scale,
-                    num_images_per_prompt=1,
-                    num_inference_steps=num_steps,
-                    generator=generator,
-                    eta=1.0,
-                )
-                
-                selected_views = result.images[:6]
-                
-                # Save multi-view composite
-                multiview_image = Image.new("RGB", (self.image_size * 6, self.image_size))
-                for i, view in enumerate(selected_views):
-                    multiview_image.paste(view, (self.image_size * i, 0))
-                
-                multiview_path = output_dir / "multiview.png"
-                multiview_image.save(multiview_path)
+                try:
+                    selected_views = self.gemini_generator.generate_multiview(input_image)
+                    original_views = [v.copy() for v in selected_views]  # For Gemini, they're already square
+                except Exception as e:
+                    raise gr.Error(f"Gemini generation failed: {str(e)}. Try uploading a pre-generated 2x3 grid instead.")
+
+            # Save processed input (for reference)
+            input_path = output_dir / "input.png"
+            input_image.save(input_path)
+            
+            # Save multi-view composite (preserve original aspect ratios)
+            # Use original_views instead of selected_views
+            max_height = max(view.size[1] for view in original_views)
+            total_width = sum(view.size[0] for view in original_views)
+            
+            multiview_image = Image.new("RGB", (total_width, max_height), (255, 255, 255))
+            x_offset = 0
+            for view in original_views:
+                # Center vertically if view is shorter than max_height
+                y_offset = (max_height - view.size[1]) // 2
+                multiview_image.paste(view, (x_offset, y_offset))
+                x_offset += view.size[0]
+            
+            multiview_path = output_dir / "multiview.png"
+            multiview_image.save(multiview_path)
             
             # Prepare 3D reconstruction input
             view_arrays = [np.array(view) for view in selected_views]
@@ -241,58 +337,46 @@ class FaceLiftPipeline:
             output_path = output_dir / "output.png"
             Image.fromarray(comp_image).save(output_path)
             
-            # Generate turntable video
-            turntable_frames = render_turntable(gaussians, rendering_resolution=self.image_size, 
-                                               num_views=180)
-            turntable_frames = rearrange(turntable_frames, "h (v w) c -> v h w c", v=180)
-            turntable_frames = np.ascontiguousarray(turntable_frames)
-            
-            turntable_path = output_dir / "turntable.mp4"
-            imageseq2video(turntable_frames, str(turntable_path), fps=30)
-            
-            return str(input_path), str(multiview_path), str(output_path), \
-                   str(turntable_path), str(ply_path)
+            return str(input_path), str(multiview_path), str(output_path), str(ply_path)
             
         except Exception as e:
             raise gr.Error(f"Generation failed: {str(e)}")
 
 
-def main():
-    """Run the FaceLift application."""
+def create_demo():
+    """Create and return the Gradio demo interface."""
     pipeline = FaceLiftPipeline()
     
-    # Load examples (Filtered for Single Image)
-    examples = []
-    if pipeline.examples_dir.exists():
-        examples = [[str(f)] for f in sorted(pipeline.examples_dir.iterdir()) 
-                   if f.suffix.lower() in {'.png', '.jpg', '.jpeg'}]
-    
-    # Create interface
     demo = gr.Interface(
         fn=pipeline.generate_3d_head,
-        title="FaceLift: Single Image 3D Face Reconstruction",
+        title="FaceLift: Single Image 3D Face Reconstruction (Gemini)",
         description="""
-        Transform a single portrait image OR a 6-view strip into a complete 3D head model.
+        Transform a single portrait into a complete 3D head model using Gemini 2.0 Flash and GS-LRM.
         """,
         inputs=[
-            gr.Image(type="filepath", label="Input Image (Portrait or 6x1 Strip)"),
-            gr.Radio(["Single Image", "6-View Strip"], value="Single Image", label="Input Mode"),
-            gr.Checkbox(value=True, label="Auto Cropping (Single Image Only)"),
-            gr.Slider(1.0, 10.0, 3.0, step=0.1, label="Guidance Scale"),
-            gr.Number(value=4, label="Random Seed"),
-            gr.Slider(10, 100, 50, step=5, label="Generation Steps"),
+            gr.Image(type="filepath", label="Input Portrait Image"),
+            gr.Textbox(label="API Key (Gemini)", type="password", placeholder="Optional - only needed if generating new views", value=""),
+            gr.Dropdown(choices=["Gemini"], value="Gemini", label="Generation Model", visible=False),
+            gr.Checkbox(value=True, label="Auto Cropping"),
+            gr.Slider(1.0, 10.0, 3.0, step=0.1, label="Guidance Scale (Unused)"),
+            gr.Number(value=4, label="Random Seed (Unused)"),
+            gr.Slider(10, 100, 50, step=5, label="Generation Steps (Unused)"),
         ],
         outputs=[
             gr.Image(label="Processed Input"),
             gr.Image(label="Multi-view Generation"),
             gr.Image(label="3D Reconstruction"),
-            gr.Video(label="Turntable Animation"),
             gr.File(label="3D Model (.ply)"),
         ],
-        examples=examples,
         allow_flagging="never",
     )
     
+    return demo
+
+
+def main():
+    """Main function for local development."""
+    demo = create_demo()
     demo.queue(max_size=10)
     demo.launch(share=True, server_name="0.0.0.0", server_port=7860, show_error=True)