import os
import sys
import subprocess

# --- Installation automatique des paquets manquants lors du démarrage du Space ---
try:
    from mage_flow import MageFlowPipeline
except ImportError:
    print("⏳ Installation du pipeline officiel de Microsoft depuis GitHub...")
    # Installation des dépendances requises
    subprocess.check_call([sys.executable, "-m", "pip", "install", "diffusers==0.38.0", "transformers>=5.3", "einops", "pydantic", "loguru", "accelerate", "gradio"])
    # Installation directe du sous-dossier mage_flow
    subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/microsoft/Mage.git#subdirectory=mage_flow", "--no-deps"])
    from mage_flow import MageFlowPipeline

import torch
import gradio as gr
from PIL import Image

# --- Initialisation du modèle Microsoft Mage-Flow-Turbo ---
MODEL_ID = "microsoft/Mage-Flow-Turbo"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"⏳ Chargement des poids de {MODEL_ID} sur {device.upper()}...")
pipe = MageFlowPipeline.from_pretrained(MODEL_ID, device=device)
print("✅ Pipeline prêt à l'emploi !")

# --- Fonction de génération liée à l'interface ---
def generate_image(prompt):
    if not prompt.strip():
        return None
    
    with torch.inference_mode():
        # Configuration optimisée officielle pour le mode Turbo (4 étapes, CFG 1.0)
        images = pipe.generate(
            [prompt],
            heights=1024,
            widths=1024,
            steps=4,
            cfg=1.0
        )
    return images[0]

# --- Interface Utilisateur Web (Gradio) ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚀 Microsoft Mage-Flow-Turbo Generator")
    gr.Markdown("Générez des images en résolution native (1024x1024) en moins d'une seconde grâce à l'architecture MMDiT de Microsoft.")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Votre Prompt", 
                placeholder="Ex: A beautiful cinematic shot of a neon cyberpunk fox...",
                lines=3
            )
            submit_btn = gr.Button("Générer l'image 🎨", variant="primary")
        
        with gr.Column():
            image_output = gr.Image(label="Image Générée", type="pil")
            
    submit_btn.click(fn=generate_image, inputs=prompt_input, outputs=image_output)

# Lancement de l'application sur le port requis par Hugging Face
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)