Update app.py

app.py

@@ -11,6 +11,9 @@ import time
 from PIL import Image
 import io
 import requests
+from diffusers import FluxPipeline
+
+
 
 # Configuration des modèles disponibles
 TEXT_MODELS = {
@@ -90,14 +93,34 @@ class PresentationGenerator:
                 token=self.token
             )
         
+#    def load_image_model(self, model_name):
+#        """Charge le modèle de génération d'images"""
+#        model_id = IMAGE_MODELS[model_name]
+#        self.image_pipeline = pipeline(
+#            "text-to-image",
+#            model=model_id,
+#            token=self.token
+#        )
+
+##Modif01 : Correction Pour Flux Non Chargé sur HFSpace
     def load_image_model(self, model_name):
         """Charge le modèle de génération d'images"""
         model_id = IMAGE_MODELS[model_name]
-        self.image_pipeline = pipeline(
-            "text-to-image",
-            model=model_id,
-            token=self.token
-        )
+        if model_id == "black-forest-labs/FLUX.1-schnell":
+            self.image_pipeline = FluxPipeline.from_pretrained(
+                model_id,
+                torch_dtype=torch.bfloat16
+            )
+            self.image_pipeline.enable_model_cpu_offload()  # Économise de la VRAM en déchargeant le modèle sur le CPU
+            print(f"Modèle d'image FLUX chargé : {model_id}")
+        else:
+            self.image_pipeline = pipeline(
+                "text-to-image",
+                model=model_id,
+                token=self.token
+            )
+        print(f"Modèle d'image chargé : {model_id}")
+
 
     def generate_text(self, prompt, temperature=0.7, max_tokens=4096):
         """Génère le texte de la présentation"""
@@ -122,18 +145,45 @@ class PresentationGenerator:
             )
             return self.text_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+#    def generate_image(self, prompt, negative_prompt="", num_inference_steps=30):
+#        """Génère une image pour la diapositive"""
+#        try:
+#            image = self.image_pipeline(
+#                prompt=prompt,
+#                negative_prompt=negative_prompt,
+#                num_inference_steps=num_inference_steps
+#            )[0]  # Pipeline retourne une liste d'images, on prend la première
+#            return image
+#        except Exception as e:
+#            print(f"Erreur lors de la génération de l'image: {str(e)}")
+#            return None
+
+
+      ##Modif01 : Correction Pour Flux Non Chargé sur HFSpace
+
     def generate_image(self, prompt, negative_prompt="", num_inference_steps=30):
-        """Génère une image pour la diapositive"""
-        try:
-            image = self.image_pipeline(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                num_inference_steps=num_inference_steps
-            )[0]  # Pipeline retourne une liste d'images, on prend la première
-            return image
-        except Exception as e:
-            print(f"Erreur lors de la génération de l'image: {str(e)}")
-            return None
+        """Génère une image pour la diapositive"""      
+        try:      
+            if isinstance(self.image_pipeline, FluxPipeline):      
+                image = self.image_pipeline(      
+                    prompt=prompt,      
+                    guidance_scale=0.0,      
+                    num_inference_steps=num_inference_steps,      
+                    max_sequence_length=256,      
+                    generator=torch.Generator("cpu").manual_seed(0)      
+                ).images[0]      
+            else:      
+                image = self.image_pipeline(      
+                    prompt=prompt,      
+                    negative_prompt=negative_prompt,      
+                    num_inference_steps=num_inference_steps      
+                )[0]  # Pipeline retourne une liste d'images, on prend la première      
+            return image      
+        except Exception as e:      
+            print(f"Erreur lors de la génération de l'image: {str(e)}")      
+        return None      
+      
+      
 
     def parse_presentation_content(self, content):
         """Parse le contenu généré en sections pour les diapositives"""
@@ -379,5 +429,4 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
-    
+    demo.launch()