pipeline/generator.py

"""
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
  SUPREMEAI — Pipeline de Génération Vidéo Principal
  Combine : Wan2.1 + CogVideoX + AnimateDiff-Lightning + SVD
  Routing intelligent selon GPU disponible et mode choisi
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
"""

import os, sys, time, gc, hashlib, json, logging
from pathlib import Path
from typing import Optional, Callable, Generator
import numpy as np

import torch
import torch.nn.functional as F
from torch.cuda.amp import autocast

logger = logging.getLogger(__name__)

# ── Constantes ──────────────────────────────────────────────────────────────
CACHE_DIR    = Path(os.getenv("SUPREMEAI_CACHE", "/tmp/supremeai_cache"))
OUTPUT_DIR   = Path(os.getenv("SUPREMEAI_OUTPUT", "/tmp/supremeai_output"))
CACHE_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# Modèles disponibles (HuggingFace IDs)
MODELS = {
    "wan2.1":        "Wan-AI/Wan2.1-T2V-14B",
    "wan2.1-fast":   "Wan-AI/Wan2.1-T2V-1.3B",   # version légère
    "cogvideox":     "THUDM/CogVideoX-5b",
    "cogvideox-2b":  "THUDM/CogVideoX-2b",         # version légère
    "hunyuan":       "tencent/HunyuanVideo",
    "svd":           "stabilityai/stable-video-diffusion-img2vid-xt",
    "animatediff":   "ByteDance/AnimateDiff-Lightning",
}


class GPUProfiler:
    """Détecte les capacités GPU et recommande le modèle optimal."""

    @staticmethod
    def detect() -> dict:
        info = {
            "has_cuda":   torch.cuda.is_available(),
            "gpu_name":   "",
            "vram_gb":    0,
            "compute_cap": 0.0,
            "recommended_model": "animatediff",  # fallback CPU-friendly
            "recommended_mode":  "speed",
        }
        if not torch.cuda.is_available():
            info["recommended_model"] = "animatediff"
            info["recommended_mode"]  = "speed"
            return info

        props = torch.cuda.get_device_properties(0)
        info["gpu_name"]    = props.name
        info["vram_gb"]     = props.total_memory / 1e9
        info["compute_cap"] = props.major + props.minor / 10

        vram = info["vram_gb"]
        if vram >= 40:          # A100, H100
            info["recommended_model"] = "wan2.1"
            info["recommended_mode"]  = "quality"
        elif vram >= 24:        # RTX 3090, 4090
            info["recommended_model"] = "cogvideox"
            info["recommended_mode"]  = "quality"
        elif vram >= 16:        # RTX 3080, 4080
            info["recommended_model"] = "cogvideox-2b"
            info["recommended_mode"]  = "balanced"
        elif vram >= 8:         # RTX 3070, 4060 Ti
            info["recommended_model"] = "animatediff"
            info["recommended_mode"]  = "speed"
        else:                   # < 8 Go ou CPU
            info["recommended_model"] = "animatediff"
            info["recommended_mode"]  = "speed"

        logger.info(f"GPU: {info['gpu_name']} | VRAM: {vram:.1f}GB → model: {info['recommended_model']}")
        return info


