preload models

app.py

@@ -141,145 +141,187 @@ css="""
 }
 """
 
-with gr.Blocks(css=css) as demo:
-    
-    # Compute the model loading status message ahead of creating the Info component.
-    model_status = "Models loaded successfully!"
-
-    info = gr.Info(model_status)
-
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# Text to Product
-        Using Mistral + Flux + Trellis
-        """)
-        
-        with gr.Row():
-            
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            
-            prompt_button = gr.Button("Refine prompt", scale=0)
-        
-        refined_prompt = gr.Text(
-            label="Refined Prompt", 
-            show_label=False,
-            max_lines=10, 
-            placeholder="Prompt refined by Mistral", 
-            container=False,
-            max_length=2048,
-            )
+def preload_models():
+    print("Préchargement des modèles...")
+    try:
+        # Préchargement du modèle de génération de texte
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        tokenizer = AutoTokenizer.from_pretrained(
+            "mistralai/Mistral-7B-Instruct-v0.3",
+            use_fast=True
+        )
+        global _text_gen_pipeline
+        _text_gen_pipeline = pipeline(
+            "text-generation",
+            model="mistralai/Mistral-7B-Instruct-v0.3",
+            tokenizer=tokenizer,
+            max_new_tokens=2048,
+            device=device,
+        )
         
+        # Préchargement du modèle de génération d'images
+        dtype = torch.bfloat16
+        global _image_gen_pipeline
+        _image_gen_pipeline = DiffusionPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-schnell",
+            torch_dtype=dtype
+        ).to(device)
         
-        run_button = gr.Button("Create visual", scale=0)
+        print("Modèles préchargés avec succès!")
+        return True
+    except Exception as e:
+        print(f"Erreur lors du préchargement des modèles: {str(e)}")
+        return False
 
-        generated_image = gr.Image(label="Generated Image", show_label=False)
-        
-        with gr.Accordion("Advanced Settings Mistral", open=False):
-            gr.Slider(
-                label="Temperature",
-                value=0.9,
-                minimum=0.0,
-                maximum=1.0,
-                step=0.05,
-                interactive=True,
-                info="Higher values produce more diverse outputs",
-            ),
-            gr.Slider(
-                label="Max new tokens",
-                value=256,
-                minimum=0,
-                maximum=1048,
-                step=64,
-                interactive=True,
-                info="The maximum numbers of new tokens",
-            ),
-            gr.Slider(
-                label="Top-p (nucleus sampling)",
-                value=0.90,
-                minimum=0.0,
-                maximum=1,
-                step=0.05,
-                interactive=True,
-                info="Higher values sample more low-probability tokens",
-            ),
-            gr.Slider(
-                label="Repetition penalty",
-                value=1.2,
-                minimum=1.0,
-                maximum=2.0,
-                step=0.05,
-                interactive=True,
-                info="Penalize repeated tokens",
-            )
+def create_interface():
+    # Préchargement des modèles
+    models_loaded = preload_models()
+    
+    if not models_loaded:
+        model_status = "⚠️ Erreur lors du chargement des modèles"
+    else:
+        model_status = "✅ Modèles chargés avec succès!"
+
+    with gr.Blocks(css=css) as demo:
         
-        with gr.Accordion("Advanced Settings Flux", 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)
+        info = gr.Info(model_status)
+
+        with gr.Column(elem_id="col-container"):
+            gr.Markdown(f"""# Text to Product
+            Using Mistral + Flux + Trellis
+            """)
             
             with gr.Row():
                 
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter your prompt",
+                    container=False,
                 )
                 
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
+                prompt_button = gr.Button("Refine prompt", scale=0)
+            
+            refined_prompt = gr.Text(
+                label="Refined Prompt", 
+                show_label=False,
+                max_lines=10, 
+                placeholder="Prompt refined by Mistral", 
+                container=False,
+                max_length=2048,
                 )
             
-            with gr.Row():
+            
+            run_button = gr.Button("Create visual", scale=0)
+
+            generated_image = gr.Image(label="Generated Image", show_label=False)
+            
+            with gr.Accordion("Advanced Settings Mistral", open=False):
+                gr.Slider(
+                    label="Temperature",
+                    value=0.9,
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.05,
+                    interactive=True,
+                    info="Higher values produce more diverse outputs",
+                ),
+                gr.Slider(
+                    label="Max new tokens",
+                    value=256,
+                    minimum=0,
+                    maximum=1048,
+                    step=64,
+                    interactive=True,
+                    info="The maximum numbers of new tokens",
+                ),
+                gr.Slider(
+                    label="Top-p (nucleus sampling)",
+                    value=0.90,
+                    minimum=0.0,
+                    maximum=1,
+                    step=0.05,
+                    interactive=True,
+                    info="Higher values sample more low-probability tokens",
+                ),
+                gr.Slider(
+                    label="Repetition penalty",
+                    value=1.2,
+                    minimum=1.0,
+                    maximum=2.0,
+                    step=0.05,
+                    interactive=True,
+                    info="Penalize repeated tokens",
+                )
+            
+            with gr.Accordion("Advanced Settings Flux", open=False):
                 
-  
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
                     step=1,
-                    value=4,
+                    value=0,
                 )
-        
-        gr.Examples(
-            examples=examples,
-            fn=infer,
-            inputs=[prompt],
-            outputs=[generated_image, seed],
-            cache_examples=True,
-            cache_mode='lazy'
-        )
+                
+                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=[generated_image, seed],
+                cache_examples=True,
+                cache_mode='lazy'
+            )
 
 
-    gr.on(
-        triggers=[prompt_button.click, prompt.submit],
-        fn = refine_prompt,
-        inputs = [prompt],
-        outputs = [refined_prompt]
-    )
+        gr.on(
+            triggers=[prompt_button.click, prompt.submit],
+            fn = refine_prompt,
+            inputs = [prompt],
+            outputs = [refined_prompt]
+        )
+
+        gr.on(
+            triggers=[run_button.click],
+            fn = infer,
+            inputs = [refined_prompt, seed, randomize_seed, width, height, num_inference_steps],
+            outputs = [generated_image, seed]
+        )
 
-    gr.on(
-        triggers=[run_button.click],
-        fn = infer,
-        inputs = [refined_prompt, seed, randomize_seed, width, height, num_inference_steps],
-        outputs = [generated_image, seed]
-    )
+    return demo
 
-demo.launch()
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()