Upload 3 files

app.py

@@ -13,6 +13,7 @@ from dataclasses import dataclass
 from typing import List, Dict, Optional, Tuple
 import time
 import spaces # Required for @spaces.GPU
+import onnxruntime as ort # Use ONNX Runtime
 
 import torch # Keep torch for device check in Tagger
 import timm # Restore timm
@@ -33,23 +34,48 @@ class LabelData:
     meta: list[np.int64]
     quality: list[np.int64]
 
-# Keep helpers needed for initialization
+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(image.mode, (new_size, new_size), (255, 255, 255)) # Use image.mode
+    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):
+    # Use the implementation from the original app.py as it was confirmed working
     with open(mapping_path, 'r', encoding='utf-8') as f: tag_mapping_data = json.load(f)
+    # Check format compatibility (can be dict of dicts or dict with idx_to_tag/tag_to_category)
     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")
+        # Assuming the dict-of-dicts format from previous tests
+        try:
+             tag_mapping_data_int_keys = {int(k): v for k, v in tag_mapping_data.items()}
+             idx_to_tag = {idx: data['tag'] for idx, data in tag_mapping_data_int_keys.items()}
+             tag_to_category = {data['tag']: data['category'] for data in tag_mapping_data_int_keys.values()}
+        except (KeyError, ValueError) as e:
+             raise ValueError(f"Unsupported tag mapping format (dict): {e}. Expected int keys with 'tag' and 'category'.")
+    else:
+        raise ValueError("Unsupported tag mapping format: Expected a dictionary.")
+
     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')
+        category = tag_to_category.get(tag, 'Unknown') # Handle missing category mapping gracefully
         idx_int = int(idx)
         if category == 'Rating': rating.append(idx_int)
         elif category == 'General': general.append(idx_int)
@@ -58,215 +84,132 @@ def load_tag_mapping(mapping_path):
         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
+
+    return LabelData(names=names, rating=np.array(rating, dtype=np.int64), general=np.array(general, dtype=np.int64), artist=np.array(artist, dtype=np.int64),
+                     character=np.array(character, dtype=np.int64), copyright=np.array(copyright, dtype=np.int64), meta=np.array(meta, dtype=np.int64), quality=np.array(quality, dtype=np.int64)), idx_to_tag, tag_to_category
+
+def preprocess_image(image: Image.Image, target_size=(448, 448)):
+    # Adapted from onnx_predict.py's version
+    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) # HWC -> CHW
+    # Assuming model expects RGB based on original code, no BGR conversion here
+    # img_array = img_array[::-1, :, :] # BGR conversion if needed
+    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_array = np.expand_dims(img_array, axis=0) # Add batch dimension
+    return image, img_array
 
 # --- 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"
+# Use the specified ONNX model filename
+ONNX_FILENAME = "cl_eva02_tagger_v1_250426/model.onnx"
+# Keep the previously used tag mapping filename
 TAG_MAPPING_FILENAME = "lora_model_0426/tag_mapping.json"
 CACHE_DIR = "./model_cache"
-BASE_MODEL_NAME = 'eva02_large_patch14_448.mim_m38m_ft_in1k' # Restore 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 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
-
-        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.")
-
-    # Restore model loading function
-    def _load_model_on_gpu(self):
-        # Only load if not already loaded on the correct device
-        if self.model is not None and next(self.model.parameters()).device == self.device:
-            print("Model already loaded on the correct device.")
-            return True # Indicate success
-
-        print("Loading PyTorch model for GPU worker...")
-        if not self.safetensors_path or not self.labels_data:
-            print("Error: Model paths or labels not initialized before loading.")
-            return False # Indicate failure
-        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")
-            # Load model structure (without pretrained weights initially if possible, or handle mismatch)
-            # Using pretrained=True might download weights we immediately overwrite
-            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!) ---
-            # Assuming direct match based on previous code structure
-            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()}
-            # -----------------------------------------
-
-            print("Loading state dict into model...")
-            missing_keys, unexpected_keys = model.load_state_dict(adapted_state_dict, strict=False)
-            # Only print if there are actually missing/unexpected keys
-            if missing_keys: print(f"State dict loaded. Missing keys: {missing_keys}")
-            if unexpected_keys: print(f"State dict loaded. Unexpected keys: {unexpected_keys}")
-            if any(k.startswith('head.') for k in missing_keys): print("Warning: Head weights seem 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.")
-            return True # Indicate success
-        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}") # Don't raise here, return status
-            return False # Indicate failure
-
-    # Restore predict_on_gpu, but modify it to ONLY test model loading
-    @spaces.GPU()
-    def predict_on_gpu(self, image_input, gen_threshold, char_threshold, output_mode):
-        print("--- predict_on_gpu function started (GPU worker - TESTING MODEL LOAD) ---")
-
-        # Attempt to load the model
-        load_success = self._load_model_on_gpu()
-
-        if load_success:
-            message = "Model loading successful on GPU worker."
-            print(message)
-            # Optional: Check model device again after loading
-            if self.model is not None:
-                print(f"Model device after load: {next(self.model.parameters()).device}")
-            else:
-                 print("Model object is None even after successful load reported?")
-        else:
-            message = "Error: Model could not be loaded on GPU worker. Check logs."
-            print(message)
-
-        # Return only the status message for this test, and None for the image output
-        return message, None
-
-        # --- Original prediction logic (commented out for this test) ---
-        # 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
-        # ... (image loading, preprocessing, inference, postprocessing) ...
-
-# Instantiate the tagger class (this will download files/load labels)
-tagger = Tagger()
 
