Create app.py

app.py

@@ -0,0 +1,388 @@
+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 # Needed for downloading models
+from tqdm import tqdm # For download progress bar
+import spaces
+
+# --- New Official Implementation Imports ---
+from stablepy import load_upscaler_model
+
+# --- New Global Constants ---
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DIRECTORY_UPSCALERS = "upscalers"
+
+# --- Configuration ---
+# Set your Hugging Face Write Token as an environment variable
+# export HF_TOKEN_ORG="hf_YourTokenHere"
+HF_TOKEN_ORG = os.getenv("HF_TOKEN_ORG")
+DATASET_REPO_ID = "TestOrganizationPleaseIgnore/test"
+DATASET_FILENAME = "upscaler_preferences.csv"
+LOCAL_CSV_PATH = "upscaler_preferences_local.csv" # Local backup for safety
+PUSH_THRESHOLD = 10 # Push after this many new votes
+
+# This dictionary remains as a global constant as it's static configuration
+UPSCALER_DICT_GUI = {
+    "RealESRGAN_x4plus": "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x4.pth",
+    "RealESRGAN_x2plus": "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x2.pth",
+    "SwinIR_x4": "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/001_classicalSR_DIV2K_s48w8_SwinIR-M_x4.pth",
+    "BSRGAN_x2": "https://huggingface.co/glassful/models/resolve/main/BSRGANx2.pth",
+    "NewModel_x4_beta": "path/to/new_model.pth" # Example of a local model
+}
+
+# --- 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
+
+def extract_exif_data(image):
+    """Placeholder function to extract EXIF data. Can be expanded later."""
+    # In a real implementation, you would use a library like piexif
+    # and return the exif bytes. For now, it does nothing.
+    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
+
+        elif not os.path.exists(model_path):
+            gr.Warning(f"Local model file not found at: {model_path}")
+            return None
+
+
+        # 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)
+
+    @spaces.GPU()
+    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=True):
+                with gr.Row():
+                    upscaler_size_slider = gr.Slider(minimum=1.1, maximum=4.0, value=2.0, step=0.1, label="Upscale Factor")
+                    tile_slider = gr.Slider(minimum=0, maximum=1024, value=192, step=16, label="Tile Size (0 is not tile)")
+                    tile_overlap_slider = gr.Slider(minimum=0, maximum=128, value=8, step=1, label="Tile Overlap")
+                    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 so we can rank them."
+                )
+                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)
+                    output_image_b = gr.Image(label="Result B", interactive=False)
+                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=self.blind_upscale,
+                    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")
+                        upscaler_model_dropdown = gr.Dropdown(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)
+
+                run_button_playground.click(
+                    fn=self.playground_upscale,
+                    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
+def dummy_gpu():
+    return None
+
+# --- Main Execution Block ---
+if __name__ == "__main__":
+    # Before launching, ensure the upscalers directory exists
+    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)