simplify cide

app.py

@@ -42,27 +42,25 @@ def get_image_gen_pipeline():
 def get_text_gen_pipeline():
     global _text_gen_pipeline
     if _text_gen_pipeline is None:
-         try:
-              device = "cuda" if torch.cuda.is_available() else "cpu"
-              tokenizer = AutoTokenizer.from_pretrained(
-                  "mistralai/Mistral-7B-Instruct-v0.3",
-                  use_fast=True  
-              )
-              # Set pad_token_id to eos_token_id if pad_token is not set
-              if tokenizer.pad_token is None:
-                  tokenizer.pad_token = tokenizer.eos_token
-                  
-              _text_gen_pipeline = pipeline(
-                   "text-generation",
-                   model="mistralai/Mistral-7B-Instruct-v0.3",
-                   tokenizer=tokenizer,
-                   max_new_tokens=2048,
-                   device=device,
-                   pad_token_id=tokenizer.pad_token_id  # Explicitly set pad_token_id
-              )
-         except Exception as e:
-              print(f"Error loading text generation model: {e}")
-              return None
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            tokenizer = AutoTokenizer.from_pretrained(
+                "mistralai/Mistral-7B-Instruct-v0.3",
+                use_fast=True
+            )
+            tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
+            
+            _text_gen_pipeline = pipeline(
+                "text-generation",
+                model="mistralai/Mistral-7B-Instruct-v0.3",
+                tokenizer=tokenizer,
+                max_new_tokens=2048,
+                device=device,
+                pad_token_id=tokenizer.pad_token_id
+            )
+        except Exception as e:
+            print(f"Error loading text generation model: {e}")
+            return None
     return _text_gen_pipeline
 
 @spaces.GPU()
@@ -127,15 +125,10 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
                 max_sequence_length=512
             )
             
-            # Ensure the image is properly normalized and converted
         image = output.images[0]
-        #    if isinstance(image, torch.Tensor):
-        #        image = (image.clamp(-1, 1) + 1) / 2
-        #        image = (image * 255).round().clamp(0, 255).to(torch.uint8).cpu().numpy()
-        #        image = Image.fromarray(image)
 
         #torch.cuda.empty_cache()  
-        return image, seed
+        return image, f"Image generated successfully with seed {seed}"
     except Exception as e:
         print(f"Error in infer: {str(e)}")  
         return None, f"Error generating image: {str(e)}"
@@ -154,150 +147,93 @@ css="""
 """
 
 def preload_models():
-    print("Préchargement des modèles...")
+    global _text_gen_pipeline, _image_gen_pipeline
+    
+    print("Preloading models...")
+    success = True
+    
     try:
-        # Préchargement du modèle de génération de texte
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        # Explicitly load the fast tokenizer LGR
-        tokenizer = AutoTokenizer.from_pretrained(
-            "mistralai/Mistral-7B-Instruct-v0.3",
-            use_fast=True  # Ensures a fast tokenizer is used
-        )
-        _text_gen_pipeline = pipeline(
-            "text-generation",
-            model="mistralai/Mistral-7B-Instruct-v0.3",
-            tokenizer=tokenizer,  # Pass the fast tokenizer in LGR
-            max_new_tokens=2048,
-            device=device,
-        )
-        
-        # Préchargement du modèle de génération d'images
-        dtype = torch.bfloat16
-        _image_gen_pipeline = DiffusionPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-schnell",
-            # "black-forest-labs/FLUX.1-dev",
-            torch_dtype=dtype
-        ).to(device)
-        
-        print("Modèles préchargés avec succès!")
-        return True
+        _text_gen_pipeline = get_text_gen_pipeline()
+        if _text_gen_pipeline is None:
+            success = False
     except Exception as e:
-        print(f"Erreur lors du préchargement des modèles: {str(e)}")
-        return False
+        print(f"Error preloading text generation model: {str(e)}")
+        success = False
+    
+    try:
+        _image_gen_pipeline = get_image_gen_pipeline()
+        if _image_gen_pipeline is None:
+            success = False
+    except Exception as e:
+        print(f"Error preloading image generation model: {str(e)}")
+        success = False
+    
+    status = "Models preloaded successfully!" if success else "Error preloading models"
+    print(status)
+    return success
 
 def create_interface():
-    # Modify the preloading logic
+    # Preload models if needed
     if PRELOAD_MODELS:
         models_loaded = preload_models()
-        model_status = "✅ Modèles chargés avec succès!" if models_loaded else "⚠️ Erreur lors du chargement des modèles"
+        model_status = "✅ Models loaded successfully!" if models_loaded else "⚠️ Error loading models"
     else:
-        model_status = "ℹ️ Modèles seront chargés à la demande"
+        model_status = "ℹ️ Models will be loaded on demand"
 
     with gr.Blocks(css=css) as demo:
-        info = gr.Info(model_status)
+        gr.Info(model_status)
 
         with gr.Column(elem_id="col-container"):
-            gr.Markdown(f"""# Text to Product
-            Using Mistral-7B-Instruct-v0.3 + FLUX.1-dev + Trellis
-            """)
+            gr.Markdown("# Text to Product\nUsing Mistral-7B-Instruct-v0.3 + FLUX.1-dev + Trellis")
             
+            # Basic inputs
             with gr.Row():
-                
                 prompt = gr.Text(
-                    label="Prompt",
                     show_label=False,
                     max_lines=1,
                     placeholder="Enter basic object prompt",
                     container=False,
                 )
-                
-                prompt_button = gr.Button("Refine prompt with Mistral", scale=0)
+                prompt_button = gr.Button("Refine prompt with Mistral")
             
             refined_prompt = gr.Text(
-                label="Refined Prompt", 
                 show_label=False,
                 max_lines=10, 
                 placeholder="Detailed object prompt", 
                 container=False,
                 max_length=2048,
-                )
-            
-            
-            run_button = gr.Button("Create visual with Flux", 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",
-                )
+            visual_button = gr.Button("Create visual with Flux")
+            generated_image = gr.Image(show_label=False)
+            error_box = gr.Textbox(
+                label="Status Messages", 
+                interactive=False,
+                placeholder="Status messages will appear here",
+            )
             
-            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)
-                
-                with gr.Row():
-                    
-                    width = gr.Slider(
-                        label="Width",
-                        minimum=256,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=32,
-                        value=512,
+            # Accordion sections for advanced settings
+            with gr.Accordion("Advanced Settings", open=False):
+                with gr.Tab("Mistral"):
+                    # Mistral settings
+                    temperature = gr.Slider(
+                        label="Temperature",
+                        value=0.9,
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.05,
+                        info="Higher values produce more diverse outputs",
                     )
                     
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=256,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=32,
-                        value=512,
-                    )
-                
-                with gr.Row():
+                with gr.Tab("Flux"):
+                    # Flux settings
+                    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=512)
+                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512)
                     
-      
                     num_inference_steps = gr.Slider(
                         label="Number of inference steps",
                         minimum=1,
@@ -305,29 +241,29 @@ def create_interface():
                         step=1,
                         value=10,
                     )
-            
+
+            # Examples section
             gr.Examples(
                 examples=examples,
                 fn=refine_prompt,
-                inputs = [prompt],
-                outputs = [refined_prompt],
+                inputs=[prompt],
+                outputs=[refined_prompt],
                 cache_examples=True,
-                cache_mode='lazy'
             )
 
-
+        # Event handlers
         gr.on(
             triggers=[prompt_button.click, prompt.submit],
-            fn = refine_prompt,
-            inputs = [prompt],
-            outputs = [refined_prompt]
+            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, prompt]
+            triggers=[visual_button.click],
+            fn=infer,
+            inputs=[refined_prompt, seed, randomize_seed, width, height, num_inference_steps],
+            outputs=[generated_image, error_box]
         )
 
     return demo