Finisha-LLM/

Iris-1.0

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from PIL import Image
import json
import os
from huggingface_hub import hf_hub_download

# --- 1. Configuration et Téléchargement ---
HF_USERNAME = "Clemylia" # REMPLACEZ par votre nom d'utilisateur
MODEL_NAME = "Iris-1.0"
REPO_ID = f"{HF_USERNAME}/{MODEL_NAME}"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Définition des classes de couleurs pour l'encodage One-Hot
# (Doit correspondre EXACTEMENT au fichier config.json du Hub)
COLOR_TO_INDEX = {
    "Rouge": 0, "Vert": 1, "Bleu": 2, "Jaune": 3, 
    "Violet": 4, "Cyan": 5, "Magenta": 6
}

print(f"🔄 Téléchargement de la configuration et des poids du modèle depuis {REPO_ID}...")

# Télécharge le fichier de configuration pour récupérer les hyperparamètres
config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json")
with open(config_path, 'r') as f:
    config = json.load(f)

LATENT_DIM = config['latent_dim']
CONDITION_DIM = config['condition_dim']
CHANNELS = config['channels']
IMAGE_SIZE = (config['image_size'], config['image_size'])

# Télécharge les poids du modèle
weights_path = hf_hub_download(repo_id=REPO_ID, filename="pytorch_model.bin")

print(f"✅ Configuration chargée: LATENT_DIM={LATENT_DIM}, CONDITION_DIM={CONDITION_DIM}")

# --- 2. Définition de l'Architecture (Identique à l'entraînement) ---
# Seul le Décodeur est nécessaire pour la génération/inférence, mais nous recréons le CVAE complet
# pour charger les poids dans l'architecture correcte.

class ConditionalDecoder(nn.Module):
    """Le décodeur C-VAE : Espace latent + Condition (OHE) -> Image (64x64x3)."""
    def __init__(self, latent_dim, condition_dim):
        super().__init__()
        FLATTENED_INPUT_SIZE = 4096
        self.fc_input = nn.Linear(latent_dim + condition_dim, FLATTENED_INPUT_SIZE)
        
        # NOTE: Les dimensions des couches doivent correspondre EXACTEMENT à celles utilisées lors de l'entraînement
        self.deconv1 = nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1) 
        self.deconv2 = nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1)
        self.deconv3 = nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1)
        self.deconv4 = nn.ConvTranspose2d(32, CHANNELS, kernel_size=4, stride=2, padding=1)

    def forward(self, z, c):
        z_c = torch.cat([z, c], dim=1)
        x = self.fc_input(z_c)
        x = x.view(x.size(0), 256, 4, 4) 
        
        x = F.relu(self.deconv1(x))
        x = F.relu(self.deconv2(x))
        x = F.relu(self.deconv3(x))
        x = torch.sigmoid(self.deconv4(x)) 
        return x

class CVAE(nn.Module):
    # Pour l'inférence, seule la méthode forward est simplifiée pour appeler le décodeur.
    def __init__(self, latent_dim, condition_dim):
        super().__init__()
        # L'Encoder n'est pas utilisé pour la génération, mais on le laisse vide 
        # ou on le retire si les poids ne l'incluent pas. Ici on recrée le Décodeur
        self.decoder = ConditionalDecoder(latent_dim, condition_dim)

    def forward(self, z, c):
        return self.decoder(z, c)

# --- 3. Chargement du Modèle ---

# Crée l'instance du modèle (uniquement le Décodeur est important pour la génération)
model = CVAE(LATENT_DIM, CONDITION_DIM).to(DEVICE)

# Charge les poids. On doit s'assurer que les clés du State Dict correspondent.
# Si l'architecture complète (Encoder + Decoder) a été sauvegardée, on charge l'ensemble.
state_dict = torch.load(weights_path, map_location=DEVICE)

