Upload app.py

app.py

@@ -1,333 +1,32 @@
-# --- Imports ---
 import gradio as gr
-import numpy as np
-from PIL import Image, ImageDraw, ImageFont
-import json
-import os
-import io
-import requests
-import matplotlib.pyplot as plt
-import matplotlib
-from huggingface_hub import hf_hub_download
-from dataclasses import dataclass
-from typing import List, Dict, Optional, Tuple
+import spaces
 import time
-import spaces # Required for @spaces.GPU
-
-import torch
-import timm
-from safetensors.torch import load_file as safe_load_file
-
-# MatplotlibのバックエンドをAggに設定
-matplotlib.use('Agg')
-
-# --- Data Classes and Helper Functions ---
-@dataclass
-class LabelData:
-    names: list[str]
-    rating: list[np.int64]
-    general: list[np.int64]
-    artist: list[np.int64]
-    character: list[np.int64]
-    copyright: list[np.int64]
-    meta: list[np.int64]
-    quality: list[np.int64]
-
-def pil_ensure_rgb(image: Image.Image) -> Image.Image:
-    if image.mode not in ["RGB", "RGBA"]:
-        image = image.convert("RGBA") if "transparency" in image.info else image.convert("RGB")
-    if image.mode == "RGBA":
-        background = Image.new("RGB", image.size, (255, 255, 255))
-        background.paste(image, mask=image.split()[3])
-        image = background
-    return image
-
-def pil_pad_square(image: Image.Image) -> Image.Image:
-    width, height = image.size
-    if width == height: return image
-    new_size = max(width, height)
-    new_image = Image.new("RGB", (new_size, new_size), (255, 255, 255))
-    paste_position = ((new_size - width) // 2, (new_size - height) // 2)
-    new_image.paste(image, paste_position)
-    return new_image
-
-def load_tag_mapping(mapping_path):
-    with open(mapping_path, 'r', encoding='utf-8') as f: tag_mapping_data = json.load(f)
-    if isinstance(tag_mapping_data, dict) and "idx_to_tag" in tag_mapping_data:
-        idx_to_tag = {int(k): v for k, v in tag_mapping_data["idx_to_tag"].items()}
-        tag_to_category = tag_mapping_data["tag_to_category"]
-    elif isinstance(tag_mapping_data, dict):
-        tag_mapping_data = {int(k): v for k, v in tag_mapping_data.items()}
-        idx_to_tag = {idx: data['tag'] for idx, data in tag_mapping_data.items()}
-        tag_to_category = {data['tag']: data['category'] for data in tag_mapping_data.values()}
-    else: raise ValueError("Unsupported tag mapping format")
-    names = [None] * (max(idx_to_tag.keys()) + 1)
-    rating, general, artist, character, copyright, meta, quality = [], [], [], [], [], [], []
-    for idx, tag in idx_to_tag.items():
-        if idx >= len(names): names.extend([None] * (idx - len(names) + 1))
-        names[idx] = tag
-        category = tag_to_category.get(tag, 'Unknown')
-        idx_int = int(idx)
-        if category == 'Rating': rating.append(idx_int)
-        elif category == 'General': general.append(idx_int)
-        elif category == 'Artist': artist.append(idx_int)
-        elif category == 'Character': character.append(idx_int)
-        elif category == 'Copyright': copyright.append(idx_int)
-        elif category == 'Meta': meta.append(idx_int)
-        elif category == 'Quality': quality.append(idx_int)
-    return LabelData(names=names, rating=np.array(rating), general=np.array(general), artist=np.array(artist),
-                     character=np.array(character), copyright=np.array(copyright), meta=np.array(meta), quality=np.array(quality)), tag_to_category
-
-def get_tags(probs, labels: LabelData, gen_threshold, char_threshold):
-    result = {"rating": [], "general": [], "character": [], "copyright": [], "artist": [], "meta": [], "quality": []}
-    if labels.rating.size > 0:
-        valid_indices = labels.rating[labels.rating < len(probs)]
-        if valid_indices.size > 0:
-            rating_probs = probs[valid_indices]
-            if rating_probs.size > 0:
-                rating_idx = np.argmax(rating_probs); original_idx = valid_indices[rating_idx]
-                if original_idx < len(labels.names): result["rating"].append((labels.names[original_idx], float(rating_probs[rating_idx])))
-    if labels.quality.size > 0:
-        valid_indices = labels.quality[labels.quality < len(probs)]
-        if valid_indices.size > 0:
-            quality_probs = probs[valid_indices]
-            if quality_probs.size > 0:
-                quality_idx = np.argmax(quality_probs); original_idx = valid_indices[quality_idx]
-                if original_idx < len(labels.names): result["quality"].append((labels.names[original_idx], float(quality_probs[quality_idx])))
-    category_map = {"general": (labels.general, gen_threshold), "character": (labels.character, char_threshold),
-                    "copyright": (labels.copyright, char_threshold), "artist": (labels.artist, char_threshold), "meta": (labels.meta, gen_threshold)}
-    for category, (indices, threshold) in category_map.items():
-        if indices.size > 0:
-            valid_indices = indices[(indices < len(probs)) & (indices < len(labels.names))]
-            if valid_indices.size > 0:
-                 category_probs = probs[valid_indices]; mask = category_probs >= threshold
-                 selected_indices = valid_indices[mask]; selected_probs = category_probs[mask]
-                 for idx, prob in zip(selected_indices, selected_probs): result[category].append((labels.names[idx], float(prob)))
-    for k in result: result[k] = sorted(result[k], key=lambda x: x[1], reverse=True)
-    return result
-
-def visualize_predictions(image: Image.Image, predictions, threshold=0.45):
-    filtered_meta = [(tag, prob) for tag, prob in predictions.get("meta", []) if not any(p in tag.lower() for p in ['id', 'commentary', 'request', 'mismatch'])]
-    predictions["meta"] = filtered_meta
-    fig = plt.figure(figsize=(20, 12), dpi=100); gs = fig.add_gridspec(1, 2, width_ratios=[1.2, 1])
-    ax_img = fig.add_subplot(gs[0, 0]); ax_img.imshow(image); ax_img.set_title("Original Image"); ax_img.axis('off')
-    ax_tags = fig.add_subplot(gs[0, 1])
-    all_tags, all_probs, all_colors = [], [], []
-    color_map = {'rating': 'red', 'character': 'blue', 'copyright': 'purple', 'artist': 'orange', 'general': 'green', 'meta': 'gray', 'quality': 'yellow'}
-    for cat, prefix, color in [('rating', 'R', 'red'), ('character', 'C', 'blue'), ('copyright', '©', 'purple'), ('artist', 'A', 'orange'), ('general', 'G', 'green'), ('meta', 'M', 'gray'), ('quality', 'Q', 'yellow')]:
-        for tag, prob in predictions.get(cat, []): all_tags.append(f"[{prefix}] {tag}"); all_probs.append(prob); all_colors.append(color)
-    if not all_tags:
-        ax_tags.text(0.5, 0.5, "No tags found", ha='center', va='center'); ax_tags.set_title(f"Tags (threshold={threshold})"); ax_tags.axis('off')
-    else:
-        sorted_indices = sorted(range(len(all_probs)), key=lambda i: all_probs[i]); all_tags = [all_tags[i] for i in sorted_indices]; all_probs = [all_probs[i] for i in sorted_indices]; all_colors = [all_colors[i] for i in sorted_indices]
-        num_tags = len(all_tags); bar_height = min(0.8, 0.8 * (30 / num_tags) if num_tags > 30 else 0.8); y_positions = np.arange(num_tags)
-        bars = ax_tags.barh(y_positions, all_probs, height=bar_height, color=all_colors); ax_tags.set_yticks(y_positions); ax_tags.set_yticklabels(all_tags)
-        fontsize = 10 if num_tags <= 40 else 8 if num_tags <= 60 else 6; [label.set_fontsize(fontsize) for label in ax_tags.get_yticklabels()]
-        for i, (bar, prob) in enumerate(zip(bars, all_probs)): ax_tags.text(min(prob + 0.02, 0.98), y_positions[i], f"{prob:.3f}", va='center', fontsize=fontsize)
-        ax_tags.set_xlim(0, 1); ax_tags.set_title(f"Tags (threshold={threshold})")
-        from matplotlib.patches import Patch; legend_elements = [Patch(facecolor=color, label=cat.capitalize()) for cat, color in color_map.items() if cat in predictions and predictions[cat]]; ax_tags.legend(handles=legend_elements, loc='lower right', fontsize=8)
-    plt.tight_layout(); plt.subplots_adjust(bottom=0.05); buf = io.BytesIO(); plt.savefig(buf, format='png', dpi=100); plt.close(fig); buf.seek(0)
-    return Image.open(buf)
-
-def preprocess_image(image: Image.Image, target_size=(448, 448)):
-    image = pil_ensure_rgb(image)
-    image = pil_pad_square(image)
-    image_resized = image.resize(target_size, Image.BICUBIC)
-    img_array = np.array(image_resized, dtype=np.float32) / 255.0
-    img_array = img_array.transpose(2, 0, 1)
-    img_array = img_array[::-1, :, :] # BGR
-    mean = np.array([0.5, 0.5, 0.5], dtype=np.float32).reshape(3, 1, 1)
-    std = np.array([0.5, 0.5, 0.5], dtype=np.float32).reshape(3, 1, 1)
-    img_array = (img_array - mean) / std
-    img_tensor = torch.from_numpy(img_array).unsqueeze(0)
-    return image, img_tensor
-
-# --- Constants ---
-REPO_ID = "cella110n/cl_tagger"
-SAFETENSORS_FILENAME = "lora_model_0426/checkpoint_epoch_4.safetensors"
-METADATA_FILENAME = "lora_model_0426/checkpoint_epoch_4_metadata.json"
-TAG_MAPPING_FILENAME = "lora_model_0426/tag_mapping.json"
-CACHE_DIR = "./model_cache"
-BASE_MODEL_NAME = 'eva02_large_patch14_448.mim_m38m_ft_in1k' # Define base model name
-
-# --- Tagger Class ---
-class Tagger:
-    def __init__(self):
-        print("Initializing Tagger...")
-        self.safetensors_path = None
-        self.metadata_path = None
-        self.tag_mapping_path = None
-        self.labels_data = None
-        self.tag_to_category = None
-        self.model = None # Model will be loaded on first predict call
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self._initialize_paths_and_labels()
-
-    def _download_files(self):
-        if self.safetensors_path and self.tag_mapping_path and os.path.exists(self.safetensors_path) and os.path.exists(self.tag_mapping_path):
-            print("Files seem already downloaded.")
-            return
-        print("Downloading model files...")
-        hf_token = os.environ.get("HF_TOKEN")
-        try:
-            self.safetensors_path = hf_hub_download(repo_id=REPO_ID, filename=SAFETENSORS_FILENAME, cache_dir=CACHE_DIR, token=hf_token, force_download=True)
-            self.tag_mapping_path = hf_hub_download(repo_id=REPO_ID, filename=TAG_MAPPING_FILENAME, cache_dir=CACHE_DIR, token=hf_token, force_download=True)
-            print(f"Safetensors: {self.safetensors_path}")
-            print(f"Tag mapping: {self.tag_mapping_path}")
-            try:
-                self.metadata_path = hf_hub_download(repo_id=REPO_ID, filename=METADATA_FILENAME, cache_dir=CACHE_DIR, token=hf_token, force_download=True)
-                print(f"Metadata: {self.metadata_path}")
-            except Exception: print(f"Metadata ({METADATA_FILENAME}) not found/download failed."); self.metadata_path = None
-        except Exception as e:
-            print(f"Error downloading files: {e}")
-            if "401 Client Error" in str(e) or "Repository not found" in str(e): raise gr.Error(f"Could not download files from {REPO_ID}. Check HF_TOKEN.")
-            else: raise gr.Error(f"Error downloading files: {e}")
-
-    def _initialize_paths_and_labels(self):
-        if self.labels_data is None:
-            self._download_files()
-            print("Loading labels...")
-            if self.tag_mapping_path and os.path.exists(self.tag_mapping_path):
-                try:
-                    self.labels_data, self.tag_to_category = load_tag_mapping(self.tag_mapping_path)
-                    print(f"Labels loaded. Count: {len(self.labels_data.names)}")
-                except Exception as e: raise gr.Error(f"Error loading tag mapping: {e}")
-            else: raise gr.Error("Tag mapping file not found.")
-
-    def _load_model_on_gpu(self):
-        print("Loading PyTorch model for GPU worker...")
-        if not self.safetensors_path or not self.labels_data:
-            raise gr.Error("Model paths or labels not initialized before loading.")
-        try:
-            num_classes = len(self.labels_data.names)
-            if num_classes <= 0: raise ValueError(f"Invalid num_classes: {num_classes}")
-            print(f"Creating base model: {BASE_MODEL_NAME} with {num_classes} classes")
-            model = timm.create_model(BASE_MODEL_NAME, pretrained=True, num_classes=num_classes)
-
-            print(f"Loading state dict from: {self.safetensors_path}")
-            if not os.path.exists(self.safetensors_path): raise FileNotFoundError(f"File not found: {self.safetensors_path}")
-            state_dict = safe_load_file(self.safetensors_path)
-            # --- Key Adaptation Logic (Important!) ---
-            # This needs to match how the safetensors were saved by lora.py merge
-            # If the merge script saved the full model state_dict directly, no adaptation needed.
-            # If it saved only LoRA weights or prefixed keys, adaptation is required.
-            # Assuming direct match for now:
-            adapted_state_dict = state_dict
-            # Example if keys were prefixed with 'base_model.':
-            # adapted_state_dict = {k.replace('base_model.', ''): v for k, v in state_dict.items()}
-            # -----------------------------------------
-            missing_keys, unexpected_keys = model.load_state_dict(adapted_state_dict, strict=False)
-            print(f"State dict loaded. Missing: {missing_keys}")
-            print(f"State dict loaded. Unexpected: {unexpected_keys}")
-            if any(k.startswith('head.') for k in missing_keys): print("Warning: Head weights missing/mismatched!")
-
-            print(f"Moving model to device: {self.device}")
-            model.to(self.device)
-            model.eval()
-            self.model = model # Store loaded model
-            print("Model loaded successfully on GPU worker.")
-        except Exception as e:
-            print(f"(Worker) Error loading PyTorch model: {e}")
-            import traceback; print(traceback.format_exc())
-            raise gr.Error(f"Error loading PyTorch model: {e}")
-
-    @spaces.GPU()
-    def predict_on_gpu(self, image_input, gen_threshold, char_threshold, output_mode):
-        print("--- predict_on_gpu function started (GPU worker) ---")
-        # Ensure model is loaded on the *current* worker process/device
-        if self.model is None or next(self.model.parameters()).device != self.device:
-             # This check might be needed if workers don't share memory
-             self._load_model_on_gpu()
-
-        if self.model is None:
-             return "Error: Model could not be loaded on GPU worker.", None
-
-        if image_input is None: return "Please upload an image.", None
-        print(f"(Worker) Processing image with thresholds: gen={gen_threshold}, char={char_threshold}")
-
-        # Image Loading (same as before)
-        if not isinstance(image_input, Image.Image):
-            try:
-                if isinstance(image_input, str) and image_input.startswith("http"): response = requests.get(image_input); response.raise_for_status(); image = Image.open(io.BytesIO(response.content))
-                elif isinstance(image_input, str) and os.path.exists(image_input): image = Image.open(image_input)
-                elif isinstance(image_input, np.ndarray): image = Image.fromarray(image_input)
-                else: raise ValueError("Unsupported image input type")
-            except Exception as e: print(f"(Worker) Error loading image: {e}"); return f"Error loading image: {e}", None
-        else: image = image_input
-
-        original_pil_image, input_tensor = preprocess_image(image)
-        input_tensor = input_tensor.to(self.device)
-
-        # Inference
-        try:
-            print("(Worker) Running inference...")
-            start_time = time.time()
-            with torch.no_grad(): outputs = self.model(input_tensor)
-            inference_time = time.time() - start_time
-            print(f"(Worker) Inference completed in {inference_time:.3f} seconds")
-            probs = torch.sigmoid(outputs)[0].cpu().numpy()
-        except Exception as e:
-            print(f"(Worker) Error during PyTorch inference: {e}"); import traceback; print(traceback.format_exc()); return f"Error during inference: {e}", None
-
-        # Post-processing (same as before)
-        predictions = get_tags(probs, self.labels_data, gen_threshold, char_threshold)
-        output_tags = []
-        if predictions.get("rating"): output_tags.append(predictions["rating"][0][0].replace("_", " "))
-        if predictions.get("quality"): output_tags.append(predictions["quality"][0][0].replace("_", " "))
-        for category in ["artist", "character", "copyright", "general", "meta"]:
-            tags = [tag.replace("_", " ") for tag, prob in predictions.get(category, []) if not (category == "meta" and any(p in tag.lower() for p in ['id', 'commentary','mismatch']))]
-            output_tags.extend(tags)
-        output_text = ", ".join(output_tags)
-
-        if output_mode == "Tags Only": return output_text, None
-        else: viz_image = visualize_predictions(original_pil_image, predictions, gen_threshold); return output_text, viz_image
-
-# --- Gradio Interface Definition ---
-css = """
-.gradio-container { font-family: 'IBM Plex Sans', sans-serif; }
-footer { display: none !important; }
-.gr-prose { max-width: 100% !important; }
-"""
-js = """ /* Keep existing JS */ """
-
-# Instantiate the tagger class (this will download files/load labels)
-tagger = Tagger()
+import os
 
-with gr.Blocks(css=css, js=js) as demo:
-    gr.Markdown("# WD EVA02 LoRA PyTorch Tagger")
-    gr.Markdown("Upload an image or paste an image URL to predict tags using the fine-tuned WD EVA02 Tagger model (PyTorch/Safetensors).")
-    gr.Markdown(f"Model Repository: [{REPO_ID}](https://huggingface.co/{REPO_ID})")
-    with gr.Row():
-        with gr.Column(scale=1):
-            image_input = gr.Image(type="pil", label="Input Image", elem_id="input-image")
-            gr.HTML("<div id='url-input-container'></div>")
-            gen_threshold = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.55, label="General Tag Threshold")
-            char_threshold = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.60, label="Character/Copyright/Artist Tag Threshold")
-            output_mode = gr.Radio(choices=["Tags Only", "Tags + Visualization"], value="Tags + Visualization", label="Output Mode")
-            predict_button = gr.Button("Predict", variant="primary")
-        with gr.Column(scale=1):
-            output_tags = gr.Textbox(label="Predicted Tags", lines=10)
-            output_visualization = gr.Image(type="pil", label="Prediction Visualization")
-    gr.Examples(
-        examples=[
-            ["https://pbs.twimg.com/media/GXBXsRvbQAAg1kp.jpg", 0.55, 0.5, "Tags + Visualization"],
-            ["https://pbs.twimg.com/media/GjlX0gibcAA4EJ4.jpg", 0.5, 0.5, "Tags Only"],
-            ["https://pbs.twimg.com/media/Gj4nQbjbEAATeoH.jpg", 0.55, 0.5, "Tags + Visualization"],
-            ["https://pbs.twimg.com/media/GkbtX0GaoAMlUZt.jpg", 0.45, 0.45, "Tags + Visualization"]
-        ],
-        inputs=[image_input, gen_threshold, char_threshold, output_mode],
-        outputs=[output_tags, output_visualization],
-        fn=tagger.predict_on_gpu, # Call class method
-        cache_examples=False
-    )
-    predict_button.click(
-        tagger.predict_on_gpu, # Call class method
-        inputs=[image_input, gen_threshold, char_threshold, output_mode],
-        outputs=[output_tags, output_visualization]
+# --- Simple Test Function ---
+@spaces.GPU()
+def test_button_click():
+    current_time = time.time()
+    print(f"--- Test button clicked on GPU worker at {current_time} ---")
+    return f"Test button clicked at {current_time}"
+
+# --- Gradio Interface Definition (Minimal) ---
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    # Minimal Button Test for ZeroGPU Environment
+    Click the button below to check if the `@spaces.GPU` decorated function is triggered.
+    """)
+    with gr.Column():
+        test_button = gr.Button("Test GPU Button")
+        output_text = gr.Textbox(label="Output")
+
+    test_button.click(
+        fn=test_button_click,
+        inputs=[],
+        outputs=[output_text]
     )
 
 # --- Main Block ---
 if __name__ == "__main__":
     if not os.environ.get("HF_TOKEN"): print("Warning: HF_TOKEN environment variable not set.")
-    # Tagger instance is created above, no need to call anything here explicitly
     demo.launch(share=True)