class PromptEnhancer:
    """Améliore automatiquement les prompts avec des tokens de qualité."""

    QUALITY_TOKENS = (
        "masterpiece, best quality, ultra detailed, sharp focus, "
        "professional cinematography, 8K resolution, HDR, "
    )
    NEGATIVE_BASE = (
        "worst quality, low quality, blurry, motion blur, noise, "
        "watermark, text, deformed, ugly, bad anatomy, flickering, "
        "inconsistent, jitter, artifacts, grain"
    )

    @classmethod
    def enhance(cls, prompt: str, style_enhancer: str = "", negative: str = "") -> tuple[str, str]:
        enhanced_pos = style_enhancer + cls.QUALITY_TOKENS + prompt
        enhanced_neg = cls.NEGATIVE_BASE + (", " + negative if negative else "")
        return enhanced_pos[:500], enhanced_neg[:300]

    @classmethod
    def auto_storyboard(cls, topic: str, n_scenes: int = 5) -> list:
        """Génère un storyboard automatique depuis un sujet."""
        templates = {
            "tutorial": [
                f"Clean desktop screen showing {topic}, cursor highlights menu, smooth zoom",
                f"Step-by-step demonstration of {topic}, highlighted UI elements, annotations",
                f"Close-up of important settings for {topic}, arrow pointing to key button",
                f"Split screen: before and after {topic}, clear visual comparison",
                f"Final result of {topic}, success checkmark animation, clean background",
            ],
            "cinematic": [
                f"Wide establishing shot, {topic}, golden hour lighting, anamorphic",
                f"Medium shot, {topic}, dramatic lighting, shallow depth of field",
                f"Close-up detail, {topic}, macro lens, bokeh background",
                f"Dynamic action shot, {topic}, motion blur, high contrast",
                f"Final reveal shot, {topic}, epic scale, cinematic color grade",
            ],
            "sport": [
                f"Stadium wide shot, crowd cheering, {topic}, dramatic atmosphere",
                f"Player close-up, intense focus, {topic}, slow motion",
                f"Action moment, {topic}, dynamic angle, freeze frame effect",
                f"Celebration moment, {topic}, confetti, crowd reaction",
                f"Highlights recap, {topic}, fast cuts, energetic music sync",
            ],
        }
        # Détecte le type automatiquement
        t = topic.lower()
        if any(k in t for k in ["tutorial", "how to", "guide", "screen", "settings"]):
            scenes = templates["tutorial"]
        elif any(k in t for k in ["sport", "football", "soccer", "game", "match"]):
            scenes = templates["sport"]
        else:
            scenes = templates["cinematic"]
        return scenes[:n_scenes]


