Update to clean FaceLift (no Gemini)

app.py

@@ -1,25 +1,5 @@
-"""
-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
 
-# 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
-    )
+demo = create_demo()
+demo.queue(max_size=10)
+demo.launch()

gradio_app.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 """
-FaceLift: Single Image 3D Face Reconstruction (Gemini Edition)
-Generates 3D head models from single images using Gemini 2.0 Flash and GS-LRM.
+FaceLift: Single Image 3D Face Reconstruction
+Generates 3D head models from single images using multi-view diffusion and GS-LRM.
 """
 
 import json
@@ -30,39 +30,45 @@ 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."""
+    """Download model weights from HuggingFace if not already present.
+    
+    Returns:
+        Path to the downloaded repository
+    """
     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 gslrm_path.exists():
+    if mvdiffusion_path.exists() and gslrm_path.exists() and prompt_embeds_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 (Gemini Only)."""
+    """Pipeline for FaceLift 3D head generation from single images."""
     
     def __init__(self):
         # Download weights from HuggingFace if needed
@@ -70,18 +76,23 @@ 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] # Front, Back, Left, Right, Top, Bottom
+        self.camera_indices = [2, 1, 0, 5, 4, 3]
         
-        # Initialize Gemini Generator
-        self.gemini_generator = GeminiGenerator()
+        # 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)
         
-        # 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))
         
@@ -97,191 +108,84 @@ 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, api_key, model_type="Gemini", auto_crop=True, guidance_scale=3.0, 
+    def generate_3d_head(self, image_path, mode="Single Image", auto_crop=True, guidance_scale=3.0, 
                          random_seed=4, num_steps=50):
-        """Generate 3D head from single image."""
+        """Generate 3D head from single image or 6-view strip."""
         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)
             
-            # 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}")
-            
-            
+            # Load input
+            original_img = Image.open(image_path)
+            input_path = output_dir / "input.png"
+            original_img.save(input_path)
+
             selected_views = []
-            original_views = []  # Keep original aspect ratios for multiview composite
             
-            if is_strip:
-                print("Detected pre-generated multiview image (6x1). Skipping generation & cropping.")
-                input_image = Image.fromarray(original_img)
-                single_view_width = w // 6
+            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
                 for i in range(6):
-                    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:
-                # Normal flow: Preprocess -> Generate
-                input_image = preprocess_image(original_img) if auto_crop else \
-                             preprocess_image_without_cropping(original_img)
+                    view = original_img.crop((self.image_size * i, 0, self.image_size * (i + 1), self.image_size))
+                    selected_views.append(view)
                 
-                # 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.")
+                # Visualization of input
+                multiview_image = original_img
+                multiview_path = output_dir / "multiview.png"
+                multiview_image.save(multiview_path)
+                
+            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)
                 
-                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))
                 
-                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)
+                # 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)
             
             # Prepare 3D reconstruction input
             view_arrays = [np.array(view) for view in selected_views]
@@ -337,46 +241,58 @@ class FaceLiftPipeline:
             output_path = output_dir / "output.png"
             Image.fromarray(comp_image).save(output_path)
             
-            return str(input_path), str(multiview_path), str(output_path), str(ply_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)
             
         except Exception as e:
             raise gr.Error(f"Generation failed: {str(e)}")
 
 
-def create_demo():
-    """Create and return the Gradio demo interface."""
+def main():
+    """Run the FaceLift application."""
     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 (Gemini)",
+        title="FaceLift: Single Image 3D Face Reconstruction",
         description="""
-        Transform a single portrait into a complete 3D head model using Gemini 2.0 Flash and GS-LRM.
+        Transform a single portrait image OR a 6-view strip into a complete 3D head model.
         """,
         inputs=[
-            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)"),
+            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"),
         ],
         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)