#!/usr/bin/env python3
"""
API Flask pour générer des images
Endpoint unique: /generate
"""

import torch
from pathlib import Path
from flask import Flask, request, jsonify, send_file
from diffusers import StableDiffusionPipeline
from datetime import datetime
import io

app = Flask(__name__)

# Configuration
MODEL_PATH = Path("/app/model")
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

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")
    
    if not MODEL_PATH.exists():
        print(f"❌ Erreur: Le modèle n'existe pas à {MODEL_PATH}")
        return False
    
    try:
        print(f"📱 Appareil: {DEVICE}")
        print(f"📁 Modèle: {MODEL_PATH}\n")
        
        print("Chargement...", end=" ", flush=True)
        
        # Sur CPU: toujours float32, sur GPU: float16
        dtype = torch.float32 if DEVICE == "cpu" else torch.float16
        
        pipeline = StableDiffusionPipeline.from_pretrained(
            str(MODEL_PATH),
            torch_dtype=dtype,
            safety_checker=None
        ).to(DEVICE)
        print("✅\n")
        
        print("Optimisations...", end=" ", flush=True)
        pipeline.enable_attention_slicing()
        print("✅\n")
        
        model_loaded = True
        print("="*70)
        print("✅ Modèle prêt!")
        print("="*70 + "\n")
        return True
    except Exception as e:
        print(f"❌ Erreur: {e}")
        return False

@app.route('/health', methods=['GET'])
def health():
    """Health check"""
    return jsonify({
        "status": "ok" if model_loaded else "loading",
        "device": DEVICE,
        "model_loaded": model_loaded
    })

@app.route('/generate', methods=['POST'])
def generate():
    """
    Endpoint unique pour générer une image
    POST /generate
    {
        "prompt": "electrical wiring schematic",
        "steps": 30,
        "guidance_scale": 7.5
    }
    """
    
    if not model_loaded:
        return jsonify({"error": "Model not loaded"}), 503
    
    try:
        data = request.get_json()
        
        if not data or "prompt" not in data:
            return jsonify({"error": "Missing 'prompt' in request"}), 400
        
        prompt = data.get("prompt", "")
        steps = int(data.get("steps", 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...")
        print(f"   Prompt: {prompt}")
        print(f"   Steps: {steps}, Guidance: {guidance_scale}")
        
        with torch.no_grad():
            image = pipeline(
                prompt,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                height=512,
                width=512
            ).images[0]
        
        # Sauvegarder
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"generated_{timestamp}.png"
        filepath = OUTPUT_DIR / filename
        image.save(filepath)
        
        print(f"   ✅ Image générée: {filename}\n")
        
        # Retourner l'image
        img_io = io.BytesIO()
        image.save(img_io, 'PNG')
        img_io.seek(0)
        
        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():
    """Info sur l'API"""
    return jsonify({
        "service": "LoRA Solar Panel Generator API",
        "version": "1.0",
        "device": DEVICE,
        "model_loaded": model_loaded,
        "endpoints": {
            "health": "GET /health",
            "generate": "POST /generate"
        }
    })

if __name__ == '__main__':
    if not load_model():
        exit(1)
    
    # HF Spaces écoute sur le port 7860
    port = 7860
    print(f"\n🚀 Serveur démarrage sur 0.0.0.0:{port}\n")
    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)
#!/usr/bin/env python3
"""
API Flask simple pour générer des images
Utilise UNIQUEMENT le modèle fusionné
"""

import torch
from pathlib import Path
from flask import Flask, request, jsonify, send_file
from diffusers import StableDiffusionPipeline
from datetime import datetime
import io
import os

app = Flask(__name__)

# Configuration
MODEL_PATH = Path("/app/model")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
OUTPUT_DIR = Path("/app/generated_images")
OUTPUT_DIR.mkdir(exist_ok=True, parents=True)

# Variable globale
pipeline = None
model_loaded = False

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")
    
    if not MODEL_PATH.exists():
        print(f"❌ Erreur: Le modèle n'existe pas à {MODEL_PATH}")
        return False
    
    try:
        print(f"📱 Appareil: {DEVICE}")
        print(f"📁 Modèle: {MODEL_PATH}\n")
        
        print("Chargement...", end=" ", flush=True)
        pipeline = StableDiffusionPipeline.from_pretrained(
            str(MODEL_PATH),
            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
            safety_checker=None
        ).to(DEVICE)
        print("✅\n")
        
        print("Optimisations...", end=" ", flush=True)
        pipeline.enable_attention_slicing()
        print("✅\n")
        
        model_loaded = True
        print("="*70)
        print("✅ Modèle prêt!")
        print("="*70 + "\n")
        return True
    except Exception as e:
        print(f"❌ Erreur: {e}")
        return False

@app.route('/health', methods=['GET'])
def health():
    """Health check"""
    return jsonify({
        "status": "ok" if model_loaded else "loading",
        "device": DEVICE,
        "model_loaded": model_loaded
    })

@app.route('/generate', methods=['POST'])
def generate():
    """
    Endpoint unique pour générer une image
    POST /generate
    {
        "prompt": "electrical wiring schematic",
        "steps": 30,
        "guidance_scale": 7.5
    }
    """
    
    if not model_loaded:
        return jsonify({"error": "Model not loaded"}), 503
    
    try:
        data = request.get_json()
        
        if not data or "prompt" not in data:
            return jsonify({"error": "Missing 'prompt' in request"}), 400
        
        prompt = data.get("prompt", "")
        steps = int(data.get("steps", 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...")
        print(f"   Prompt: {prompt}")
        print(f"   Steps: {steps}, Guidance: {guidance_scale}")
        
        with torch.no_grad():
            image = pipeline(
                prompt,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                height=512,
                width=512
            ).images[0]
        
        # Sauvegarder
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"generated_{timestamp}.png"
        filepath = OUTPUT_DIR / filename
        image.save(filepath)
        
        print(f"   ✅ Image générée: {filename}\n")
        
        # Retourner l'image
        img_io = io.BytesIO()
        image.save(img_io, 'PNG')
        img_io.seek(0)
        
        return send_file(img_io, mimetype='image/png')
    
    except Exception as e:
        print(f"❌ Erreur: {str(e)}\n")
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    if not load_model():
        exit(1)
    
    print("\n🚀 Serveur démarrage sur 0.0.0.0:5000\n")
    app.run(host='0.0.0.0', port=5000, debug=False)