import streamlit as st import pandas as pd import os from datetime import datetime import csv import io import chardet import zipfile import tempfile import requests from PIL import Image from io import BytesIO import concurrent.futures import threading from queue import Queue import time import uuid import logging import shutil # Configuration du logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Paramètres API_URL = "https://router.huggingface.co/hf-inference/models/black-forest-labs/FLUX.1-dev" HF_API_KEY = os.getenv('HF_API_KEY') if not HF_API_KEY: st.error("Erreur: Le token API Hugging Face n'est pas configuré. Veuillez définir la variable d'environnement HF_API_KEY.") st.stop() # Configuration des headers pour l'API Hugging Face headers = { "Authorization": f"Bearer {HF_API_KEY}", } # File d'attente pour les résultats results_queue = Queue() # File d'attente pour les mises à jour de statut des threads status_queue = Queue() # File d'attente pour les messages d'erreur error_queue = Queue() # Compteur pour suivre la progression progress_counter = 0 progress_lock = threading.Lock() # Dictionnaire pour suivre l'état des threads thread_status = {} thread_status_lock = threading.Lock() # Chemin du dossier de sortie et du ZIP output_dir = None zip_path = None zip_lock = threading.Lock() def detect_encoding(csv_content): """Détecte l'encodage du contenu CSV.""" result = chardet.detect(csv_content) return result['encoding'] def clean_csv_content(csv_content): """Nettoie le contenu CSV en supprimant les virgules à l'intérieur des guillemets.""" # Détection de l'encodage encoding = detect_encoding(csv_content) st.info(f"Encodage détecté : {encoding}") try: # Tentative de décodage avec l'encodage détecté decoded_content = csv_content.decode(encoding) except UnicodeDecodeError: # Si échec, essayer avec latin-1 qui peut lire tous les octets st.warning("Impossible de décoder avec l'encodage détecté, utilisation de latin-1") decoded_content = csv_content.decode('latin-1') output = io.StringIO() reader = csv.reader(io.StringIO(decoded_content), delimiter=';') writer = csv.writer(output) for row in reader: cleaned_row = [] for cell in row: # Supprime les virgules à l'intérieur des guillemets if cell.startswith('"') and cell.endswith('"'): cell = cell.replace(',', '') cleaned_row.append(cell) writer.writerow(cleaned_row) return output.getvalue().encode('utf-8') def update_thread_status(thread_id, status, prompt=None): """Met à jour le statut d'un thread.""" with thread_status_lock: if thread_id not in thread_status: thread_status[thread_id] = { "status": status, "prompt": prompt, "start_time": time.time(), "end_time": None } else: thread_status[thread_id]["status"] = status if prompt: thread_status[thread_id]["prompt"] = prompt # Si le statut n'est plus "En cours", enregistrer le temps de fin if status != "En cours" and thread_status[thread_id]["end_time"] is None: thread_status[thread_id]["end_time"] = time.time() # Ajout du statut à la file d'attente pour l'affichage status_queue.put((thread_id, status, prompt)) def generate_image(prompt, idx, total, thread_id): """Génère une image à partir d'un prompt et met à jour la progression.""" try: # Mise à jour du statut du thread update_thread_status(thread_id, "En cours", prompt) # Utilisation de l'API Hugging Face pour la génération d'images response = requests.post(API_URL, headers=headers, json={"inputs": prompt}) if response.status_code == 200: # Mise à jour de la progression with progress_lock: global progress_counter progress_counter += 1 progress = progress_counter / total # Mise à jour du statut du thread update_thread_status(thread_id, "Terminé", prompt) # Ajout du résultat à la file d'attente results_queue.put((idx, response.content, prompt, thread_id)) return True else: # Mise à jour du statut du thread error_msg = f"Erreur lors de la génération de l'image {idx+1}: {response.status_code}" update_thread_status(thread_id, f"Erreur: {response.status_code}", prompt) # Utilisation de logging au lieu de st.error logger.error(error_msg) # Ajout de l'erreur à la file d'attente error_queue.put((idx, error_msg)) return False except Exception as e: # Mise à jour du statut du thread error_msg = f"Erreur lors de la génération de l'image {idx+1}: {str(e)}" update_thread_status(thread_id, f"Erreur: {str(e)}", prompt) # Utilisation de logging au lieu de st.error logger.error(error_msg) # Ajout de l'erreur à la file d'attente error_queue.put((idx, error_msg)) return False def save_image(image_bytes, output_path): """Sauvegarde une image et met à jour le ZIP.""" if image_bytes is None: return False try: # Sauvegarde de l'image image = Image.open(BytesIO(image_bytes)) image.save(output_path) # Mise à jour du ZIP with zip_lock: global zip_path if zip_path and os.path.exists(zip_path): # Supprimer l'ancien ZIP os.remove(zip_path) # Créer un nouveau ZIP avec toutes les images if os.path.exists(os.path.dirname(output_path)): with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: for root, dirs, files in os.walk(os.path.dirname(output_path)): for file in files: if file.endswith('.png'): file_path = os.path.join(root, file) arcname = os.path.relpath(file_path, os.path.dirname(output_path)) zipf.write(file_path, arcname) return True except Exception as e: error_msg = f"Erreur lors de la sauvegarde de l'image: {str(e)}" logger.error(error_msg) error_queue.put((-1, error_msg)) return False def create_zip_file(directory_path): """Crée un fichier ZIP à partir du dossier d'images.""" zip_path = f"{directory_path}.zip" with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: for root, dirs, files in os.walk(directory_path): for file in files: if file.endswith('.png'): file_path = os.path.join(root, file) arcname = os.path.relpath(file_path, directory_path) zipf.write(file_path, arcname) return zip_path def display_thread_status(thread_status_container): """Affiche le statut de tous les threads.""" with thread_status_lock: # Création d'un DataFrame pour afficher les statuts status_data = [] for thread_id, info in thread_status.items(): # Calcul du temps écoulé if info["end_time"] is not None: # Si le thread a terminé, utiliser le temps de fin elapsed_time = info["end_time"] - info["start_time"] else: # Sinon, calculer le temps écoulé jusqu'à maintenant elapsed_time = time.time() - info["start_time"] status_data.append({ "Thread ID": thread_id[:8], # Afficher seulement les 8 premiers caractères "Statut": info["status"], "Prompt": info["prompt"] if info["prompt"] else "N/A", "Temps écoulé": f"{elapsed_time:.1f}s" }) if status_data: df = pd.DataFrame(status_data) thread_status_container.dataframe(df, use_container_width=True) else: thread_status_container.info("Aucun thread en cours d'exécution") def process_results(total_images, progress_bar, image_placeholder, thread_status_container, error_container, download_button_placeholder): """Traite les résultats de la file d'attente et met à jour l'interface.""" success_count = 0 processed_count = 0 errors = [] # Boucle principale pour traiter les résultats et mettre à jour l'interface while processed_count < total_images: # Traitement des mises à jour de statut while not status_queue.empty(): thread_id, status, prompt = status_queue.get() # L'affichage sera mis à jour dans la prochaine itération # Traitement des erreurs while not error_queue.empty(): idx, error_msg = error_queue.get() errors.append(error_msg) # Mise à jour de l'affichage des erreurs error_container.error("\n".join(errors)) # Affichage du statut des threads display_thread_status(thread_status_container) # Traitement des résultats d'images if not results_queue.empty(): idx, image_bytes, prompt, thread_id = results_queue.get() # Sauvegarde de l'image output_path = os.path.join(output_dir, f"image_{idx}.png") if save_image(image_bytes, output_path): success_count += 1 # Affichage de l'image générée with image_placeholder.container(): st.image(image_bytes, caption=f"Image {idx + 1}/{total_images} - Prompt: {prompt}") # Mise à jour du bouton de téléchargement with download_button_placeholder.container(): if os.path.exists(zip_path): with open(zip_path, 'rb') as f: st.download_button( label="Télécharger toutes les images (ZIP)", data=f, file_name=os.path.basename(zip_path), mime="application/zip" ) processed_count += 1 progress_bar.progress(processed_count / total_images) else: time.sleep(0.1) # Petit délai pour éviter de surcharger le CPU return success_count def main(): st.title("Générateur d'images à partir de CSV avec Hugging Face (Multi-thread) 🧠🎨") # Variables globales global output_dir, zip_path # Upload du fichier CSV uploaded_file = st.file_uploader("Choisissez un fichier CSV", type=['csv']) if uploaded_file is not None: try: # Récupération du nom du fichier CSV (sans extension) csv_filename = os.path.splitext(uploaded_file.name)[0] # Lecture et nettoyage du CSV st.info("Étape 1/5 : Nettoyage du fichier CSV...") csv_content = uploaded_file.read() cleaned_csv_content = clean_csv_content(csv_content) # Lecture du CSV nettoyé df = pd.read_csv(io.BytesIO(cleaned_csv_content)) # Ignorer la première ligne (en-tête) df = df.iloc[1:].reset_index(drop=True) # Création du dossier de sortie avec le nom du fichier CSV st.info("Étape 2/5 : Création du dossier de sortie...") timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") # Création d'un dossier permanent dans le répertoire courant output_dir = f"{csv_filename}_{timestamp}" os.makedirs(output_dir, exist_ok=True) # Création du fichier ZIP zip_path = f"{output_dir}.zip" # Affichage des colonnes du CSV st.info("Étape 3/5 : Affichage des colonnes disponibles...") st.write("Colonnes disponibles dans le CSV :") st.write(df.columns.tolist()) # Sélection de la colonne contenant les prompts prompt_column = st.selectbox("Sélectionnez la colonne contenant les prompts", df.columns.tolist()) # Sélection du nombre de threads num_threads = st.slider("Nombre de threads parallèles", min_value=1, max_value=10, value=3) if st.button("Générer les images"): st.info(f"Étape 4/5 : Début de la génération des images avec Hugging Face (utilisant {num_threads} threads)...") # Création des éléments d'interface progress_bar = st.progress(0) image_placeholder = st.empty() # Conteneur pour l'état des threads st.subheader("État des threads") thread_status_container = st.empty() # Conteneur pour les erreurs error_container = st.empty() # Conteneur pour le bouton de téléchargement download_button_placeholder = st.empty() # Réinitialisation du compteur de progression et des statuts global progress_counter progress_counter = 0 with thread_status_lock: thread_status.clear() # Création du thread pool with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor: # Soumission des tâches futures = [] for idx, row in df.iterrows(): prompt = row[prompt_column] thread_id = str(uuid.uuid4()) # Génération d'un ID unique pour chaque thread futures.append(executor.submit(generate_image, prompt, idx, len(df), thread_id)) # Traitement des résultats en parallèle success_count = process_results(len(df), progress_bar, image_placeholder, thread_status_container, error_container, download_button_placeholder) if success_count > 0: st.success(f"Étape 5/5 : {success_count} images ont été générées avec succès sur {len(df)} tentatives") st.info(f"Les images ont été sauvegardées dans le dossier '{output_dir}' et le fichier ZIP '{zip_path}'") else: st.error("Aucune image n'a pu être générée. Veuillez vérifier votre token API Hugging Face et réessayer.") except Exception as e: st.error(f"Une erreur est survenue : {str(e)}") if __name__ == "__main__": main()