app.py

import gradio as gr
import pandas as pd
import random
import os
from PIL import Image
from huggingface_hub import HfApi
from io import StringIO
from transformers import pipeline
import torch
import requests
from tqdm import tqdm
import spaces
import functools
from constants import updated_upscaler_dict as UPSCALER_DICT_GUI
from stablepy import load_upscaler_model

# --- New Global Constants ---
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DIRECTORY_UPSCALERS = "upscalers"

# --- Configuration ---
# export HF_TOKEN_ORG="hf_YourTokenHere"
HF_TOKEN_ORG = os.getenv("HF_TOKEN_ORG")
DATASET_REPO_ID = "TestOrganizationPleaseIgnore/upscale_board_data"
DATASET_FILENAME = "upscaler_preferences.csv"
LOCAL_CSV_PATH = "upscaler_preferences_local.csv" # Local backup for safety
PUSH_THRESHOLD = 5 # Push after this many new votes

# --- Helper Functions for New Implementation ---
def download_model(directory, url):
    """Downloads a file from a URL to a specified directory with a progress bar."""
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"Created directory: {directory}")

    filename = url.split('/')[-1]
    filepath = os.path.join(directory, filename)

    if os.path.exists(filepath):
        print(f"Model '{filename}' already exists. Skipping download.")
        return filepath

    try:
        print(f"Downloading model '{filename}' from {url}...")
        response = requests.get(url, stream=True)
        response.raise_for_status()  # Raise an exception for bad status codes
        total_size_in_bytes = int(response.headers.get('content-length', 0))
        block_size = 1024  # 1 Kibibyte

        with tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True, desc=f"Downloading {filename}") as progress_bar:
            with open(filepath, 'wb') as file:
                for data in response.iter_content(block_size):
                    progress_bar.update(len(data))
                    file.write(data)

        if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
            print("ERROR, something went wrong during download.")
            return None

        print(f"Model '{filename}' downloaded successfully to '{filepath}'.")
        return filepath
    except requests.exceptions.RequestException as e:
        print(f"Error downloading model: {e}")
        return None


