Update app.py

app.py

@@ -1,6 +1,4 @@
-###TEST03 JUSTE CHARGER FLUX-SCHNELL
-###https://huggingface.co/spaces/black-forest-labs/FLUX.1-schnell/blob/main/app.py
-###
+##TEST125 J'en peu plus de FLUX
 import os
 import gradio as gr
 from huggingface_hub import login
@@ -8,186 +6,18 @@ from diffusers import FluxPipeline
 import torch
 from PIL import Image
 import fitz  # PyMuPDF pour la gestion des PDF
-import sentencepiece
-
-import numpy as np
-import random
-import spaces
-
-
-
-
-
-#
-#import gradio as gr
-#import numpy as np
-#import random
-#import spaces
-#import torch
-#from diffusers import DiffusionPipeline
-#
-#dtype = torch.bfloat16
-#device = "cuda" if torch.cuda.is_available() else "cpu"
-#
-#pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
-#
-#MAX_SEED = np.iinfo(np.int32).max
-#MAX_IMAGE_SIZE = 2048
-#
-#@spaces.GPU()
-#def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
-#    if randomize_seed:
-#        seed = random.randint(0, MAX_SEED)
-#    generator = torch.Generator().manual_seed(seed)
-#    image = pipe(
-#            prompt = prompt, 
-#            width = width,
-#            height = height,
-#            num_inference_steps = num_inference_steps, 
-#            generator = generator,
-#            guidance_scale=0.0
-#    ).images[0] 
-#    return image, seed
-# 
-#examples = [
-#    "a tiny astronaut hatching from an egg on the moon",
-#    "a cat holding a sign that says hello world",
-#    "an anime illustration of a wiener schnitzel",
-#]
-#
-#css="""
-##col-container {
-#    margin: 0 auto;
-#    max-width: 520px;
-#}
-#"""
-#
-#with gr.Blocks(css=css) as demo:
-#    
-#    with gr.Column(elem_id="col-container"):
-#        gr.Markdown(f"""# FLUX.1 [schnell]
-#12B param rectified flow transformer distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/) for 4 step generation
-#[[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-schnell)]
-#        """)
-#        
-#        with gr.Row():
-#            
-#            prompt = gr.Text(
-#                label="Prompt",
-#                show_label=False,
-#                max_lines=1,
-#                placeholder="Enter your prompt",
-#                container=False,
-#            )
-#            
-#            run_button = gr.Button("Run", scale=0)
-#        
-#        result = gr.Image(label="Result", show_label=False)
-#        
-#        with gr.Accordion("Advanced Settings", open=False):
-#            
-#            seed = gr.Slider(
-#                label="Seed",
-#                minimum=0,
-#                maximum=MAX_SEED,
-#                step=1,
-#                value=0,
-#            )
-#            
-#            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-#            
-#            with gr.Row():
-#                
-#                width = gr.Slider(
-#                    label="Width",
-#                    minimum=256,
-#                    maximum=MAX_IMAGE_SIZE,
-#                    step=32,
-#                    value=1024,
-#                )
-#                
-#                height = gr.Slider(
-#                    label="Height",
-#                    minimum=256,
-#                    maximum=MAX_IMAGE_SIZE,
-#                    step=32,
-#                    value=1024,
-#                )
-#            
-#            with gr.Row():
-#                
-#  
-#                num_inference_steps = gr.Slider(
-#                    label="Number of inference steps",
-#                    minimum=1,
-#                    maximum=50,
-#                    step=1,
-#                    value=4,
-#                )
-#        
-#        gr.Examples(
-#            examples = examples,
-#            fn = infer,
-#            inputs = [prompt],
-#            outputs = [result, seed],
-#            cache_examples="lazy"
-#        )
-#
-#    gr.on(
-#        triggers=[run_button.click, prompt.submit],
-#        fn = infer,
-#        inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps],
-#        outputs = [result, seed]
-#    )
-#
-#demo.launch()
-#
-#
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-# Force l'utilisation du CPU pour tout PyTorch
-#torch.set_default_device("cpu")
-
-
-#dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-#
-#pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
-
-
-
-
-
-
-
-
-
+import gc  # Pour le garbage collector
+import psutil  # Pour monitorer la mémoire
 
+# Configuration globale pour réduire l'utilisation de la mémoire
+torch.set_default_device("cpu")
+torch.set_num_threads(2)  # Limite le nombre de threads CPU
+torch.set_grad_enabled(False)  # Désactive complètement le calcul des gradients
 
+def get_memory_usage():
+    """Retourne l'utilisation actuelle de la mémoire en GB"""
+    process = psutil.Process(os.getpid())
+    return process.memory_info().rss / 1024 / 1024 / 1024  # Conversion en GB
 
 def load_pdf(pdf_path):
     """Traite le texte d'un fichier PDF"""
@@ -211,69 +41,90 @@ class FluxGenerator:
             raise ValueError("Token d'authentification HuggingFace non trouvé")
         login(self.token)
         self.pipeline = None
-        self.device = "cpu"  # Force l'utilisation du CPU
+        self.device = "cpu"
         self.load_model()
 
     def load_model(self):
-        """Charge le modèle FLUX avec des paramètres optimisés pour CPU"""
+        """Charge le modèle FLUX avec des paramètres optimisés pour faible mémoire"""
         try:
