<script>window.huggingface={variables:{"SPACE_CREATOR_USER_ID":"6aa804219e37b3cab39ab28f"}};</script><meta property="og:title" content="Golden Star Suit System2"><meta property="og:description" content="Automatic student head-to-suit template compositing"><meta property="og:image" content="https://cdn-thumbnails.huggingface.co/social-thumbnails/spaces/sheerab/golden-star-suit-system2.png"><meta name="twitter:card" content="summary_large_image"><meta name="twitter:image" content="https://cdn-thumbnails.huggingface.co/social-thumbnails/spaces/sheerab/golden-star-suit-system2.png">import os
import gradio as gr
import numpy as np
import random
import spaces
import torch

from diffusers import Flux2KleinPipeline
from PIL import Image


# ============================================================
# DEVICE / DTYPE
# ============================================================

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

MAX_SEED = np.iinfo(np.int32).max


# ============================================================
# MODELS
# ============================================================

REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-9B"

LORA_REPO_ID = "Alissonerdx/BFS-Best-Face-Swap"

LORA_FILENAME = (
    "bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors"
)


# ============================================================
# FACE-SWAP PROMPT
# ============================================================

FACE_SWAP_PROMPT = """
head_swap:

Picture 1 is the TARGET TEMPLATE.
Picture 2 is the REFERENCE PERSON.

Use Picture 1 as the absolute base composition.

Remove the original head and face of the person in Picture 1 completely.
Replace it with the head and face of the person from Picture 2.

STRICTLY PRESERVE FROM PICTURE 1:
- body
- clothing
- suit
- shoulders
- neck position
- body proportions
- pose
- camera perspective
- framing
- composition
- background
- environment
- lighting direction
- shadows
- highlights
- color temperature
- depth of field
- image geometry

STRICTLY PRESERVE FROM PICTURE 2:
- identity
- facial structure
- face shape
- jawline
- chin
- eyes
- eyebrows
- nose
- lips
- ears
- skin tone
- skin texture
- hair
- hairline
- natural facial details

The final head must naturally fit the body in Picture 1.

Match the head scale, neck connection, perspective, lighting,
color, shadows and skin integration to Picture 1.

Do not create a new body.

Do not redesign the clothing.

Do not change the background.

Do not change the person's body.

Do not change the pose.

Do not change the camera framing.

Do not enlarge or shrink the composition.

The result must look like a real professional studio photograph.

Photorealistic.
Natural skin texture.
Natural hair strands.
Sharp facial details.
Seamless head-to-neck transition.
Professional DSLR studio quality.
"""


# ============================================================
# LOAD MODEL
# ============================================================

print("Loading FLUX.2 Klein 9B Distilled model...")

pipe = Flux2KleinPipeline.from_pretrained(
    REPO_ID_DISTILLED,
    torch_dtype=dtype
)

pipe.to(device)

print("Loading Face Swap LoRA...")

pipe.load_lora_weights(
    LORA_REPO_ID,
    weight_name=LORA_FILENAME
)

print("LoRA loaded successfully!")


# ============================================================
# GET MODEL-SAFE DIMENSIONS
# ============================================================

def get_generation_dimensions(target_image):
    """
    Get dimensions used internally by FLUX.

    IMPORTANT:
    These dimensions are ONLY the temporary generation dimensions.

    The final output will ALWAYS be restored to the
    EXACT ORIGINAL TARGET IMAGE DIMENSIONS.
    """

    if target_image is None:
        return 1024, 1024

    original_width, original_height = target_image.size

    aspect_ratio = original_width / original_height

    # Keep FLUX generation within its practical range.
    # Do NOT alter the original target image itself.

    if aspect_ratio >= 1:
        generation_width = 1024
        generation_height = round(
            1024 / aspect_ratio / 8
        ) * 8
    else:
        generation_height = 1024
        generation_width = round(
            1024 * aspect_ratio / 8
        ) * 8

    generation_width = max(
        256,
        min(1024, generation_width)
    )

    generation_height = max(
        256,
        min(1024, generation_height)
    )

    return generation_width, generation_height


# ============================================================
# FINAL SIZE RESTORATION
# ============================================================

def restore_original_size(
    generated_image,
    target_image
):
    """
    Restore the generated image to the EXACT
    dimensions of the original target/template.

    Example:

    Target:
        1200 x 1800

    Generated internally:
        680 x 1024

    Final:
        1200 x 1800
    """

    original_size = target_image.size

    if generated_image.size != original_size:

        generated_image = generated_image.resize(
            original_size,
            Image.Resampling.LANCZOS
        )

    return generated_image


# ============================================================
# FACE SWAP
# ============================================================

