Spaces:
Running
on
Zero
Running
on
Zero
| 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 | |
| import functools | |
| # --- 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 | |
| # def on_gpu_configurable(duration=60): | |
| # def decorator(func): | |
| # @functools.wraps(func) | |
| # @spaces.GPU(duration=duration) | |
| # def wrapper(*args, **kwargs): | |
| # return func(*args, **kwargs) | |
| # return wrapper | |
| # return decorator | |
| 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) | |
| # @on_gpu_configurable(duration=59) | |
| 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=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") | |
| 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=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) | |
| def gpu_tab1(*args, **kwargs): | |
| return app.blind_upscale(*args, **kwargs) | |
| def gpu_tab2(*args, **kwargs): | |
| return app.playground_upscale(*args, **kwargs) | |
| # --- 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) |