app.py

import torch
import gradio as gr
from diffusers import FluxPipeline, FluxTransformer2DModel
import gc
import random
import glob
from pathlib import Path
from PIL import Image
import os
import time
import json
from fasteners import InterProcessLock
import spaces

AGG_FILE = Path(__file__).parent / "agg_stats.json"
LOCK_FILE = AGG_FILE.with_suffix(".lock")

def _load_agg_stats() -> dict:
    if AGG_FILE.exists():
        with open(AGG_FILE, "r") as f:
            try:
                return json.load(f)
            except json.JSONDecodeError:
                print(f"Warning: {AGG_FILE} is corrupted. Starting with empty stats.")
                return {"8-bit": {"attempts": 0, "correct": 0}, "4-bit": {"attempts": 0, "correct": 0}}
    return {"8-bit": {"attempts": 0, "correct": 0},
            "4-bit": {"attempts": 0, "correct": 0}}

def _save_agg_stats(stats: dict) -> None:
    with InterProcessLock(str(LOCK_FILE)):
        with open(AGG_FILE, "w") as f:
            json.dump(stats, f, indent=2)

USER_STATS_FILE = Path(__file__).parent / "user_stats.json"
USER_STATS_LOCK_FILE = USER_STATS_FILE.with_suffix(".lock")

def _load_user_stats() -> dict:
    if USER_STATS_FILE.exists():
        with open(USER_STATS_FILE, "r") as f:
            try:
                return json.load(f)
            except json.JSONDecodeError:
                print(f"Warning: {USER_STATS_FILE} is corrupted. Starting with empty user stats.")
                return {}
    return {}

def _save_user_stats(stats: dict) -> None:
    with InterProcessLock(str(USER_STATS_LOCK_FILE)):
        with open(USER_STATS_FILE, "w") as f:
            json.dump(stats, f, indent=2)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")

DEFAULT_HEIGHT = 1024
DEFAULT_WIDTH = 1024
DEFAULT_GUIDANCE_SCALE = 3.5
DEFAULT_NUM_INFERENCE_STEPS = 15
DEFAULT_MAX_SEQUENCE_LENGTH = 512
HF_TOKEN = os.environ.get("HF_ACCESS_TOKEN")

CACHED_PIPES = {}
def load_bf16_pipeline():
    print("Loading BF16 pipeline...")
    MODEL_ID = "black-forest-labs/FLUX.1-dev"
    if MODEL_ID in CACHED_PIPES:
        return CACHED_PIPES[MODEL_ID]
    start_time = time.time()
    try:
        pipe = FluxPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16,
                token=HF_TOKEN
        )
        pipe.to(DEVICE)
        end_time = time.time()
        mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
        print(f"BF16 Pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")
        CACHED_PIPES[MODEL_ID] = pipe
        return pipe
    except Exception as e:
        print(f"Error loading BF16 pipeline: {e}")
        raise

def load_bnb_8bit_pipeline():
    print("Loading 8-bit BNB pipeline...")
    MODEL_ID = "derekl35/FLUX.1-dev-bnb-8bit"
    if MODEL_ID in CACHED_PIPES:
        return CACHED_PIPES[MODEL_ID]
    start_time = time.time()
    try:
        pipe = FluxPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16
        )
        pipe.to(DEVICE)
        # pipe.enable_model_cpu_offload()
        end_time = time.time()
        mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
        print(f"8-bit BNB pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")
        CACHED_PIPES[MODEL_ID] = pipe
        return pipe
    except Exception as e:
        print(f"Error loading 8-bit BNB pipeline: {e}")
        raise

def load_bnb_4bit_pipeline():
    print("Loading 4-bit BNB pipeline...")
    MODEL_ID = "derekl35/FLUX.1-dev-nf4"
    if MODEL_ID in CACHED_PIPES:
        return CACHED_PIPES[MODEL_ID]
    start_time = time.time()
    try:
        pipe = FluxPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16
        )
        pipe.to(DEVICE)
        # pipe.enable_model_cpu_offload()
        end_time = time.time()
        mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
        print(f"4-bit BNB pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")
        CACHED_PIPES[MODEL_ID] = pipe
        return pipe
    except Exception as e:
        print(f"Error loading 4-bit BNB pipeline: {e}")
        raise

