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)