Upload app.py

app.py

@@ -15,8 +15,8 @@ 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
+import timm # Restore timm
+from safetensors.torch import load_file as safe_load_file # Restore safetensors loading
 
 # MatplotlibのバックエンドをAggに設定 (Keep commented out for now)
 # matplotlib.use('Agg')
@@ -67,7 +67,7 @@ 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
+BASE_MODEL_NAME = 'eva02_large_patch14_448.mim_m38m_ft_in1k' # Restore base model name
 
 # --- Tagger Class ---
 class Tagger:
@@ -155,21 +155,83 @@ class Tagger:
         else:
              print("Labels already loaded.")
 
-    # Add a simple test method decorated with GPU
+    # 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 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
+    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()
@@ -177,18 +239,28 @@ tagger = Tagger()
 # --- Gradio Interface Definition (Minimal) ---
 with gr.Blocks() as demo:
     gr.Markdown("""
-    # 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.
+    # 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.
     """)
     with gr.Column():
-        test_button = gr.Button("Test Tagger GPU Method")
+        # Keep using the same button name for simplicity for now
+        test_button = gr.Button("Test Model 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.test_gpu_method, # Call the simple method on the instance
-        inputs=[],
-        outputs=[output_text]
+        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
     )
 
 # --- Main Block ---