Update app.py

app.py

@@ -7,21 +7,18 @@ import torch
 from diffusers import Flux2KleinPipeline
 from PIL import Image
 
+# NEW: for reading first frame
+import cv2
+
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 MAX_SEED = np.iinfo(np.int32).max
 
-# Model repository ID for 9B distilled
 REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-9B"
-
-# LoRA repository and file
 LORA_REPO_ID = "Alissonerdx/BFS-Best-Face-Swap"
 LORA_FILENAME = "bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors"
 
-# Fixed prompt for face swapping
-#FACE_SWAP_PROMPT = "head_swap: start with Picture 1 as the base image, keeping its lighting, environment, and background. remove the head from Picture 1 completely and replace it with the head from Picture 2, strictly preserving the hair, eye color, nose structure of Picture 2. copy the direction of the eye, head rotation, micro expressions from Picture 1, high quality, sharp details, 4k."
-
 FACE_SWAP_PROMPT = """head_swap: start with Picture 1 as the base image, keeping its lighting, environment, and background. Remove the head from Picture 1 completely and replace it with the head from Picture 2.
 
 FROM PICTURE 1 (strictly preserve):
@@ -45,24 +42,60 @@ print(f"Loading LoRA from {LORA_REPO_ID}...")
 pipe.load_lora_weights(LORA_REPO_ID, weight_name=LORA_FILENAME)
 print("LoRA loaded successfully!")
 
-def update_dimensions_from_image(target_image):
+
+def first_frame_from_video(video_value) -> Image.Image:
     """
-    Update width/height based on target image aspect ratio.
+    Convert a Gradio Video input into the first frame as a PIL Image.
 
-    Keeps one side at 1024 and scales the other proportionally,
-    with both sides as multiples of 8.
+    Gradio Video value is typically either:
+      - a filepath string
+      - or a dict-like object with a path
+      - or (older versions) a tuple
+    We handle the common cases robustly.
+    """
+    if video_value is None:
+        return None
+
+    video_path = None
+
+    # Common shapes across gradio versions:
+    if isinstance(video_value, str):
+        video_path = video_value
+    elif isinstance(video_value, dict) and "path" in video_value:
+        video_path = video_value["path"]
+    elif isinstance(video_value, (list, tuple)) and len(video_value) > 0:
+        # sometimes it's (path, metadata) or [path, ...]
+        video_path = video_value[0]
+    else:
+        # last attempt: attribute access
+        video_path = getattr(video_value, "path", None)
 
-    Args:
-        target_image: PIL Image of the target/body image.
+    if not video_path or not os.path.exists(video_path):
+        raise gr.Error("Could not read the uploaded video file.")
 
-    Returns:
-        tuple: A tuple of (width, height) integers, both multiples of 8.
+    cap = cv2.VideoCapture(video_path)
+    ok, frame_bgr = cap.read()
+    cap.release()
+
+    if not ok or frame_bgr is None:
+        raise gr.Error("Could not extract the first frame from the video.")
+
+    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(frame_rgb)
+
+
+def update_dimensions_from_video(target_video):
     """
+    Update width/height based on the first frame of the uploaded video.
+
+    Keeps one side at 1024 and scales the other proportionally,
+    with both sides as multiples of 8.
+    """
+    target_image = first_frame_from_video(target_video)
     if target_image is None:
-        return 1024, 1024  # Default dimensions
+        return 1024, 1024
 
     img_width, img_height = target_image.size
-
     aspect_ratio = img_width / img_height
 
     if aspect_ratio >= 1:  # Landscape or square
@@ -72,11 +105,9 @@ def update_dimensions_from_image(target_image):
         new_height = 1024
         new_width = int(1024 * aspect_ratio)
 
-    # Round to nearest multiple of 8
     new_width = round(new_width / 8) * 8
     new_height = round(new_height / 8) * 8
 
-    # Ensure within valid range (minimum 256, maximum 1024)
     new_width = max(256, min(1024, new_width))
     new_height = max(256, min(1024, new_height))
 
@@ -86,7 +117,7 @@ def update_dimensions_from_image(target_image):
 @spaces.GPU(duration=85)
 def face_swap(
     reference_face: Image.Image,
-    target_image: Image.Image,
+    target_video,  # CHANGED: now a video input
     seed: int = 42,
     randomize_seed: bool = False,
     width: int = 1024,
@@ -95,32 +126,17 @@ def face_swap(
     guidance_scale: float = 1.0,
     progress=gr.Progress(track_tqdm=True)
 ):
-    """
-    Perform face swapping using FLUX.2 Klein 9B with LoRA.
-
-    Args:
-        reference_face: The face image to swap in (Picture 2).
-        target_image: The target body/base image (Picture 1).
-        seed: Random seed for reproducible generation.
-        randomize_seed: Set to True to use a random seed.
-        width: Output image width in pixels (256-1024, must be multiple of 8).
-        height: Output image height in pixels (256-1024, must be multiple of 8).
-        num_inference_steps: Number of denoising steps (default 4 for distilled).
-        guidance_scale: How closely to follow the prompt (default 1.0 for distilled).
-
-    Returns:
-        tuple: A tuple containing the generated PIL Image and the seed used.
-    """
-    if reference_face is None or target_image is None:
-        raise gr.Error("Please provide both a reference face and a target image!")
+    if reference_face is None or target_video is None:
+        raise gr.Error("Please provide both a reference face and a target video!")
+
+    # Extract first frame to use as target image
+    target_image = first_frame_from_video(target_video)
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    # Important: Pass target image (body) first, then reference face
-    # This matches the prompt structure: Picture 1 = target, Picture 2 = reference
     image_list = [target_image, reference_face]
 
     progress(0.2, desc="Swapping face...")
@@ -135,7 +151,6 @@ def face_swap(
         generator=generator,
     ).images[0]
 
-    # Return slider comparison (before, after) and seed
     return (target_image, image), seed
 
 
@@ -150,7 +165,6 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-
     with gr.Column(elem_id="col-container"):
         gr.Markdown("""# Face Swap with FLUX.2 Klein 9B
 