@spaces.GPU(duration=240)
def generate_images(prompt, quantization_choice, progress=gr.Progress(track_tqdm=True)):
    if not prompt:
        return None, {}, gr.update(value="Please enter a prompt.", interactive=False), gr.update(choices=[], value=None), gr.update(interactive=True), gr.update(interactive=True)

    if not quantization_choice:
        return None, {}, gr.update(value="Please select a quantization method.", interactive=False), gr.update(choices=[], value=None), gr.update(interactive=True), gr.update(interactive=True)

    if quantization_choice == "8-bit":
        quantized_load_func = load_bnb_8bit_pipeline
        quantized_label = "Quantized (8-bit)"
    elif quantization_choice == "4-bit":
        quantized_load_func = load_bnb_4bit_pipeline
        quantized_label = "Quantized (4-bit)"
    else:
        return None, {}, gr.update(value="Invalid quantization choice.", interactive=False), gr.update(choices=[], value=None), gr.update(interactive=True), gr.update(interactive=True)

    model_configs = [
        ("Original", load_bf16_pipeline),
        (quantized_label, quantized_load_func),
    ]

    results = []
    pipe_kwargs = {
        "prompt": prompt,
        "height": DEFAULT_HEIGHT,
        "width": DEFAULT_WIDTH,
        "guidance_scale": DEFAULT_GUIDANCE_SCALE,
        "num_inference_steps": DEFAULT_NUM_INFERENCE_STEPS,
        "max_sequence_length": DEFAULT_MAX_SEQUENCE_LENGTH,
    }

    seed = random.getrandbits(64)
    print(f"Using seed: {seed}")

    for i, (label, load_func) in enumerate(model_configs):
        progress(i / len(model_configs), desc=f"Loading {label} model...")
        print(f"\n--- Loading {label} Model ---")
        load_start_time = time.time()
        try:
            current_pipe = load_func()
            load_end_time = time.time()
            print(f"{label} model loaded in {load_end_time - load_start_time:.2f} seconds.")

            progress((i + 0.5) / len(model_configs), desc=f"Generating with {label} model...")
            print(f"--- Generating with {label} Model ---")
            gen_start_time = time.time()
            image_list = current_pipe(**pipe_kwargs, generator=torch.manual_seed(seed)).images
            image = image_list[0]
            gen_end_time = time.time()
            results.append({"label": label, "image": image})
            print(f"--- Finished Generation with {label} Model in {gen_end_time - gen_start_time:.2f} seconds ---")
            mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
            print(f"Memory reserved: {mem_reserved:.2f} GB")

        except Exception as e:
            print(f"Error during {label} model processing: {e}")
            return None, {}, gr.update(value=f"Error processing {label} model: {e}", interactive=False), gr.update(choices=[], value=None), gr.update(interactive=True), gr.update(interactive=True)


    if len(results) != len(model_configs):
        return None, {}, gr.update(value="Failed to generate images for all model types.", interactive=False), gr.update(choices=[], value=None), gr.update(interactive=True), gr.update(interactive=True)

    shuffled_results = results.copy()
    random.shuffle(shuffled_results)
    shuffled_data_for_gallery = [(res["image"], f"Image {i+1}") for i, res in enumerate(shuffled_results)]
    correct_mapping = {i: res["label"] for i, res in enumerate(shuffled_results)}
    print("Correct mapping (hidden):", correct_mapping)

    return shuffled_data_for_gallery, correct_mapping, "Generation complete! Make your guess.", None, gr.update(interactive=True), gr.update(interactive=True)


def check_guess(user_guess, correct_mapping_state):
    if not isinstance(correct_mapping_state, dict) or not correct_mapping_state:
        return "Please generate images first (state is empty or invalid)."
    if user_guess is None:
        return "Please select which image you think is quantized."

    quantized_image_index = -1
    quantized_label_actual = ""
    for index, label in correct_mapping_state.items():
        if "Quantized" in label:
            quantized_image_index = index
            quantized_label_actual = label
            break
    if quantized_image_index == -1:
        return "Error: Could not find the quantized image in the mapping data."

    correct_guess_label = f"Image {quantized_image_index + 1}"
    if user_guess == correct_guess_label:
        feedback = f"Correct! {correct_guess_label} used the {quantized_label_actual} model."
    else:
        feedback = f"Incorrect. The quantized image ({quantized_label_actual}) was {correct_guess_label}."
    return feedback

