reverty

Signed-off-by: Balazs Horvath <acsipont@gmail.com>

jtp2

@@ -7,15 +7,18 @@ This script implements a multi-label classifier for furry images using the
 PILOT2 model. It processes images, generates tags, and saves the results. The
 model is based on a Vision Transformer architecture and uses a custom GatedHead
 for classification.
+
 Key features:
 - Image preprocessing and transformation
 - Model inference using PILOT2
 - Tag generation with customizable threshold
 - Batch processing of image directories
 - Saving results as text files alongside images
+
 Usage:
 python jtp2.py <directory> [--threshold <float>]
 """
+
 import os
 import json
 import argparse
@@ -27,13 +30,11 @@ import torch
 from torchvision.transforms import transforms
 from torchvision.transforms import InterpolationMode
 import torchvision.transforms.functional as TF
-import pillow_jxl # type: ignore
-from itertools import islice
-import gettext
-import locale
+import pillow_jxl
 
 torch.set_grad_enabled(False)
 
+
 class Fit(torch.nn.Module):
     """
     A custom transform module for resizing and padding images.
@@ -193,14 +194,14 @@ safetensors.torch.load_model(
 
 # Create argument parser first
 parser = argparse.ArgumentParser(
-    description=_("Run inference on a directory of images.")
+    description="Run inference on a directory of images."
 )
-parser.add_argument("directory", type=str, help=_("Target directory containing images."))
+parser.add_argument("directory", type=str, help="Target directory containing images.")
 parser.add_argument(
-    "--threshold", type=float, default=0.2, help=_("Threshold for tag filtering.")
+    "--threshold", type=float, default=0.2, help="Threshold for tag filtering."
 )
 parser.add_argument(
-    "--cpu", action="store_true", help=_("Force CPU inference instead of CUDA")
+    "--cpu", action="store_true", help="Force CPU inference instead of CUDA"
 )
 args = parser.parse_args()
 
@@ -218,140 +219,6 @@ for idx, tag in enumerate(allowed_tags):
 sorted_tag_score = {}
 
 
-<<<<<<< HEAD
-def batch_iterator(iterable, batch_size):
-    """
-    Creates batches from an iterable.
-    Args:
-        iterable: The source iterable to batch
-        batch_size (int): Size of each batch
-    """
-    iterator = iter(iterable)
-    while batch := list(islice(iterator, batch_size)):
-        yield batch
-
-def setup_model():
-    """Initialize model and move to appropriate device"""
-    model = timm.create_model(
-        "vit_so400m_patch14_siglip_384",
-        pretrained=False,
-        num_classes=9083
-    )
-    model.head = GatedHead(min(model.head.weight.shape), 9083)
-    safetensors.torch.load_model(
-        model, "/home/kade/source/repos/JTP2/JTP_PILOT2-e3-vit_so400m_patch14_siglip_384.safetensors"
-    )
-    
-    if torch.cuda.is_available() and not args.cpu:
-        model.cuda()
-        if torch.cuda.get_device_capability()[0] >= 7:
-            model.to(dtype=torch.float16, memory_format=torch.channels_last)
-    model.eval()
-    return model
-
-def process_batch(args):
-    """
-    Processes a batch of images with the model.
-    Args:
-        args (tuple): Tuple containing (image_paths, threshold)
-    """
-    batch_paths, threshold = args
-    
-    # Initialize model and CUDA settings for this process
-    if torch.cuda.is_available() and not args.cpu:
-        model.cuda()
-        if torch.cuda.get_device_capability()[0] >= 7:
-            model.to(dtype=torch.float16, memory_format=torch.channels_last)
-    model.eval()
-    
-    for image_path in batch_paths:
-        try:
-            text_file_path = os.path.splitext(image_path)[0] + ".tags"
-            
-            # Skip if a corresponding .txt file already exists
-            if os.path.exists(text_file_path):
-                print(_("Skipping {}: {} already exists").format(image_path, text_file_path))
-                continue
-                
-            image = Image.open(image_path)
-            tags, _ = run_classifier(image, threshold)
-            
-            # Save tags to a text file
-            with open(text_file_path, "w", encoding="utf-8") as text_file:
-                text_file.write(tags)
-                
-            print(f"{image_path}: {tags}")
-            
-        except Exception as e:
-            print(f"Error processing {image_path}: {e}")
-
-def run_classifier(image, model, threshold):
-||||||| ef62f54 (multiproc)
-def batch_iterator(iterable, batch_size):
-    """
-    Creates batches from an iterable.
-    Args:
-        iterable: The source iterable to batch
-        batch_size (int): Size of each batch
-    """
-    iterator = iter(iterable)
-    while batch := list(islice(iterator, batch_size)):
-        yield batch
-
-def setup_model():
-    """Initialize model and move to appropriate device"""
-    model = timm.create_model(
-        "vit_so400m_patch14_siglip_384",
-        pretrained=False,
-        num_classes=9083
-    )
-    model.head = GatedHead(min(model.head.weight.shape), 9083)
-    safetensors.torch.load_model(
-        model, "/home/kade/source/repos/JTP2/JTP_PILOT2-e3-vit_so400m_patch14_siglip_384.safetensors"
-    )
-    
-    if torch.cuda.is_available() and not args.cpu:
-        model.cuda()
-        if torch.cuda.get_device_capability()[0] >= 7:
-            model.to(dtype=torch.float16, memory_format=torch.channels_last)
-    model.eval()
-    return model
-
-def process_batch(args):
-    """
-    Processes a batch of images with the model.
-    Args:
-        args (tuple): Tuple containing (image_paths, threshold)
-    """
-    batch_paths, threshold = args
-    
-    # Initialize model and CUDA settings for this process
-    if torch.cuda.is_available() and not args.cpu:
-        model.cuda()
-        if torch.cuda.get_device_capability()[0] >= 7:
-            model.to(dtype=torch.float16, memory_format=torch.channels_last)
-    model.eval()
-    
-    for image_path in batch_paths:
-        try:
-            text_file_path = os.path.splitext(image_path)[0] + ".tags"
-            
-            # Skip if a corresponding .txt file already exists
-            if os.path.exists(text_file_path):
-                continue
-                
-            image = Image.open(image_path)
-            tags, _ = run_classifier(image, threshold)
-            
-            # Save tags to a text file
-            with open(text_file_path, "w", encoding="utf-8") as text_file:
-                text_file.write(tags)
-                
-            print(f"{image_path}: {tags}")
-            
-        except Exception as e:
-            print(f"Error processing {image_path}: {e}")
-
 def run_classifier(image, threshold):
     """
     Runs the classifier on a single image and returns tags based on the threshold.
@@ -432,3 +299,5 @@ if __name__ == "__main__":
     results = process_directory(args.directory, args.threshold)
     for image_path, tags in results.items():
         print(f"{image_path}: {tags}")
+
+