@@ -166,14 +180,15 @@ Swap faces using Flux.2 Klein 9B [Alissonerdx/BFS-Best-Face-Swap](https://huggin
                         sources=["upload"],
                         elem_classes="image-container"
                     )
-        
-                    target_image = gr.Image(
-                        label="Target Image (Body/Scene)",
-                        type="pil",
-                        sources=["upload"],
-                        elem_classes="image-container"
+
+                    # CHANGED: target image -> target video
+                    target_video = gr.Video(
+                        label="Target Video (Body/Scene) - first frame will be used",
+                        sources=["upload"]
                     )
+
                 run_button = gr.Button("Swap Face", visible=False)
+
                 with gr.Accordion("Advanced Settings", open=False):
                     seed = gr.Slider(
                         label="Seed",
@@ -182,9 +197,9 @@ Swap faces using Flux.2 Klein 9B [Alissonerdx/BFS-Best-Face-Swap](https://huggin
                         step=1,
                         value=0,
                     )
-        
+
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        
+
                     with gr.Row():
                         width = gr.Slider(
                             label="Width",
@@ -193,7 +208,7 @@ Swap faces using Flux.2 Klein 9B [Alissonerdx/BFS-Best-Face-Swap](https://huggin
                             step=8,
                             value=1024,
                         )
-        
+
                         height = gr.Slider(
                             label="Height",
                             minimum=256,
@@ -201,7 +216,7 @@ Swap faces using Flux.2 Klein 9B [Alissonerdx/BFS-Best-Face-Swap](https://huggin
                             step=8,
                             value=1024,
                         )
-        
+
                     with gr.Row():
                         num_inference_steps = gr.Slider(
                             label="Inference Steps",
@@ -211,7 +226,7 @@ Swap faces using Flux.2 Klein 9B [Alissonerdx/BFS-Best-Face-Swap](https://huggin
                             value=4,
                             info="Number of denoising steps (4 is optimal for distilled model)"
                         )
-        
+
                         guidance_scale = gr.Slider(
                             label="Guidance Scale",
                             minimum=0.0,
@@ -221,28 +236,24 @@ Swap faces using Flux.2 Klein 9B [Alissonerdx/BFS-Best-Face-Swap](https://huggin
                             info="How closely to follow the prompt (1.0 is optimal for distilled model)"
                         )
 
-
             comparison_slider = gr.ImageSlider(
                 label="Before / After",
                 type="pil"
             )
 
-        
-
-        
         seed_output = gr.Number(label="Seed Used", visible=False)
 
-    # Auto-update dimensions when target image is uploaded
-    target_image.upload(
-        fn=update_dimensions_from_image,
-        inputs=[target_image],
+    # CHANGED: auto-update dimensions when target video is uploaded/changed
+    target_video.change(
+        fn=update_dimensions_from_video,
+        inputs=[target_video],
         outputs=[width, height]
     )
 
-    # Create a shared input/output configuration
+    # CHANGED: swap inputs use target_video instead of target_image
     swap_inputs = [
         reference_face,
-        target_image,
+        target_video,
         seed,
         randomize_seed,
         width,
@@ -252,31 +263,26 @@ Swap faces using Flux.2 Klein 9B [Alissonerdx/BFS-Best-Face-Swap](https://huggin
     ]
     swap_outputs = [comparison_slider, seed_output]
 
-    # Manual trigger via button
     run_button.click(
         fn=face_swap,
         inputs=swap_inputs,
         outputs=swap_outputs,
     )
 
-    # Auto-trigger when both images are uploaded
-    def auto_swap_wrapper(ref_face, target_img, s, rand_s, w, h, steps, cfg):
-        """Only run face swap if both images are provided"""
-        if ref_face is not None and target_img is not None:
-            result = face_swap(ref_face, target_img, s, rand_s, w, h, steps, cfg)
-            # Show the button after first generation
+    def auto_swap_wrapper(ref_face, targ_vid, s, rand_s, w, h, steps, cfg):
+        if ref_face is not None and targ_vid is not None:
+            result = face_swap(ref_face, targ_vid, s, rand_s, w, h, steps, cfg)
             return result[0], result[1], gr.update(visible=True)
         return None, s, gr.update(visible=False)
 
-    # Trigger on reference face upload/change
     reference_face.change(
         fn=auto_swap_wrapper,
         inputs=swap_inputs,
         outputs=[comparison_slider, seed_output, run_button],
     )
 
-    # Trigger on target image upload/change
-    target_image.change(
+    # CHANGED: trigger on target video upload/change
+    target_video.change(
         fn=auto_swap_wrapper,
         inputs=swap_inputs,
         outputs=[comparison_slider, seed_output, run_button],