EXAMPLE_DIR = Path(__file__).parent / "examples"
EXAMPLES = [
    {
        "prompt": "A photorealistic portrait of an astronaut on Mars",
        "files": ["astronauts_seed_6456306350371904162.png", "astronauts_bnb_8bit.png"],
        "quantized_idx": 1,
        "quant_method": "bnb 8-bit",
    },
    {
        "prompt": "Water-color painting of a cat wearing sunglasses",
        "files": ["watercolor_cat_bnb_8bit.png", "watercolor_cat_seed_14269059182221286790.png"],
        "quantized_idx": 0,
        "quant_method": "bnb 8-bit",
    },
    # {
    #     "prompt": "Neo-tokyo cyberpunk cityscape at night, rain-soaked streets, 8-K",
    #     "files": ["cyber_city_q.jpg", "cyber_city.jpg"],
    #     "quantized_idx": 0,
    # },
]

def load_example(idx):
    ex = EXAMPLES[idx]
    imgs = [Image.open(EXAMPLE_DIR / f) for f in ex["files"]]
    gallery_items = [(img, f"Image {i+1}") for i, img in enumerate(imgs)]
    mapping = {i: (f"Quantized {ex['quant_method']}" if i == ex["quantized_idx"] else "Original")
               for i in range(2)}
    return gallery_items, mapping, f"{ex['prompt']}"

def _accuracy_string(correct: int, attempts: int) -> tuple[str, float]:
    if attempts:
        pct = 100 * correct / attempts
        return f"{pct:.1f}%", pct
    return "N/A", -1.0

def _add_medals(user_rows):
    MEDALS = {0: "🥇 ", 1: "🥈 ", 2: "🥉 "}
    return [
        [MEDALS.get(i, "") + row[0], *row[1:]]
        for i, row in enumerate(user_rows)
    ]

def update_leaderboards_data():
    agg = _load_agg_stats()
    quant_rows = []
    for method, stats in agg.items():
        acc_str, acc_val = _accuracy_string(stats["correct"], stats["attempts"])
        quant_rows.append([
            method,
            stats["correct"],
            stats["attempts"],
            acc_str
        ])
    quant_rows.sort(key=lambda r: r[1]/r[2] if r[2] != 0 else 1e9)

    user_stats = _load_user_stats()
    user_rows = []
    for user, st in user_stats.items():
        acc_str, acc_val = _accuracy_string(st["total_correct"], st["total_attempts"])
        user_rows.append([user, st["total_correct"], st["total_attempts"], acc_str])
    user_rows.sort(key=lambda r: (-float(r[3].rstrip('%')) if r[3] != "N/A" else float('-inf'), -r[2]))
    user_rows = _add_medals(user_rows)

    return quant_rows, user_rows

