import os import gradio as gr import numpy as np import random import spaces import torch from diffusers import Flux2KleinPipeline from PIL import Image 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" 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): - Scene: lighting conditions, shadows, highlights, color temperature, environment, background - Head positioning: exact rotation angle, tilt, direction the head is facing - Expression: facial expression, micro-expressions, eye gaze direction, mouth position, emotion FROM PICTURE 2 (strictly preserve identity): - Facial structure: face shape, bone structure, jawline, chin - All facial features: eye color, eye shape, nose structure, lip shape and fullness, eyebrows - Hair: color, style, texture, hairline - Skin: texture, tone, complexion The replaced head must seamlessly match Picture 1's lighting and expression while maintaining the complete identity from Picture 2. High quality, photorealistic, sharp details, 4k.""" print("Loading FLUX.2 Klein 9B Distilled model...") pipe = Flux2KleinPipeline.from_pretrained( REPO_ID_DISTILLED, torch_dtype=dtype, ) pipe.to(device) 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): """ Update width and height based on the target image aspect ratio. """ if target_image is None: return 1024, 1024 img_width, img_height = target_image.size aspect_ratio = img_width / img_height if aspect_ratio >= 1: new_width = 1024 new_height = int(1024 / aspect_ratio) else: new_height = 1024 new_width = int(1024 * aspect_ratio) new_width = round(new_width / 8) * 8 new_height = round(new_height / 8) * 8 new_width = max(256, min(1024, new_width)) new_height = max(256, min(1024, new_height)) return new_width, new_height @spaces.GPU(duration=85) def face_swap( reference_face: Image.Image, target_image: Image.Image, additional_prompt: str = "", seed: int = 42, randomize_seed: bool = False, width: int = 1024, height: int = 1024, num_inference_steps: int = 4, guidance_scale: float = 1.0, progress=gr.Progress(track_tqdm=True), ): if reference_face is None or target_image is None: raise gr.Error( "Please provide both a reference face and a target image!" ) if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator(device=device).manual_seed(seed) image_list = [ target_image, reference_face, ] final_prompt = FACE_SWAP_PROMPT if additional_prompt.strip(): final_prompt += ( f"\n\nADDITIONAL DETAILS:\n" f"{additional_prompt.strip()}" ) progress(0.2, desc="Swapping face...") image = pipe( prompt=final_prompt, image=image_list, height=height, width=width, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator, ).images[0] return (target_image, image), seed css = """ #col-container { margin: 0 auto; max-width: 1200px; } .image-container img { object-fit: contain; } """ with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """# Face Swap with FLUX.2 Klein 9B Swap faces using Flux.2 Klein 9B [Alissonerdx/BFS-Best-Face-Swap](https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap) LoRA """ ) with gr.Row(): reference_face = gr.Image( label="Reference Face", type="pil", sources=[ "upload", "clipboard", ], elem_classes="image-container", ) target_image = gr.Image( label="Target Image (Body/Scene)", type="pil", sources=[ "upload", "clipboard", ], elem_classes="image-container", ) additional_prompt = gr.Textbox( label="Additional Prompt (Optional)", placeholder=( "Add extra details here. " "Press Enter to generate, Shift+Enter for a new line." ), lines=2, max_lines=6, ) run_button = gr.Button( "Generate Face Swap", variant="primary", ) with gr.Accordion("Advanced Settings", open=False): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, ) randomize_seed = gr.Checkbox( label="Randomize seed", value=True, ) with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=1024, step=8, value=1024, ) height = gr.Slider( label="Height", minimum=256, maximum=1024, step=8, value=1024, ) with gr.Row(): num_inference_steps = gr.Slider( label="Inference Steps", minimum=1, maximum=20, step=1, value=4, info=( "Number of denoising steps " "(4 is optimal for distilled model)" ), ) guidance_scale = gr.Slider( label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=1.0, 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, ) # Update dimensions when the target image is uploaded, # pasted from the clipboard, or dropped into the component. target_image.input( fn=update_dimensions_from_image, inputs=[target_image], outputs=[width, height], queue=False, ) swap_inputs = [ reference_face, target_image, additional_prompt, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, ] swap_outputs = [ comparison_slider, seed_output, ] # Generate through the button. run_button.click( fn=face_swap, inputs=swap_inputs, outputs=swap_outputs, ) # Enter generates the image. # Shift+Enter inserts a new line in the textbox. additional_prompt.submit( fn=face_swap, inputs=swap_inputs, outputs=swap_outputs, ) if __name__ == "__main__": demo.launch( share=True, theme=gr.themes.Citrus(), mcp_server=True, )