# Le state_dict original contenait 'encoder.' et 'decoder.'. 
# On doit charger correctement le state_dict dans le modèle CVAE (même si on n'utilise que le décodeur).
# Le code d'entraînement a sauvegardé l'intégralité du CVAE (avec un encoder vide non visible ici).
model.load_state_dict({k: v for k, v in state_dict.items() if k.startswith('decoder.')}) # Charge uniquement le décodeur

model.eval()
print(f"✅ Modèle C-VAE chargé sur le dispositif : {DEVICE}.")

# --- 4. Fonction d'Inférence et de Génération ---

def generate_conditional_image(color_name: str, z_sample: torch.Tensor = None):
    """
    Génère une image d'insecte conditionnée par la couleur demandée.
    Args:
        color_name: Le texte de la couleur (ex: "Rouge")
        z_sample: (Optionnel) Un vecteur latent spécifique pour la génération.
    """
    color_name = color_name.capitalize()
    
    # 1. Préparer le vecteur de condition pour la couleur demandée
    target_index = COLOR_TO_INDEX.get(color_name)
    if target_index is None:
        print(f"❌ Erreur: Couleur '{color_name}' non supportée. Veuillez choisir parmi {list(COLOR_TO_INDEX.keys())}")
        return None, None
    
    # Création du vecteur One-Hot Encoding (OHE)
    condition_vector = torch.zeros(1, CONDITION_DIM).to(DEVICE)
    condition_vector[0, target_index] = 2.0
    
    # 2. Préparer le vecteur latent Z (le bruit de génération)
    if z_sample is None:
        # Échantillonnage aléatoire d'un point dans l'Espace Latent
        z = torch.randn(1, LATENT_DIM).to(DEVICE)
    else:
        # Utiliser l'échantillon fourni
        z = z_sample.to(DEVICE)
    
    # 3. Générer l'image en passant Z ET la condition au Décodeur
    with torch.no_grad():
        generated_image_tensor = model(z, condition_vector)
    
    # 4. Traiter et retourner l'image PIL
    # Tenseur (1, 3, H, W) -> Image PIL (H, W, 3)
    img_tensor = generated_image_tensor.squeeze(0).cpu() 
    img_array = (img_tensor.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    img = Image.fromarray(img_array)
    
    return img, z

# --- 5. Test du Modèle avec une Demande Textuelle ---

# Simule une demande utilisateur
demande_utilisateur = "Je veux une abeille de couleur Violette"
couleur_demandee = "Rouge" # Extraction simple de la couleur

print("-" * 50)
print(f"🤖 Traitement de la demande utilisateur: '{demande_utilisateur}'")
print(f"🔬 Couleur extraite: {couleur_demandee}")
print("-" * 50)

# Génération de l'image
generated_img, _ = generate_conditional_image(couleur_demandee)

if generated_img:
    output_filename = f"generated_{couleur_demandee}_iris1_0.png"
    generated_img.save(output_filename)
    
    print(f"\n✅ Génération réussie! L'image conditionnelle a été sauvegardée sous: {output_filename}")
    print(f"   (Taille: {IMAGE_SIZE[0]}x{IMAGE_SIZE[1]})")

# --- Test de la cohérence: Générer la même image mais avec une couleur différente ---
# Réutilisons le même bruit latent 'z' pour voir la transition de couleur
if generated_img:
    couleur_test = "Rouge"
    
    # Générons un Z initial (bruit)
    _, z_initial = generate_conditional_image("Cyan") # On jette l'image, on garde le Z
    
    print(f"\n🔬 Test d'interpolation: Génération d'un insecte {couleur_test} en utilisant le même 'bruit' (Z).")
    
    # Génération conditionnelle
    test_img, _ = generate_conditional_image(couleur_test, z_sample=z_initial)
    
    if test_img:
        test_output_filename = f"fa{couleur_test}_TEST.png"
        test_img.save(test_output_filename)
        print(f"✅ Test réussi. L'image de test a été sauvegardée sous: {test_output_filename}")
        print("   -> L'insecte devrait avoir la même forme mais une couleur différente (Cyan).")