Update app.py

app.py

@@ -153,6 +153,10 @@
 
 # demo.queue(api_open=True)
 # demo.launch(debug=True)
+
+
+
+
 import gradio as gr
 import spaces
 from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
@@ -161,36 +165,34 @@ from PIL import Image
 from datetime import datetime
 import os
 
-# subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
 DESCRIPTION = "[Sparrow Qwen2-VL-7B Backend](https://github.com/katanaml/sparrow)"
 
+# ============================================================================
+# IMPORTANT: NE PAS charger le modèle ici (scope global)
+# Le modèle doit être chargé UNIQUEMENT dans la fonction @spaces.GPU
+# ============================================================================
+
+# Variables globales pour le cache (sans charger le modèle)
+_model = None
+_processor = None
+
 
 def array_to_image_path(image_filepath, max_width=1250, max_height=1750):
     if image_filepath is None:
         raise ValueError("No image provided. Please upload an image before submitting.")
 
-    # Open the uploaded image using its filepath
     img = Image.open(image_filepath)
+    input_image_extension = image_filepath.split('.')[-1].lower()
 
-    # Extract the file extension from the uploaded file
-    input_image_extension = image_filepath.split('.')[-1].lower()  # Extract extension from filepath
-
-    # Set file extension based on the original file, otherwise default to PNG
     if input_image_extension in ['jpg', 'jpeg', 'png']:
         file_extension = input_image_extension
     else:
-        file_extension = 'png'  # Default to PNG if extension is unavailable or invalid
+        file_extension = 'png'
 
-    # Get the current dimensions of the image
     width, height = img.size
-
-    # Initialize new dimensions to current size
     new_width, new_height = width, height
 
-    # Check if the image exceeds the maximum dimensions
     if width > max_width or height > max_height:
-        # Calculate the new size, maintaining the aspect ratio
         aspect_ratio = width / height
 
         if width > max_width:
@@ -201,60 +203,38 @@ def array_to_image_path(image_filepath, max_width=1250, max_height=1750):
             new_height = max_height
             new_width = int(new_height * aspect_ratio)
 
-    # Generate a unique filename using timestamp
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     filename = f"image_{timestamp}.{file_extension}"
-
-    # Save the image
     img.save(filename)
-
-    # Get the full path of the saved image
     full_path = os.path.abspath(filename)
 
     return full_path, new_width, new_height
 
 
-# CORRECTION: Ne pas initialiser le modèle dans le scope global
-# À la place, on va le charger dans la fonction décorée avec @spaces.GPU
-
-# Cache global pour éviter de recharger le modèle à chaque appel
-_model_cache = {}
-
-def get_model_and_processor():
+@spaces.GPU
+def run_inference(input_imgs, text_input):
     """
-    Charge le modèle et le processeur une seule fois et les met en cache
-    Cette fonction doit être appelée UNIQUEMENT dans une fonction @spaces.GPU
+    CORRECTION CRITIQUE: Le modèle est chargé ICI, pas dans le scope global
     """
-    if 'model' not in _model_cache:
-        print("Chargement du modèle Qwen2-VL-7B-Instruct...")
-        _model_cache['model'] = Qwen2VLForConditionalGeneration.from_pretrained(
+    global _model, _processor
+    
+    # Charger le modèle uniquement la première fois (lazy loading)
+    if _model is None or _processor is None:
+        print("🔄 Chargement du modèle Qwen2-VL-7B-Instruct...")
+        _model = Qwen2VLForConditionalGeneration.from_pretrained(
             "Qwen/Qwen2-VL-7B-Instruct",
             torch_dtype="auto",
             device_map="auto"
         )
-        _model_cache['processor'] = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
-        print("Modèle chargé avec succès!")
-    
-    return _model_cache['model'], _model_cache['processor']
-
-
-@spaces.GPU
-def run_inference(input_imgs, text_input):
-    """
-    CORRECTION: Le modèle est maintenant chargé ICI, à l'intérieur de la fonction @spaces.GPU
-    Cela évite l'erreur "CUDA must not be initialized in the main process"
-    """
-    # Charger le modèle et le processeur dans la fonction GPU
-    model, processor = get_model_and_processor()
+        _processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
+        print("✅ Modèle chargé avec succès!")
     
     results = []
 
     for image in input_imgs:
-        # Convert each image to the required format
         image_path, width, height = array_to_image_path(image)
 
         try:
-            # Prepare messages for each image
             messages = [
                 {
                     "role": "user",
@@ -273,13 +253,12 @@ def run_inference(input_imgs, text_input):
                 }
             ]
 
-            # Prepare inputs for the model
-            text = processor.apply_chat_template(
+            text = _processor.apply_chat_template(
                 messages, tokenize=False, add_generation_prompt=True
             )
 
             image_inputs, video_inputs = process_vision_info(messages)
-            inputs = processor(
+            inputs = _processor(
                 text=[text],
                 images=image_inputs,
                 videos=video_inputs,
@@ -288,23 +267,24 @@ def run_inference(input_imgs, text_input):
             )
             inputs = inputs.to("cuda")
 
-            # Generate inference output
-            generated_ids = model.generate(**inputs, max_new_tokens=4096)
+            generated_ids = _model.generate(**inputs, max_new_tokens=4096)
             generated_ids_trimmed = [
                 out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
             ]
-            raw_output = processor.batch_decode(
+            raw_output = _processor.batch_decode(
                 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=True
             )
 
             results.append(raw_output[0])
-            print("Processed: " + image)
+            print(f"✅ Processed: {image}")
+        except Exception as e:
+            print(f"❌ Error processing {image}: {str(e)}")
+            results.append(f"Error: {str(e)}")
         finally:
-            # Clean up the temporary image file
             if os.path.exists(image_path):
                 os.remove(image_path)
 
-    return results
+    return "\n\n---\n\n".join(results)
 
 
 css = """
@@ -321,10 +301,10 @@ with gr.Blocks(css=css) as demo:
         with gr.Row():
             with gr.Column():
                 input_imgs = gr.Files(file_types=["image"], label="Upload Document Images")
-                text_input = gr.Textbox(label="Query")
+                text_input = gr.Textbox(label="Query", placeholder="Enter your query here...")
                 submit_btn = gr.Button(value="Submit", variant="primary")
             with gr.Column():
-                output_text = gr.Textbox(label="Response")
+                output_text = gr.Textbox(label="Response", elem_id="output")
 
         submit_btn.click(run_inference, [input_imgs, text_input], [output_text])