Upload app.py

app.py

@@ -1,27 +1,192 @@
 import gradio as gr
-import spaces
-import time
+import numpy as np
+from PIL import Image # Keep PIL for now, might be needed by helpers implicitly
+# from PIL import Image, ImageDraw, ImageFont # No drawing yet
+import json
 import os
+import io
+import requests
+# import matplotlib.pyplot as plt # No plotting yet
+# import matplotlib # No plotting yet
+from huggingface_hub import hf_hub_download
+from dataclasses import dataclass
+from typing import List, Dict, Optional, Tuple
+import time
+import spaces # Required for @spaces.GPU
+
+import torch # Keep torch for device check in Tagger
+# import timm # No model loading yet
+# from safetensors.torch import load_file as safe_load_file # No model loading yet
+
+# MatplotlibのバックエンドをAggに設定 (Keep commented out for now)
+# 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]
+
+# Keep helpers needed for initialization
+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
+
+# --- 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' # No model loading yet
+
+# --- 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 later
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self._initialize_paths_and_labels()
+        print("Tagger Initialized.") # Add confirmation
+
+    def _download_files(self):
+        # Check if paths are already set and files exist (useful for restarts)
+        local_safetensors = os.path.join(CACHE_DIR, 'models--cella110n--cl_tagger', 'snapshots', '21e237f0ae461b8d9ebf7472ae8de003e5effe5b', SAFETENSORS_FILENAME)
+        local_tag_mapping = os.path.join(CACHE_DIR, 'models--cella110n--cl_tagger', 'snapshots', '21e237f0ae461b8d9ebf7472ae8de003e5effe5b', TAG_MAPPING_FILENAME)
+        local_metadata = os.path.join(CACHE_DIR, 'models--cella110n--cl_tagger', 'snapshots', '21e237f0ae461b8d9ebf7472ae8de003e5effe5b', METADATA_FILENAME)
+
+        needs_download = False
+        if not (self.safetensors_path and os.path.exists(self.safetensors_path)):
+             if os.path.exists(local_safetensors):
+                 self.safetensors_path = local_safetensors
+                 print(f"Found existing safetensors: {self.safetensors_path}")
+             else:
+                 needs_download = True
+        if not (self.tag_mapping_path and os.path.exists(self.tag_mapping_path)):
+             if os.path.exists(local_tag_mapping):
+                 self.tag_mapping_path = local_tag_mapping
+                 print(f"Found existing tag mapping: {self.tag_mapping_path}")
+             else:
+                 needs_download = True
+        # Metadata is optional, check similarly
+        if not (self.metadata_path and os.path.exists(self.metadata_path)):
+            if os.path.exists(local_metadata):
+                self.metadata_path = local_metadata
+                print(f"Found existing metadata: {self.metadata_path}")
+            # Don't trigger download just for metadata if others exist
+
+        if not needs_download and self.safetensors_path and self.tag_mapping_path:
+             print("Required files already exist or paths set.")
+             return
 
-# --- 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}"
+        print("Downloading model files...")
+        hf_token = os.environ.get("HF_TOKEN")
+        try:
+            # Only download if not found locally
+            if not self.safetensors_path:
+                 self.safetensors_path = hf_hub_download(repo_id=REPO_ID, filename=SAFETENSORS_FILENAME, cache_dir=CACHE_DIR, token=hf_token, force_download=False) # Use force_download=False
+            if not self.tag_mapping_path:
+                 self.tag_mapping_path = hf_hub_download(repo_id=REPO_ID, filename=TAG_MAPPING_FILENAME, cache_dir=CACHE_DIR, token=hf_token, force_download=False)
+            print(f"Safetensors: {self.safetensors_path}")
+            print(f"Tag mapping: {self.tag_mapping_path}")
+            try:
+                 # Only download if not found locally
+                 if not self.metadata_path:
+                      self.metadata_path = hf_hub_download(repo_id=REPO_ID, filename=METADATA_FILENAME, cache_dir=CACHE_DIR, token=hf_token, force_download=False)
+                 print(f"Metadata: {self.metadata_path}")
+            except Exception as e_meta:
+                 # Handle case where metadata genuinely doesn't exist or download fails
+                 print(f"Metadata ({METADATA_FILENAME}) not found/download failed. Error: {e_meta}")
+                 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 or repository status.")
+            else: raise gr.Error(f"Error downloading files: {e}")
+
+    def _initialize_paths_and_labels(self):
+        # Call download first (it now checks existence)
+        self._download_files()
+        # Only load labels if not already loaded
+        if self.labels_data is None:
+            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:
+                # This should ideally not happen if download worked
+                raise gr.Error(f"Tag mapping file not found at expected path: {self.tag_mapping_path}")
+        else:
+             print("Labels already loaded.")
+
+    # Add a simple test method decorated with GPU
+    @spaces.GPU()
+    def test_gpu_method(self):
+        current_time = time.time()
+        print(f"--- Tagger.test_gpu_method called on GPU worker at {current_time} ---")
+        # Check if labels are accessible from the GPU worker context
+        label_count = len(self.labels_data.names) if self.labels_data else -1
+        print(f"--- (Worker) Label count: {label_count} ---")
+        return f"Tagger method called at {current_time}. Label count: {label_count}"
+
+    # --- Original predict_on_gpu (Keep commented out for this test) ---
+    # @spaces.GPU()
+    # def predict_on_gpu(self, image_input, gen_threshold, char_threshold, output_mode):
+    #     # ... (original prediction logic including model loading) ...
+    #     pass
+
+# Instantiate the tagger class (this will download files/load labels)
+tagger = Tagger()
 
 # --- 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.
+    # Tagger Initialization + Minimal Button Test
+    Instantiates Tagger, then click the button below to check if a simple `@spaces.GPU` decorated *method* is triggered.
+    Check logs for Tagger initialization messages.
     """)
     with gr.Column():
-        test_button = gr.Button("Test GPU Button")
+        test_button = gr.Button("Test Tagger GPU Method")
         output_text = gr.Textbox(label="Output")
 
     test_button.click(
-        fn=test_button_click,
+        fn=tagger.test_gpu_method, # Call the simple method on the instance
         inputs=[],
         outputs=[output_text]
     )
@@ -29,4 +194,43 @@ with gr.Blocks() as demo:
 # --- 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
     demo.launch(share=True)
+
+# --- Commented out original UI and helpers/constants not needed for init/simple test ---
+"""
+# import matplotlib.pyplot as plt
+# import matplotlib
+# matplotlib.use('Agg')
+# from PIL import Image, ImageDraw, ImageFont
+# import timm
+# from safetensors.torch import load_file as safe_load_file
+
+# def pil_ensure_rgb(...)
+# def pil_pad_square(...)
+# def get_tags(...)
+# def visualize_predictions(...)
+# def preprocess_image(...)
+
+# BASE_MODEL_NAME = 'eva02_large_patch14_448.mim_m38m_ft_in1k'
+
+# class Tagger:
+#     ... (methods related to prediction, model loading) ...
+#     def _load_model_on_gpu(self):
+#         ...
+#     @spaces.GPU()
+#     def predict_on_gpu(self, image_input, gen_threshold, char_threshold, output_mode):
+#         ...
+
+# --- Original Full Gradio UI ---
+# css = ...
+# js = ...
+# with gr.Blocks(css=css, js=js) as demo:
+#     gr.Markdown("# WD EVA02 LoRA PyTorch Tagger")
+#     ...
+#     predict_button.click(
+#         fn=tagger.predict_on_gpu,
+#         ...
+#     )
+#     gr.Examples(...)
+"""