class UpscalerApp:
    def __init__(self, repo_id, filename, local_path, push_threshold):
        """
        Initializes the application, loads data, and sets up state.
        """
        self.repo_id = repo_id
        self.filename = filename
        self.local_path = local_path
        self.push_threshold = push_threshold

        self.results_df = None
        self.new_votes_count = 0

        # Initialize the image classifier on the correct device (GPU or CPU)
        print(f"Initializing classifier on device: {DEVICE}")
        self.classifier = pipeline(
            "zero-shot-image-classification",
            model="laion/CLIP-ViT-L-14-laion2B-s32B-b82K",
            device=DEVICE
        )
        self.disambiguation_dict = {
            "Modern photo or photorealistic CGI": "modern_photo_cgi",
            "Old vintage photograph": "vintage_photo",
            "Anime illustration": "anime_illustration",
            "Manga": "manga",
            "Cartoon, Comic book": "cartoon_comic",
            "In-game screenshot with heads-up display HUD or UI elements": "in_game_screenshot_hud",
            "Pixel art or low-resolution retro graphics": "pixel_art_retro",
            "Text document or code": "text_document_code"
        }
        self.candidate_labels = list(self.disambiguation_dict.keys())

        self.initialize_dataset()
        self.ui = self.build_gradio_ui()

    def initialize_dataset(self):
        """
        Loads the dataset from the Hub, falling back to a local file,
        and finally creating a new one if necessary.
        """
        if HF_TOKEN_ORG is None:
            print("WARNING: HF_TOKEN_ORG not set. Results will only be saved locally.")

        # 1. Try to load from Hugging Face Hub first
        try:
            api = HfApi()
            file_path = api.hf_hub_download(
                repo_id=self.repo_id,
                filename=self.filename,
                repo_type="dataset",
                token=HF_TOKEN_ORG
            )
            self.results_df = pd.read_csv(file_path).set_index("model_name")
            print(f"Successfully loaded results from '{self.repo_id}'.")
        except Exception as e:
            print(f"Could not load from Hub (may not exist yet): {e}")
            # 2. If Hub fails, try to load from local backup
            if os.path.exists(self.local_path):
                print(f"Loading results from local file: '{self.local_path}'")
                self.results_df = pd.read_csv(self.local_path).set_index("model_name")
            else:
                # 3. If no local file, create a new DataFrame
                print("No local CSV found. Creating a new preference count DataFrame.")
                model_names = list(UPSCALER_DICT_GUI.keys())
                columns = ['model_name', 'count'] + list(self.disambiguation_dict.values())
                self.results_df = pd.DataFrame(columns=columns).set_index('model_name')


        # Ensure all current models and columns exist in the DataFrame
        for model in UPSCALER_DICT_GUI:
            if model not in self.results_df.index:
                print(f"Adding new model '{model}' to the DataFrame.")
                self.results_df.loc[model] = 0

        for col in list(self.disambiguation_dict.values()):
            if col not in self.results_df.columns:
                self.results_df[col] = 0

        # Save a clean local copy on startup
        self.save_results_to_local_csv()


    def push_results_to_hub(self):
        """
        Pushes the current results DataFrame to the Hugging Face Hub.
        This is a BLOCKING operation and will freeze the UI.
        """
        if HF_TOKEN_ORG is None:
            print("Skipping push: HF_TOKEN_ORG not available.")
            return

        if self.results_df is None or self.results_df.empty:
            return

        print(f"Blocking UI to push results to '{self.repo_id}'...")
        try:
            csv_buffer = StringIO()
            # reset_index() makes 'model_name' a column again before saving
            self.results_df.reset_index().to_csv(csv_buffer, index=False)

            api = HfApi()
            api.upload_file(
                path_or_fileobj=csv_buffer.getvalue().encode("utf-8"),
                path_in_repo=self.filename,
                repo_id=self.repo_id,
                repo_type="dataset",
                token=HF_TOKEN_ORG,
                commit_message="Automated preference count update"
            )
            print("Successfully pushed updated results to the Hub.")
        except Exception as e:
            print(f"Error pushing results to the Hub: {e}")

    def save_results_to_local_csv(self):
        """Saves the current DataFrame to a local CSV file for persistence."""
        if self.results_df is not None:
            self.results_df.reset_index().to_csv(self.local_path, index=False)

    # --- Official upscale function ---
    def process_upscale(self, image, upscaler_name, upscaler_size, tile, tile_overlap, half):
        """
        Processes an image using the specified upscaler model and settings.
        """
        if image is None:
            return None

        print(f"Upscaling with: {upscaler_name}, Size: {upscaler_size}, Tile: {tile}, Overlap: {tile_overlap}, Half: {half}")

        image = image.convert("RGB")
        # exif_image = extract_exif_data(image) # Placeholder for future use

        model_path = UPSCALER_DICT_GUI[upscaler_name]

        # Check if the model is a URL and download it if it doesn't exist locally
        if "https://" in str(model_path) or "http://" in str(model_path):
            local_model_path = download_model(DIRECTORY_UPSCALERS, model_path)
            if local_model_path is None:
                # Handle download failure
                gr.Warning("Failed to download the upscaler model. Please check the console for errors.")
                return None
            model_path = local_model_path

        # Load the upscaler model with specified tile and precision settings
        scaler_beta = load_upscaler_model(model=model_path, tile=tile, tile_overlap=tile_overlap, device=DEVICE, half=half)

        # Perform the upscale
        image_up = scaler_beta.upscale(image, upscaler_size, True)

        return image_up

    # --- Gradio Callback Functions ---
    def blind_upscale(self, image, upscaler_size, tile, tile_overlap, half):
        if image is None:
            return None, None, "Please upload an image.", "", "", "", gr.Button(interactive=False), gr.Button(interactive=False)

        # Classify the image
        predictions = self.classifier(image, candidate_labels=self.candidate_labels)
        top_prediction_label = predictions[0]['label']
        top_prediction_key = self.disambiguation_dict[top_prediction_label]

        model_keys = list(UPSCALER_DICT_GUI.keys())
        if len(model_keys) < 2:
            return None, None, "Not enough models to compare.", "", "", "", gr.Button(interactive=False), gr.Button(interactive=False)

        model_a_name, model_b_name = random.sample(model_keys, 2)

        # Process both images with the same settings from the UI
        upscaled_a = self.process_upscale(image, model_a_name, upscaler_size, tile, tile_overlap, half)
        upscaled_b = self.process_upscale(image, model_b_name, upscaler_size, tile, tile_overlap, half)

        if upscaled_a is None or upscaled_b is None:
            # Handle case where upscaling failed (e.g., model download error)
             return None, None, "Upscaling failed. Check console for details.", "", "", "", gr.Button(interactive=False), gr.Button(interactive=False)

        result_text = f"Image classified as: **{top_prediction_label}**. Which result do you prefer?"

        return upscaled_a, upscaled_b, result_text, model_a_name, model_b_name, top_prediction_key, gr.Button(interactive=True), gr.Button(interactive=True)

    def handle_choice(self, choice, model_a, model_b, image_category):
        if not model_a or not model_b:
            return "Please start a comparison first.", gr.Button(interactive=False), gr.Button(interactive=False)

        winner = model_a if choice == "Result A" else model_b

        if winner not in self.results_df.index:
            self.results_df.loc[winner] = 0

        # Increment the main count and the category-specific count
        self.results_df.loc[winner, 'count'] += 1
        if image_category in self.results_df.columns:
            self.results_df.loc[winner, image_category] += 1

        new_count = self.results_df.loc[winner, 'count']
        self.new_votes_count += 1

        print(f"Recorded preference for '{winner}' in category '{image_category}'. New count: {new_count}. Total new votes: {self.new_votes_count}")

        # Always save locally for safety
        self.save_results_to_local_csv()

        # If threshold is met, trigger a BLOCKING push
        if self.new_votes_count >= self.push_threshold:
            print(f"Vote threshold reached. Initiating blocking push to Hub...")
            self.push_results_to_hub() # This is a direct, blocking call
            self.new_votes_count = 0 # Reset counter

        reveal_text = f"Thank you! Your preference for **{choice}** has been recorded.\n\n- **Image A was:** {model_a}\n- **Image B was:** {model_b}"
        return reveal_text, gr.Button(interactive=False), gr.Button(interactive=False)

    def playground_upscale(self, image, upscaler_name, upscaler_size, tile, tile_overlap, half):
        if image is None or upscaler_name is None: return None
        return self.process_upscale(image, upscaler_name, upscaler_size, tile, tile_overlap, half)

    def build_gradio_ui(self):
        """Constructs the Gradio interface."""
        with gr.Blocks(theme=gr.themes.Soft()) as demo:
            gr.Markdown("# Image Upscaler GUI with A/B Testing")

            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    upscaler_size_slider = gr.Slider(minimum=1.1, maximum=4.0, value=2.0, step=0.1, label="Upscale Factor")
                with gr.Row():
                    tile_slider = gr.Slider(minimum=0, maximum=1024, value=192, step=16, label="Tile Size (0 is not tile)")
                with gr.Row():
                    tile_overlap_slider = gr.Slider(minimum=0, maximum=128, value=8, step=1, label="Tile Overlap")
                with gr.Row():
                    half_checkbox = gr.Checkbox(label="Use Half Precision (FP16)", value=True)


            with gr.Tab("Blind Test Comparison"):
                gr.Markdown("Upload an image, compare the results, and select your favorite. Your vote is recorded to rank the models.")
                gr.Markdown(
                    "> **Disclaimer:** This application **does not store your uploaded images**. "
                    "It only anonymously records which upscaler you prefer to rank them. "
                    "The collected statistics are publicly available at "
                    "[upscaler_preferences.csv](https://huggingface.co/datasets/TestOrganizationPleaseIgnore/upscale_board_data/blob/main/upscaler_preferences.csv)."
                )
                model_a_state = gr.State("")
                model_b_state = gr.State("")
                image_category_state = gr.State("")
                with gr.Row():
                    input_image_blind = gr.Image(type="pil", label="Source Image")
                compare_button = gr.Button("Compare Upscalers")
                with gr.Row():
                    output_image_a = gr.Image(label="Result A", interactive=False, format="png")
                    output_image_b = gr.Image(label="Result B", interactive=False, format="png")
                with gr.Row():
                    choose_a_button = gr.Button("I prefer Result A", interactive=False)
                    choose_b_button = gr.Button("I prefer Result B", interactive=False)
                result_text_blind = gr.Markdown("")

                compare_button.click(
                    fn=gpu_tab1,
                    inputs=[input_image_blind, upscaler_size_slider, tile_slider, tile_overlap_slider, half_checkbox],
                    outputs=[output_image_a, output_image_b, result_text_blind, model_a_state, model_b_state, image_category_state, choose_a_button, choose_b_button]
                )
                choose_a_button.click(
                    fn=lambda a, b, c: self.handle_choice("Result A", a, b, c),
                    inputs=[model_a_state, model_b_state, image_category_state],
                    outputs=[result_text_blind, choose_a_button, choose_b_button]
                )
                choose_b_button.click(
                    fn=lambda a, b, c: self.handle_choice("Result B", a, b, c),
                    inputs=[model_a_state, model_b_state, image_category_state],
                    outputs=[result_text_blind, choose_a_button, choose_b_button]
                )

            with gr.Tab("Upscaler Playground"):
                gr.Markdown("Select an upscaler model, choose a scaling factor, and process your image.")
                with gr.Row():
                    with gr.Column(scale=1):
                        input_image_playground = gr.Image(type="pil", label="Source Image", format="png")
                        upscaler_model_dropdown = gr.Dropdown(value="R-ESRGAN_4x+", choices=list(UPSCALER_DICT_GUI.keys()), label="Upscaler Model")
                        run_button_playground = gr.Button("Run Upscale")
                    with gr.Column(scale=2):
                        output_image_playground = gr.Image(label="Upscaled Result", interactive=False, format="png")

                run_button_playground.click(
                    fn=gpu_tab2,
                    inputs=[input_image_playground, upscaler_model_dropdown, upscaler_size_slider, tile_slider, tile_overlap_slider, half_checkbox],
                    outputs=[output_image_playground]
                )

        return demo

    def launch(self, **kwargs):
        self.ui.launch(**kwargs)

@spaces.GPU(duration=70)
def gpu_tab1(*args, **kwargs):
    return app.blind_upscale(*args, **kwargs)

@spaces.GPU(duration=60)
def gpu_tab2(*args, **kwargs):
    return app.playground_upscale(*args, **kwargs)

if __name__ == "__main__":
    if not os.path.exists(DIRECTORY_UPSCALERS):
        os.makedirs(DIRECTORY_UPSCALERS)

    app = UpscalerApp(
        repo_id=DATASET_REPO_ID,
        filename=DATASET_FILENAME,
        local_path=LOCAL_CSV_PATH,
        push_threshold=PUSH_THRESHOLD
    )
    app.launch(debug=True, show_error=True)