-            print("Chargement du modèle FLUX sur CPU...")
-            # Configuration spécifique pour CPU
-            torch.set_grad_enabled(False)  # Désactive le calcul des gradients
-            
+            print("Chargement du modèle FLUX avec optimisations mémoire...")
+            print(f"Mémoire utilisée avant chargement: {get_memory_usage():.2f} GB")
+
+            # Configuration pour minimiser l'utilisation de la mémoire
+            model_kwargs = {
+                "low_cpu_mem_usage": True,
+                "torch_dtype": torch.float16,  # Utilise float16 pour réduire la mémoire
+                "use_safetensors": True,  # Utilise safetensors pour un chargement plus efficace
+            }
+
             self.pipeline = FluxPipeline.from_pretrained(
                 "black-forest-labs/FLUX.1-schnell",
                 revision="refs/pr/1",
-                torch_dtype=torch.float32  # Utilise float32 au lieu de bfloat16 pour meilleure compatibilité CPU
+                device_map={"": self.device},
+                **model_kwargs
             )
-#                device_map={"cpu": self.device}  # Force tous les composants sur CPU
-#            )device
-            
-            # Désactive les optimisations GPU
-            self.pipeline.to(self.device)
-            print(f"Utilisation forcée du CPU")
-            print("Modèle FLUX chargé avec succès!")
-            
+
+            # Optimisations supplémentaires
+            self.pipeline.enable_sequential_cpu_offload()  # Décharge les composants non utilisés
+            self.pipeline.enable_attention_slicing(slice_size=1)  # Réduit l'utilisation de la mémoire pendant l'inférence
+
+            # Force le garbage collector
+            gc.collect()
+            torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+            print(f"Mémoire utilisée après chargement: {get_memory_usage():.2f} GB")
+            print("Modèle FLUX chargé avec succès en mode basse consommation!")
+
         except Exception as e:
             print(f"Erreur lors du chargement du modèle: {str(e)}")
             raise
 
     def generate_image(self, prompt, reference_image=None, pdf_file=None):
-        """Génère une image à partir d'un prompt et optionnellement une référence"""
+        """Génère une image avec paramètres optimisés pour la mémoire"""
         try:
-            # Si un PDF est fourni, ajoute son contenu au prompt
             if pdf_file is not None:
                 pdf_text = load_pdf(pdf_file)
                 if pdf_text:
                     prompt = f"{prompt}\nContexte du PDF:\n{pdf_text}"
 
-            # Configuration pour génération sur CPU
-            with torch.no_grad():  # Désactive le calcul des gradients pendant la génération
+            # Paramètres optimisés pour réduire l'utilisation de la mémoire
+            with torch.no_grad():
                 image = self.pipeline(
                     prompt=prompt,
-                    num_inference_steps=20,  # Réduit le nombre d'étapes pour accélérer sur CPU
+                    num_inference_steps=4,  # Minimum d'étapes pour économiser la mémoire
+                    height=512,  # Taille réduite
+                    width=512,   # Taille réduite
                     guidance_scale=0.0,
-                    max_sequence_length=256,
+                    max_sequence_length=128,  # Réduit la longueur de séquence
                     generator=torch.Generator(device=self.device).manual_seed(0)
                 ).images[0]
 
-            return image
+                # Force le nettoyage de la mémoire après génération
+                gc.collect()
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+                return image
 
         except Exception as e:
             print(f"Erreur lors de la génération de l'image: {str(e)}")
             return None
 
 # Instance globale du générateur
-generator = FluxGenerator()
+generator = None  # On initialise plus tard pour économiser la mémoire
 
 def generate(prompt, reference_file):
     """Fonction de génération pour l'interface Gradio"""
+    global generator
     try:
+        # Initialisation tardive du générateur
+        if generator is None:
+            generator = FluxGenerator()
+
         # Gestion du fichier de référence
         if reference_file is not None:
-            if isinstance(reference_file, dict):  # Si le fichier est fourni par Gradio
+            if isinstance(reference_file, dict):
                 file_path = reference_file.name
-            else:  # Si c'est un chemin direct
+            else:
                 file_path = reference_file
                 
             file_type = file_path.split('.')[-1].lower()
@@ -282,14 +133,13 @@ def generate(prompt, reference_file):
             elif file_type in ['png', 'jpg', 'jpeg']:
                 return generator.generate_image(prompt, reference_image=file_path)
 
-        # Génération sans référence
         return generator.generate_image(prompt)
 
     except Exception as e:
         print(f"Erreur détaillée: {str(e)}")
         return None
 
-# Interface Gradio simple
+# Interface Gradio minimaliste
 demo = gr.Interface(
     fn=generate,
     inputs=[
@@ -297,8 +147,8 @@ demo = gr.Interface(
         gr.File(label="Image ou PDF de référence (optionnel)", type="file")
     ],
     outputs=gr.Image(label="Image générée"),
-    title="Test du modèle FLUX (CPU)",
-    description="Interface simple pour tester la génération d'images avec FLUX (optimisé pour CPU)"
+    title="FLUX (Mode économique)",
+    description="Génération d'images optimisée pour systèmes à ressources limitées"
 )
 
 if __name__ == "__main__":