class VideoGenerationPipeline:
    """
    Pipeline principal de génération vidéo SupremeAI.
    Routing intelligent entre 4 backends selon GPU + mode.

    Backends :
      A) Wan 2.1 (14B)         → QUALITY  sur GPU ≥ 24GB
      B) CogVideoX-5B          → QUALITY  sur GPU ≥ 16GB
      C) AnimateDiff-Lightning → SPEED    sur GPU ≥ 8GB
      D) MoviePy Enhanced      → CPU/DEMO sans GPU
    """

    def __init__(self, device: str = "auto", offload: bool = True):
        self.device   = self._resolve_device(device)
        self.offload  = offload       # CPU offload pour économiser VRAM
        self.gpu_info = GPUProfiler.detect()
        self._pipe    = None          # pipeline chargé en mémoire
        self._model_id = None         # modèle actuellement chargé
        self.enhancer  = PromptEnhancer()

    def _resolve_device(self, device: str) -> str:
        if device == "auto":
            return "cuda" if torch.cuda.is_available() else "cpu"
        return device

    # ── Chargement lazy des modèles ────────────────────────────────────────
    def _load_wan2(self, fast: bool = False):
        """Charge Wan 2.1 (Flow Matching DiT)."""
        try:
            from diffusers import WanPipeline, WanVideoToVideoPipeline
            model_id = MODELS["wan2.1-fast"] if fast else MODELS["wan2.1"]
            dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
            pipe = WanPipeline.from_pretrained(
                model_id,
                torch_dtype=dtype,
                cache_dir=str(CACHE_DIR),
            )
            if self.offload and self.device == "cuda":
                pipe.enable_model_cpu_offload()
            else:
                pipe = pipe.to(self.device)
            # Flash Attention 2 si disponible
            try:
                pipe.enable_xformers_memory_efficient_attention()
            except Exception:
                pass
            return pipe
        except ImportError as e:
            logger.warning(f"Wan2.1 non disponible: {e}")
            return None

    def _load_cogvideox(self, size: str = "5b"):
        """Charge CogVideoX (Expert DiT avec 3D causal attention)."""
        try:
            from diffusers import CogVideoXPipeline, CogVideoXImageToVideoPipeline
            model_id = MODELS[f"cogvideox{'-2b' if size == '2b' else ''}"]
            dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
            pipe = CogVideoXPipeline.from_pretrained(
                model_id,
                torch_dtype=dtype,
                cache_dir=str(CACHE_DIR),
            )
            if self.offload:
                pipe.enable_model_cpu_offload()
                pipe.enable_sequential_cpu_offload()
            try:
                pipe.vae.enable_tiling()
                pipe.vae.enable_slicing()
            except Exception:
                pass
            return pipe
        except ImportError as e:
            logger.warning(f"CogVideoX non disponible: {e}")
            return None

    def _load_animatediff(self):
        """Charge AnimateDiff-Lightning (4 steps, ultra rapide)."""
        try:
            from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
            from diffusers.utils import load_image
            from huggingface_hub import hf_hub_download
            import safetensors.torch

            adapter = MotionAdapter.from_pretrained(
                "guoyww/animatediff-motion-adapter-v1-5-2",
                torch_dtype=torch.float16,
                cache_dir=str(CACHE_DIR),
            )
            # Charge les poids Lightning (4 steps)
            ckpt_path = hf_hub_download(
                MODELS["animatediff"], "animatediff_lightning_4step_diffusers.safetensors",
                cache_dir=str(CACHE_DIR)
            )
            adapter.load_state_dict(safetensors.torch.load_file(ckpt_path), strict=True)

            pipe = AnimateDiffPipeline.from_pretrained(
                "emilianJR/epiCRealism",
                motion_adapter=adapter,
                torch_dtype=torch.float16,
                cache_dir=str(CACHE_DIR),
            ).to(self.device)

            pipe.scheduler = EulerDiscreteScheduler.from_config(
                pipe.scheduler.config,
                timestep_spacing="trailing",
                beta_schedule="linear"
            )
            try:
                pipe.enable_vae_slicing()
                pipe.enable_xformers_memory_efficient_attention()
            except Exception:
                pass
            return pipe
        except Exception as e:
            logger.warning(f"AnimateDiff-Lightning non disponible: {e}")
            return None

    def _select_backend(self, mode: str) -> str:
        """Sélectionne le backend optimal selon GPU et mode."""
        vram = self.gpu_info.get("vram_gb", 0)
        if mode == "quality":
            if vram >= 24:   return "wan2.1"
            if vram >= 16:   return "cogvideox"
            if vram >= 8:    return "animatediff"
            return "enhanced_moviepy"
        elif mode == "speed":
            if vram >= 6:    return "animatediff"
            return "enhanced_moviepy"
        else:  # balanced
            if vram >= 16:   return "cogvideox"
            if vram >= 8:    return "animatediff"
            return "enhanced_moviepy"

    # ── Génération principale ──────────────────────────────────────────────
    def generate(
        self,
        config,        # VideoGenerationConfig
        progress_cb: Optional[Callable] = None,
    ):
        """
        Point d'entrée principal. Sélectionne automatiquement le meilleur backend.
        Retourne un GenerationResult.
        """
        from core.architecture import GenerationResult

        t_start = time.time()
        backend = self._select_backend(config.mode.value if hasattr(config.mode, "value") else config.mode)
        logger.info(f"Backend sélectionné: {backend}")

        try:
            if backend == "wan2.1":
                result = self._generate_wan2(config, progress_cb)
            elif backend == "cogvideox":
                result = self._generate_cogvideox(config, progress_cb)
            elif backend == "animatediff":
                result = self._generate_animatediff(config, progress_cb)
            else:
                result = self._generate_enhanced_moviepy(config, progress_cb)

            result.generation_time = time.time() - t_start
            result.model_used      = backend
            return result

        except Exception as e:
            logger.error(f"Erreur génération ({backend}): {e}")
            logger.info("Fallback vers enhanced_moviepy...")
            try:
                result = self._generate_enhanced_moviepy(config, progress_cb)
                result.generation_time = time.time() - t_start
                result.model_used      = "enhanced_moviepy_fallback"
                return result
            except Exception as e2:
                return GenerationResult(
                    success=False, video_path=None, preview_path=None,
                    duration=0, fps=0, resolution=(0, 0),
                    model_used=backend, generation_time=time.time()-t_start,
                    error=str(e2)
                )

    def _generate_wan2(self, config, progress_cb=None):
        """Génération via Wan 2.1 (Flow Matching DiT)."""
        from core.architecture import GenerationResult, STYLE_ENHANCERS, NEGATIVE_PROMPTS

        style_enh = STYLE_ENHANCERS.get(config.style, "")
        neg_base  = NEGATIVE_PROMPTS.get(config.style, "")
        pos, neg  = self.enhancer.enhance(config.prompt, style_enh, neg_base)

        if self._model_id != "wan2.1":
            if progress_cb: progress_cb(5, "Chargement Wan 2.1...")
            self._pipe     = self._load_wan2(fast=config.mode.value == "speed")
            self._model_id = "wan2.1"

        if progress_cb: progress_cb(20, "Génération en cours (Wan 2.1 Flow Matching)...")

        seed = config.seed if config.seed >= 0 else torch.randint(0, 2**32, (1,)).item()
        gen  = torch.Generator(device=self.device).manual_seed(seed)

        num_frames = int(config.duration * config.fps)

        with autocast(enabled=self.device == "cuda"):
            output = self._pipe(
                prompt=pos,
                negative_prompt=neg,
                num_frames=num_frames,
                height=config.height,
                width=config.width,
                num_inference_steps=config.num_inference_steps,
                guidance_scale=config.guidance_scale,
                generator=gen,
            )

        if progress_cb: progress_cb(80, "Post-traitement...")
        frames = output.frames[0]
        out_path = self._save_video(frames, config)

        if progress_cb: progress_cb(100, "✅ Terminé !")
        return GenerationResult(
            success=True, video_path=str(out_path), preview_path=None,
            duration=config.duration, fps=config.fps,
            resolution=(config.width, config.height), model_used="wan2.1",
            generation_time=0, metadata={"seed": seed, "steps": config.num_inference_steps}
        )

    def _generate_cogvideox(self, config, progress_cb=None):
        """Génération via CogVideoX (3D causal DiT)."""
        from core.architecture import GenerationResult, STYLE_ENHANCERS, NEGATIVE_PROMPTS

        style_enh = STYLE_ENHANCERS.get(config.style, "")
        neg_base  = NEGATIVE_PROMPTS.get(config.style, "")
        pos, neg  = self.enhancer.enhance(config.prompt, style_enh, neg_base)

        vram = self.gpu_info.get("vram_gb", 0)
        size = "2b" if vram < 20 else "5b"

        if self._model_id != f"cogvideox-{size}":
            if progress_cb: progress_cb(5, f"Chargement CogVideoX-{size.upper()}...")
            self._pipe     = self._load_cogvideox(size)
            self._model_id = f"cogvideox-{size}"

        if progress_cb: progress_cb(20, "Génération (CogVideoX 3D Attention)...")

        seed = config.seed if config.seed >= 0 else torch.randint(0, 2**32, (1,)).item()
        gen  = torch.Generator(device="cpu").manual_seed(seed)

        num_frames = min(int(config.duration * config.fps), 49)  # CogVideoX max 49 frames

        with autocast(enabled=self.device == "cuda"):
            output = self._pipe(
                prompt=pos,
                negative_prompt=neg,
                num_frames=num_frames,
                height=config.height,
                width=config.width,
                num_inference_steps=config.num_inference_steps,
                guidance_scale=config.guidance_scale,
                generator=gen,
            )

        if progress_cb: progress_cb(80, "Post-traitement...")
        from diffusers.utils import export_to_video
        out_path = OUTPUT_DIR / f"supremeai_{int(time.time())}.mp4"
        export_to_video(output.frames[0], str(out_path), fps=config.fps)

        if progress_cb: progress_cb(100, "✅ Terminé !")
        return GenerationResult(
            success=True, video_path=str(out_path), preview_path=None,
            duration=config.duration, fps=config.fps,
            resolution=(config.width, config.height), model_used=f"cogvideox-{size}",
            generation_time=0, metadata={"seed": seed}
        )

    def _generate_animatediff(self, config, progress_cb=None):
        """Génération via AnimateDiff-Lightning (4 steps, ultra rapide)."""
        from core.architecture import GenerationResult, STYLE_ENHANCERS, NEGATIVE_PROMPTS

        style_enh = STYLE_ENHANCERS.get(config.style, "")
        neg_base  = NEGATIVE_PROMPTS.get(config.style, "")
        pos, neg  = self.enhancer.enhance(config.prompt, style_enh, neg_base)

        if self._model_id != "animatediff":
            if progress_cb: progress_cb(5, "Chargement AnimateDiff-Lightning (4 steps)...")
            self._pipe     = self._load_animatediff()
            self._model_id = "animatediff"

        if progress_cb: progress_cb(30, "Génération rapide (4 steps LCM)...")

        seed       = config.seed if config.seed >= 0 else torch.randint(0, 2**32, (1,)).item()
        gen        = torch.Generator(device=self.device).manual_seed(seed)
        num_frames = min(int(config.duration * config.fps), 32)

        output = self._pipe(
            prompt=pos,
            negative_prompt=neg,
            num_frames=num_frames,
            height=min(config.height, 512),   # AnimateDiff limite
            width=min(config.width, 512),
            num_inference_steps=4,            # Lightning = 4 steps
            guidance_scale=1.0,               # LCM = guidance faible
            generator=gen,
        )

        if progress_cb: progress_cb(80, "Export MP4...")
        from diffusers.utils import export_to_video
        out_path = OUTPUT_DIR / f"supremeai_{int(time.time())}.mp4"
        export_to_video(output.frames[0], str(out_path), fps=config.fps)

        if progress_cb: progress_cb(100, "✅ Terminé !")
        return GenerationResult(
            success=True, video_path=str(out_path), preview_path=None,
            duration=config.duration, fps=config.fps,
            resolution=(min(config.width, 512), min(config.height, 512)),
            model_used="animatediff-lightning", generation_time=0,
            metadata={"seed": seed, "steps": 4}
        )

    def _generate_enhanced_moviepy(self, config, progress_cb=None):
        """
        Backend CPU (sans GPU) — Enhanced MoviePy avec effets avancés.
        Utilisé comme fallback et pour les tutoriels informatiques.
        Qualité professionnelle même sans GPU.
        """
        from core.architecture import GenerationResult
        from pipeline.processor import EnhancedVideoProcessor

        if progress_cb: progress_cb(10, "Initialisation moteur vidéo avancé...")
        processor = EnhancedVideoProcessor(config)
        out_path  = processor.render()

        if progress_cb: progress_cb(100, "✅ Vidéo rendue !")
        return GenerationResult(
            success=True, video_path=str(out_path), preview_path=None,
            duration=config.duration, fps=config.fps,
            resolution=(config.width, config.height),
            model_used="enhanced_moviepy", generation_time=0,
        )

    def _save_video(self, frames, config) -> Path:
        """Sauvegarde une liste de frames PIL en MP4."""
        import imageio
        out_path = OUTPUT_DIR / f"supremeai_{int(time.time())}.mp4"
        writer = imageio.get_writer(str(out_path), fps=config.fps, codec="libx264", quality=9)
        for frame in frames:
            writer.append_data(np.array(frame))
        writer.close()
        return out_path

    def generate_director_mode(self, topic: str, n_scenes: int, config, progress_cb=None):
        """
        Mode Director : topic → storyboard auto → génération scène par scène → assemblage.
        Crée une vidéo longue et cohérente automatiquement.
        """
        from moviepy import VideoFileClip, concatenate_videoclips

        storyboard = PromptEnhancer.auto_storyboard(topic, n_scenes)
        clips = []

        for i, scene_prompt in enumerate(storyboard):
            if progress_cb:
                pct = int(i / n_scenes * 90)
                progress_cb(pct, f"Scène {i+1}/{n_scenes}: {scene_prompt[:50]}...")

            scene_cfg = config.__class__(
                prompt=scene_prompt,
                style=config.style,
                mode=config.mode,
                width=config.width, height=config.height,
                fps=config.fps,
                duration=config.duration / n_scenes,
            )
            result = self.generate(scene_cfg)
            if result.success and result.video_path:
                clips.append(VideoFileClip(result.video_path))

        if not clips:
            raise RuntimeError("Aucune scène générée")

        final = concatenate_videoclips(clips, method="compose")
        out_path = OUTPUT_DIR / f"director_{int(time.time())}.mp4"
        final.write_videofile(str(out_path), fps=config.fps, codec="libx264",
                              audio_codec="aac", logger=None)
        if progress_cb: progress_cb(100, "✅ Film généré !")

        from core.architecture import GenerationResult
        return GenerationResult(
            success=True, video_path=str(out_path), preview_path=None,
            duration=config.duration, fps=config.fps,
            resolution=(config.width, config.height),
            model_used="director_mode", generation_time=0,
        )

    def unload(self):
        """Libère la mémoire GPU."""
        self._pipe     = None
        self._model_id = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()