@spaces.GPU(duration=45)
def face_swap(
    reference_face: Image.Image,
    target_image: Image.Image,
    seed: int = 42,
    randomize_seed: bool = False,
    num_inference_steps: int = 4,
    guidance_scale: float = 1.0,
    progress=gr.Progress(track_tqdm=True)
):

    if reference_face is None:
        raise gr.Error(
            "Please upload the student/person image."
        )

    if target_image is None:
        raise gr.Error(
            "Please upload the 10×15 template image."
        )

    # --------------------------------------------------------
    # SAVE ORIGINAL TEMPLATE DIMENSIONS
    # --------------------------------------------------------

    original_width, original_height = target_image.size

    original_size = (
        original_width,
        original_height
    )

    print(
        f"Original template size: "
        f"{original_width} x {original_height}"
    )

    # --------------------------------------------------------
    # SEED
    # --------------------------------------------------------

    if randomize_seed:
        seed = random.randint(
            0,
            MAX_SEED
        )

    generator = torch.Generator(
        device=device
    ).manual_seed(seed)

    # --------------------------------------------------------
    # INTERNAL GENERATION SIZE
    # --------------------------------------------------------

    generation_width, generation_height = (
        get_generation_dimensions(target_image)
    )

    print(
        f"Internal generation size: "
        f"{generation_width} x {generation_height}"
    )

    # --------------------------------------------------------
    # IMPORTANT:
    #
    # Picture 1 = TEMPLATE
    # Picture 2 = STUDENT
    # --------------------------------------------------------

    image_list = [
        target_image,
        reference_face
    ]

    progress(
        0.15,
        desc="Preparing template..."
    )

    # --------------------------------------------------------
    # GENERATION
    # --------------------------------------------------------

    progress(
        0.25,
        desc="Performing AI head replacement..."
    )

    result = pipe(
        prompt=FACE_SWAP_PROMPT,

        image=image_list,

        width=generation_width,
        height=generation_height,

        num_inference_steps=num_inference_steps,

        guidance_scale=guidance_scale,

        generator=generator,
    )

    generated_image = result.images[0]

    progress(
        0.80,
        desc="Restoring original template dimensions..."
    )

    # --------------------------------------------------------
    # RESTORE EXACT ORIGINAL TEMPLATE SIZE
    # --------------------------------------------------------

    final_image = restore_original_size(
        generated_image,
        target_image
    )

    # --------------------------------------------------------
    # HARD SIZE CHECK
    # --------------------------------------------------------

    if final_image.size != original_size:

        raise gr.Error(
            "Final image size mismatch. "
            "The output was not allowed to change the "
            "original template dimensions."
        )

    print(
        f"Final output size: "
        f"{final_image.width} x {final_image.height}"
    )

    progress(
        1.0,
        desc="Completed."
    )

    return (
        (target_image, final_image),
        seed
    )


# ============================================================
# UI
# ============================================================

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(
            """
# Golden Star Studio — AI Face Swap

### Template + Student → Professional 10×15 Portrait

The output keeps the exact original dimensions of
the uploaded template.
"""
        )

        with gr.Row():

            with gr.Column():

                with gr.Row():

                    reference_face = gr.Image(
                        label="Student / Reference Person",
                        type="pil",
                        sources=["upload"],
                        elem_classes="image-container"
                    )

                    target_image = gr.Image(
                        label="10×15 Template",
                        type="pil",
                        sources=["upload"],
                        elem_classes="image-container"
                    )

                run_button = gr.Button(
                    "Generate",
                    variant="primary"
                )

                with gr.Accordion(
                    "Advanced Settings",
                    open=False
                ):

                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42
                    )

                    randomize_seed = gr.Checkbox(
                        label="Randomize seed",
                        value=True
                    )

                    with gr.Row():

                        num_inference_steps = gr.Slider(
                            label="Inference Steps",
                            minimum=1,
                            maximum=20,
                            step=1,
                            value=4,
                            info=(
                                "4 is optimized for the "
                                "distilled model."
                            )
                        )

                        guidance_scale = gr.Slider(
                            label="Guidance Scale",
                            minimum=0.0,
                            maximum=5.0,
                            step=0.1,
                            value=1.0
                        )

            comparison_slider = gr.ImageSlider(
                label="Before / After",
                type="pil"
            )

        seed_output = gr.Number(
            label="Seed Used",
            visible=False
        )


    # ========================================================
    # INPUTS
    # ========================================================

    swap_inputs = [
        reference_face,
        target_image,
        seed,
        randomize_seed,
        num_inference_steps,
        guidance_scale
    ]

    swap_outputs = [
        comparison_slider,
        seed_output
    ]


    # ========================================================
    # BUTTON
    # ========================================================

    run_button.click(
        fn=face_swap,
        inputs=swap_inputs,
        outputs=swap_outputs
    )


# ============================================================
# LAUNCH
# ============================================================

if __name__ == "__main__":

    demo.launch(
        share=True,
        theme=gr.themes.Citrus()
    )