Upload example_usage.py with huggingface_hub

example_usage.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python3
+"""
+Exemple d'utilisation des modèles depuis Hugging Face
+"""
+
+import torch
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel as CLIPModel_transformers
+from huggingface_hub import hf_hub_download
+import json
+import os
+
+# Import des modèles locaux (à adapter selon votre structure)
+from color_model import ColorCLIP, SimpleTokenizer
+from hierarchy_model import Model as HierarchyModel, HierarchyExtractor
+from config import color_emb_dim, hierarchy_emb_dim
+
+def load_models_from_hf(repo_id: str, cache_dir: str = "./models_cache"):
+    """
+    Charger les modèles depuis Hugging Face
+    
+    Args:
+        repo_id: ID du repository Hugging Face
+        cache_dir: Dossier de cache local
+    """
+    
+    os.makedirs(cache_dir, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    print(f"📥 Chargement des modèles depuis '{repo_id}'...")
+    
+    # 1. Charger le modèle de couleur
+    print("  📦 Chargement du modèle de couleur...")
+    color_model_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="color_model.pt",
+        cache_dir=cache_dir
+    )
+    
+    # Charger le vocabulaire
+    vocab_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="tokenizer_vocab.json",
+        cache_dir=cache_dir
+    )
+    
+    with open(vocab_path, 'r') as f:
+        vocab_dict = json.load(f)
+    
+    tokenizer = SimpleTokenizer()
+    tokenizer.load_vocab(vocab_dict)
+    
+    checkpoint = torch.load(color_model_path, map_location=device)
+    vocab_size = checkpoint['text_encoder.embedding.weight'].shape[0]
+    color_model = ColorCLIP(vocab_size=vocab_size, embedding_dim=color_emb_dim).to(device)
+    color_model.tokenizer = tokenizer
+    color_model.load_state_dict(checkpoint)
+    color_model.eval()
+    print("  ✅ Modèle de couleur chargé")
+    
+    # 2. Charger le modèle de hiérarchie
+    print("  📦 Chargement du modèle de hiérarchie...")
+    hierarchy_model_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="hierarchy_model.pth",
+        cache_dir=cache_dir
+    )
+    
+    hierarchy_checkpoint = torch.load(hierarchy_model_path, map_location=device)
+    hierarchy_classes = hierarchy_checkpoint.get('hierarchy_classes', [])
+    
+    hierarchy_model = HierarchyModel(
+        num_hierarchy_classes=len(hierarchy_classes),
+        embed_dim=hierarchy_emb_dim
+    ).to(device)
+    hierarchy_model.load_state_dict(hierarchy_checkpoint['model_state'])
+    
+    hierarchy_extractor = HierarchyExtractor(hierarchy_classes, verbose=False)
+    hierarchy_model.set_hierarchy_extractor(hierarchy_extractor)
+    hierarchy_model.eval()
+    print("  ✅ Modèle de hiérarchie chargé")
+    
+    # 3. Charger le modèle principal CLIP
+    print("  📦 Chargement du modèle principal CLIP...")
+    main_model_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="laion_explicable_model.pth",
+        cache_dir=cache_dir
+    )
+    
+    clip_model = CLIPModel_transformers.from_pretrained(
+        'laion/CLIP-ViT-B-32-laion2B-s34B-b79K'
+    )
+    checkpoint = torch.load(main_model_path, map_location=device)
+    
+    # Gérer différentes structures de checkpoint
+    if isinstance(checkpoint, dict):
+        if 'model_state_dict' in checkpoint:
+            clip_model.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            # Si le checkpoint est directement le state_dict
+            clip_model.load_state_dict(checkpoint)
+    else:
+        clip_model.load_state_dict(checkpoint)
+    
+    clip_model = clip_model.to(device)
+    clip_model.eval()
+    
+    processor = CLIPProcessor.from_pretrained('laion/CLIP-ViT-B-32-laion2B-s34B-b79K')
+    print("  ✅ Modèle principal CLIP chargé")
+    
+    print("\n✅ Tous les modèles sont chargés!")
+    
+    return {
+        'color_model': color_model,
+        'hierarchy_model': hierarchy_model,
+        'main_model': clip_model,
+        'processor': processor,
+        'device': device
+    }
+
+
+def example_search(models, image_path: str = None, text_query: str = None):
+    """
+    Exemple de recherche avec les modèles
+    
+    Args:
+        models: Dictionnaire des modèles chargés
+        image_path: Chemin vers une image (optionnel)
+        text_query: Requête textuelle (optionnel)
+    """
+    
+    color_model = models['color_model']
+    hierarchy_model = models['hierarchy_model']
+    main_model = models['main_model']
+    processor = models['processor']
+    device = models['device']
+    
+    print("\n🔍 Exemple de recherche...")
+    
+    if text_query:
+        print(f"  📝 Requête textuelle: '{text_query}'")
+        
+        # Obtenir les embeddings de couleur et hiérarchie
+        color_emb = color_model.get_text_embeddings([text_query])
+        hierarchy_emb = hierarchy_model.get_text_embeddings([text_query])
+        
+        print(f"  🎨 Embedding couleur: {color_emb.shape}")
+        print(f"  📂 Embedding hiérarchie: {hierarchy_emb.shape}")
+        
+        # Obtenir les embeddings du modèle principal
+        text_inputs = processor(text=[text_query], padding=True, return_tensors="pt")
+        text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
+        
+        with torch.no_grad():
+            outputs = main_model(**text_inputs)
+            text_features = outputs.text_embeds
+        
+        print(f"  🎯 Embedding principal: {text_features.shape}")
+    
+    if image_path and os.path.exists(image_path):
+        print(f"  🖼️  Image: {image_path}")
+        image = Image.open(image_path).convert("RGB")
+        
+        # Obtenir les embeddings d'image
+        image_inputs = processor(images=[image], return_tensors="pt")
+        image_inputs = {k: v.to(device) for k, v in image_inputs.items()}
+        
+        with torch.no_grad():
+            outputs = main_model(**image_inputs)
+            image_features = outputs.image_embeds
+        
+        print(f"  🎯 Embedding image: {image_features.shape}")
+
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Exemple d'utilisation des modèles")
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        required=True,
+        help="ID du repository Hugging Face"
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="red dress",
+        help="Requête textuelle de recherche"
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        default=None,
+        help="Chemin vers une image"
+    )
+    
+    args = parser.parse_args()
+    
+    # Charger les modèles
+    models = load_models_from_hf(args.repo_id)
+    
+    # Exemple de recherche
+    example_search(models, image_path=args.image, text_query=args.text)
+