with gr.Blocks(title="FLUX Quantization Challenge", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# FLUX Model Quantization Challenge")
    with gr.Tabs():
        with gr.TabItem("Challenge"):
            gr.Markdown(
                "Compare the original FLUX.1-dev (BF16) model against a quantized version (4-bit or 8-bit). "
                "Enter a prompt, choose the quantization method, and generate two images. "
                "The images will be shuffled, can you spot which one was quantized?"
            )

            gr.Markdown("### Examples")
            ex_selector = gr.Radio(
                choices=[f"Example {i+1}" for i in range(len(EXAMPLES))],
                label="Choose an example prompt",
                interactive=True,
            )
            gr.Markdown("### …or create your own comparison")
            with gr.Row():
                prompt_input = gr.Textbox(label="Enter Prompt", scale=3)
                quantization_choice_radio = gr.Radio(
                    choices=["8-bit", "4-bit"],
                    label="Select Quantization",
                    value="8-bit",
                    scale=1
                )
                generate_button = gr.Button("Generate & Compare", variant="primary", scale=1)

            output_gallery = gr.Gallery(
                label="Generated Images",
                columns=2,
                height=606,
                object_fit="contain",
                allow_preview=True,
                show_label=True,
            )

            gr.Markdown("### Which image used the selected quantization method?")
            with gr.Row():
                image1_btn = gr.Button("Image 1")
                image2_btn = gr.Button("Image 2")

            feedback_box = gr.Textbox(label="Feedback", interactive=False, lines=1)

            with gr.Row():
                session_score_box  = gr.Textbox(label="Your accuracy this session", interactive=False)

            with gr.Row(equal_height=False):
                username_input = gr.Textbox(
                    label="Enter Your Name for Leaderboard",
                    placeholder="YourName",
                    visible=False,
                    interactive=True,
                    scale=2
                )
                add_score_button = gr.Button(
                    "Add My Score to Leaderboard",
                    visible=False,
                    variant="secondary",
                    scale=1
                )
            add_score_feedback = gr.Textbox(
                label="Leaderboard Update",
                visible=False,
                interactive=False,
                lines=1
            )

            correct_mapping_state = gr.State({})
            session_stats_state = gr.State(
                {"8-bit": {"attempts": 0, "correct": 0},
                "4-bit": {"attempts": 0, "correct": 0}}
            )
            is_example_state = gr.State(False)
            has_added_score_state = gr.State(False)

            def _load_example(sel):
                idx = int(sel.split()[-1]) - 1
                gallery_items, mapping, prompt = load_example(idx)
                quant_data, user_data = update_leaderboards_data()
                return gallery_items, mapping, prompt, True, quant_data, user_data

            ex_selector.change(
                fn=_load_example,
                inputs=ex_selector,
                outputs=[output_gallery, correct_mapping_state, prompt_input, is_example_state, quant_df, user_df],
            ).then(
                lambda: (gr.update(interactive=True), gr.update(interactive=True)),
                outputs=[image1_btn, image2_btn],
            )

            generate_button.click(
                fn=generate_images,
                inputs=[prompt_input, quantization_choice_radio],
                outputs=[output_gallery, correct_mapping_state]
            ).then(
                lambda: (False, # for is_example_state
                         False, # for has_added_score_state
                         gr.update(visible=False, value="", interactive=True), # username_input reset
                         gr.update(visible=False), # add_score_button reset
                         gr.update(visible=False, value="")), # add_score_feedback reset
                outputs=[is_example_state,
                         has_added_score_state,
                         username_input,
                         add_score_button,
                         add_score_feedback]
            ).then(
                lambda: (gr.update(interactive=True),
                        gr.update(interactive=True),
                        ""),
                outputs=[image1_btn, image2_btn, feedback_box],
            )

            def choose(choice_string, mapping, session_stats, is_example, has_added_score_curr):
                feedback = check_guess(choice_string, mapping)

                quant_label = next(label for label in mapping.values() if "Quantized" in label)
                quant_key = "8-bit" if "8-bit" in quant_label else "4-bit"

                got_it_right = "Correct!" in feedback

                sess = session_stats.copy()
                if not is_example and not has_added_score_curr:
                    sess[quant_key]["attempts"] += 1
                    if got_it_right:
                        sess[quant_key]["correct"] += 1
                    session_stats = sess

                    AGG_STATS = _load_agg_stats()
                    AGG_STATS[quant_key]["attempts"] += 1
                    if got_it_right:
                        AGG_STATS[quant_key]["correct"] += 1
                    _save_agg_stats(AGG_STATS)

                def _fmt(d):
                    a, c = d["attempts"], d["correct"]
                    pct = 100 * c / a if a else 0
                    return f"{c} / {a}  ({pct:.1f}%)"

                session_msg = ", ".join(
                    f"{k}: {_fmt(v)}" for k, v in sess.items()
                )
                current_agg_stats = _load_agg_stats()
                global_msg = ", ".join(
                    f"{k}: {_fmt(v)}" for k, v in current_agg_stats.items()
                )

                username_input_update = gr.update(visible=False, interactive=True)
                add_score_button_update = gr.update(visible=False)
                # Keep existing feedback if score was already added and feedback is visible
                current_feedback_text = add_score_feedback.value if hasattr(add_score_feedback, 'value') and add_score_feedback.value else ""
                add_score_feedback_update = gr.update(visible=has_added_score_curr, value=current_feedback_text)

                session_total_attempts = sum(stats["attempts"] for stats in sess.values())

                if not is_example and not has_added_score_curr:
                    if session_total_attempts >= 1 : # Show button if more than 1 attempt
                        username_input_update = gr.update(visible=True, interactive=True)
                        add_score_button_update = gr.update(visible=True, interactive=True)
                        add_score_feedback_update = gr.update(visible=False, value="")
                    else: # Less than 1 attempts, keep hidden
                        username_input_update = gr.update(visible=False, value=username_input.value if hasattr(username_input, 'value') else "")
                        add_score_button_update = gr.update(visible=False)
                        add_score_feedback_update = gr.update(visible=False, value="")
                elif has_added_score_curr:
                    username_input_update = gr.update(visible=True, interactive=False, value=username_input.value if hasattr(username_input, 'value') else "")
                    add_score_button_update = gr.update(visible=True, interactive=False)
                    add_score_feedback_update = gr.update(visible=True)

                # disable the buttons so the user can't vote twice
                quant_data, user_data = update_leaderboards_data() # Get updated leaderboard data
                return (feedback,
                        gr.update(interactive=False),
                        gr.update(interactive=False),
                        session_msg,
                        session_stats,
                        quant_data,
                        user_data,
                        username_input_update,
                        add_score_button_update,
                        add_score_feedback_update)


            image1_btn.click(
                fn=lambda mapping, sess, is_ex, has_added: choose("Image 1", mapping, sess, is_ex, has_added),
                inputs=[correct_mapping_state, session_stats_state, is_example_state, has_added_score_state],
                outputs=[feedback_box, image1_btn, image2_btn,
                        session_score_box, session_stats_state,
                        quant_df, user_df,
                        username_input, add_score_button, add_score_feedback],
            )
            image2_btn.click(
                fn=lambda mapping, sess, is_ex, has_added: choose("Image 2", mapping, sess, is_ex, has_added),
                inputs=[correct_mapping_state, session_stats_state, is_example_state, has_added_score_state],
                outputs=[feedback_box, image1_btn, image2_btn,
                        session_score_box, session_stats_state,
                        quant_df, user_df,
                        username_input, add_score_button, add_score_feedback],
            )

            def handle_add_score_to_leaderboard(username_str, current_session_stats_dict):
                if not username_str or not username_str.strip():
                    return ("Username is required.",  # Feedback for add_score_feedback
                            gr.update(interactive=True),    # username_input
                            gr.update(interactive=True),    # add_score_button
                            False,                          # has_added_score_state
                            None, None)                     # quant_df, user_df

                user_stats = _load_user_stats()
                user_key = username_str.strip()

                session_total_correct = sum(stats["correct"] for stats in current_session_stats_dict.values())
                session_total_attempts = sum(stats["attempts"] for stats in current_session_stats_dict.values())

                if session_total_attempts == 0:
                     return ("No attempts made in this session to add to leaderboard.",
                            gr.update(interactive=True),
                            gr.update(interactive=True),
                            False, None, None)

                if user_key in user_stats:
                    user_stats[user_key]["total_correct"] += session_total_correct
                    user_stats[user_key]["total_attempts"] += session_total_attempts
                else:
                    user_stats[user_key] = {
                        "total_correct": session_total_correct,
                        "total_attempts": session_total_attempts
                    }
                _save_user_stats(user_stats)
                
                new_quant_data, new_user_data = update_leaderboards_data()
                feedback_msg = f"Score for '{user_key}' submitted to leaderboard!"
                return (feedback_msg,                   # To add_score_feedback
                        gr.update(interactive=False),   # username_input
                        gr.update(interactive=False),   # add_score_button
                        True,                           # has_added_score_state (set to true)
                        new_quant_data,                 # To quant_df
                        new_user_data)                  # To user_df

            add_score_button.click(
                fn=handle_add_score_to_leaderboard,
                inputs=[username_input, session_stats_state],
                outputs=[add_score_feedback, username_input, add_score_button, has_added_score_state, quant_df, user_df]
            )
        with gr.TabItem("Leaderboard"):
            gr.Markdown("## Quantization Method Leaderboard  *(Lower % ⇒ harder to detect)*")
            quant_df = gr.DataFrame(
                headers=["Method", "Correct Guesses", "Total Attempts", "Detectability %"],
                interactive=False, col_count=(4, "fixed")
            )
            gr.Markdown("## User Leaderboard  *(Higher % ⇒ better spotter)*")
            user_df = gr.DataFrame(
                headers=["User", "Correct Guesses", "Total Attempts", "Accuracy %"],
                interactive=False, col_count=(4, "fixed")
            )
            demo.load(update_leaderboards_data, outputs=[quant_df, user_df])

if __name__ == "__main__":
    demo.launch(share=True)