Fix threading lock

api_server.py

@@ -1,10 +1,6 @@
 #!/usr/bin/env python3
-"""
-API Flask pour générer des images
-Endpoint unique: /generate
-"""
-
 import torch
+import threading
 from pathlib import Path
 from flask import Flask, request, jsonify, send_file
 from diffusers import StableDiffusionPipeline
@@ -13,41 +9,26 @@ import io
 
 app = Flask(__name__)
 
-# Configuration
 OUTPUT_DIR = Path("/app/generated_images")
 OUTPUT_DIR.mkdir(exist_ok=True, parents=True)
-
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Variable globale
 pipeline = None
 model_loaded = False
+lock = threading.Lock()
 
 def load_model():
-    """Charge le modèle au démarrage"""
     global pipeline, model_loaded
-    
-    print("\n" + "="*70)
-    print("🤖 Chargement du modèle fusionné...")
-    print("="*70 + "\n")
-    
     try:
         print(f"📱 Appareil: {DEVICE}")
-        
         dtype = torch.float32 if DEVICE == "cpu" else torch.float16
-        
         pipeline = StableDiffusionPipeline.from_pretrained(
             "K2MAR/mon-modele-sd",
             torch_dtype=dtype,
             safety_checker=None
         ).to(DEVICE)
-        
         pipeline.enable_attention_slicing()
-        
         model_loaded = True
-        print("="*70)
         print("✅ Modèle prêt!")
-        print("="*70 + "\n")
         return True
     except Exception as e:
         print(f"❌ Erreur: {e}")
@@ -58,71 +39,59 @@ def health():
     return jsonify({
         "status": "ok" if model_loaded else "loading",
         "device": DEVICE,
-        "model_loaded": model_loaded
+        "model_loaded": model_loaded,
+        "busy": lock.locked()
     })
 
 @app.route('/generate', methods=['POST'])
 def generate():
     if not model_loaded:
         return jsonify({"error": "Model not loaded"}), 503
-    
+
+    if lock.locked():
+        return jsonify({"error": "Server busy, try again later"}), 503
+
     try:
         data = request.get_json()
-        
         if not data or "prompt" not in data:
-            return jsonify({"error": "Missing 'prompt' in request"}), 400
-        
+            return jsonify({"error": "Missing 'prompt'"}), 400
+
         prompt = data.get("prompt", "")
-        steps = int(data.get("steps", 30))
+        steps = min(int(data.get("steps", 20)), 30)
         guidance_scale = float(data.get("guidance_scale", 7.5))
-        
+
         if not prompt:
             return jsonify({"error": "Prompt cannot be empty"}), 400
-        
-        if steps < 1 or steps > 50:
-            return jsonify({"error": "Steps must be 1-50"}), 400
-        
-        print(f"\n🎨 Génération: {prompt}")
-        
-        with torch.no_grad():
-            image = pipeline(
-                prompt,
-                num_inference_steps=steps,
-                guidance_scale=guidance_scale,
-                height=512,
-                width=512
-            ).images[0]
-        
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        filepath = OUTPUT_DIR / f"generated_{timestamp}.png"
-        image.save(filepath)
-        
+
+        print(f"\n🎨 Génération: {prompt} ({steps} steps)")
+
+        with lock:
+            with torch.no_grad():
+                image = pipeline(
+                    prompt,
+                    num_inference_steps=steps,
+                    guidance_scale=guidance_scale,
+                    height=512,
+                    width=512
+                ).images[0]
+
         img_io = io.BytesIO()
         image.save(img_io, 'PNG')
         img_io.seek(0)
-        
+
+        print(f"✅ Image générée!\n")
         return send_file(img_io, mimetype='image/png')
-    
+
     except Exception as e:
         print(f"❌ Erreur: {str(e)}\n")
         return jsonify({"error": str(e)}), 500
 
 @app.route('/', methods=['GET'])
 def home():
-    return jsonify({
-        "service": "LoRA Solar Panel Generator API",
-        "version": "1.0",
-        "device": DEVICE,
-        "model_loaded": model_loaded,
-        "endpoints": {
-            "health": "GET /health",
-            "generate": "POST /generate"
-        }
-    })
+    return jsonify({"service": "SD API", "model_loaded": model_loaded, "device": DEVICE})
 
 if __name__ == '__main__':
     if not load_model():
         exit(1)
-    
-    print("\n🚀 Serveur démarrage sur 0.0.0.0:7860\n")
+    print("\n🚀 Serveur sur 0.0.0.0:7860\n")
     app.run(host='0.0.0.0', port=7860, debug=False, threaded=True)