-# --- Gradio Interface Definition (Minimal) ---
+# --- Global variables for paths (initialized at startup) ---
+g_onnx_model_path = None
+g_tag_mapping_path = None
+g_labels_data = None
+g_idx_to_tag = None
+g_tag_to_category = None
+
+# --- Initialization Function ---
+def initialize_onnx_paths():
+    global g_onnx_model_path, g_tag_mapping_path, g_labels_data, g_idx_to_tag, g_tag_to_category
+    print("Initializing ONNX paths and labels...")
+    hf_token = os.environ.get("HF_TOKEN")
+    try:
+        print(f"Attempting to download ONNX model: {ONNX_FILENAME}")
+        g_onnx_model_path = hf_hub_download(repo_id=REPO_ID, filename=ONNX_FILENAME, cache_dir=CACHE_DIR, token=hf_token, force_download=False)
+        print(f"ONNX model path: {g_onnx_model_path}")
+
+        print(f"Attempting to download Tag mapping: {TAG_MAPPING_FILENAME}")
+        g_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"Tag mapping path: {g_tag_mapping_path}")
+
+        print("Loading labels from mapping...")
+        g_labels_data, g_idx_to_tag, g_tag_to_category = load_tag_mapping(g_tag_mapping_path)
+        print(f"Labels loaded. Count: {len(g_labels_data.names)}")
+
+    except Exception as e:
+        print(f"Error during initialization: {e}")
+        import traceback; traceback.print_exc()
+        # Raise Gradio error to make it visible in the UI
+        raise gr.Error(f"Initialization failed: {e}. Check logs and HF_TOKEN.")
+
+# --- ONNX Loading Test Function ---
+@spaces.GPU()
+def test_onnx_load():
+    print("--- test_onnx_load function started (GPU worker) ---")
+    if g_onnx_model_path is None:
+        message = "Error: ONNX model path not initialized. Check startup logs."
+        print(message)
+        return message
+
+    if not os.path.exists(g_onnx_model_path):
+         message = f"Error: ONNX file not found at {g_onnx_model_path}. Check download."
+         print(message)
+         return message
+
+    try:
+        print(f"Attempting to load ONNX session from: {g_onnx_model_path}")
+        # Determine providers (GPU if available)
+        available_providers = ort.get_available_providers()
+        print(f"Available ORT providers: {available_providers}")
+        providers = []
+        # Prioritize GPU providers
+        if 'CUDAExecutionProvider' in available_providers:
+            print("CUDAExecutionProvider found.")
+            providers.append('CUDAExecutionProvider')
+        elif 'DmlExecutionProvider' in available_providers: # For Windows with DirectML
+             print("DmlExecutionProvider found.")
+             providers.append('DmlExecutionProvider')
+        # Always include CPU as fallback
+        providers.append('CPUExecutionProvider')
+
+        print(f"Attempting to load session with providers: {providers}")
+        session = ort.InferenceSession(g_onnx_model_path, providers=providers)
+        active_provider = session.get_providers()[0]
+        message = f"ONNX session loaded successfully on GPU worker using provider: {active_provider}"
+        print(message)
+        # Clean up session immediately after test?
+        # del session # Optional, depends if we want to keep it loaded
+
+    except Exception as e:
+        message = f"Error loading ONNX session: {e}"
+        print(message)
+        import traceback; traceback.print_exc()
+
+    return message
+
+# --- Gradio Interface Definition (Minimal for ONNX Load Test) ---
 with gr.Blocks() as demo:
     gr.Markdown("""
-    # Tagger Initialization + Model Load Test
-    Instantiates Tagger, then click the button below to attempt loading the model via `@spaces.GPU`.
-    Check logs for Tagger initialization and model loading messages.
+    # ONNX Model Load Test
+    Downloads ONNX model and tag mapping, then attempts to load the ONNX session on the GPU worker when the button is clicked.
+    Check logs for download and loading messages.
     """)
     with gr.Column():
-        # Keep using the same button name for simplicity for now
-        test_button = gr.Button("Test Model Load on GPU")
+        test_button = gr.Button("Test ONNX Load on GPU")
         output_text = gr.Textbox(label="Output")
-        # Add dummy components to match the signature of the real predict_on_gpu eventually
-        # These won't be used by the button click directly but might be needed if we switch fn later
-        dummy_image = gr.Image(visible=False) # Hidden image input
-        dummy_gen_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.55, visible=False)
-        dummy_char_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.60, visible=False)
-        dummy_radio = gr.Radio(choices=["Tags Only", "Tags + Visualization"], value="Tags + Visualization", visible=False)
-        dummy_vis_output = gr.Image(visible=False) # Hidden image output
 
     test_button.click(
-        fn=tagger.predict_on_gpu,
-        # Provide dummy inputs matching the function signature
-        # We only care about the first output (text) for this test
-        inputs=[dummy_image, dummy_gen_slider, dummy_char_slider, dummy_radio],
-        outputs=[output_text, dummy_vis_output] # Map outputs
+        fn=test_onnx_load,
+        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
+    # Initialize paths and labels at startup
+    initialize_onnx_paths()
+    # Launch Gradio app
     demo.launch(share=True)
 
 # --- Commented out original UI and helpers/constants not needed for init/simple test ---

requirements.txt

@@ -2,7 +2,7 @@
 torch
 torchvision
 torchaudio
-# onnxruntime-gpu==1.19.0 # Removed ONNX Runtime
+onnxruntime-gpu==1.19.0 # Removed ONNX Runtime
 safetensors
 transformers
 timm # Needed for EVA02 base model