vnassets/vnassets/session.py

"""Persistent model session — SDXL and Qwen Image Edit held in GPU memory.

Models are loaded eagerly at construction and reused across generate()/edit()
calls. On 128 GB unified memory (Strix Halo), everything fits simultaneously.
"""
import gc
import json
import random
import time
from pathlib import Path

import safetensors.torch
import torch
from accelerate import init_empty_weights
from PIL import Image
from diffusers import (
    FlowMatchEulerDiscreteScheduler,
    QwenImageEditPlusPipeline,
    StableDiffusionXLPipeline,
)
from diffusers.models.autoencoders import AutoencoderKLQwenImage
from diffusers.models.transformers import QwenImageTransformer2DModel
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor

from .attention import patch_qwen_transformer, patch_unet_attention
from .prompt import build_compel, encode_prompts

TEXT_ENCODER_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
VAE_ID = "Qwen/Qwen-Image"


class VnAssetsSession:
    """Holds SDXL and/or Qwen Image Edit models in GPU memory for reuse.

    Usage as context manager::

        with VnAssetsSession(
            sdxl_checkpoint="models/novaAnimeXL.safetensors",
            edit_model="models/qwen_image_edit.safetensors",
            edit_lora="models/lightning-4steps.safetensors",
        ) as vna:
            vna.generate("1girl, red hair", output="base.png")
            vna.edit("base.png", "make her smile", output="happy.png")
            vna.edit("base.png", "make her sad", output="sad.png")

    Or manual lifecycle::

        vna = VnAssetsSession(sdxl_checkpoint=...)
        vna.generate(...)
        vna.close()
    """

    def __init__(
        self,
        sdxl_checkpoint: str | None = None,
        edit_model: str | None = None,
        edit_lora: str | None = None,
    ):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        # bfloat16 avoids ROCm kernel crashes on RDNA 3.5; float16 segfaults
        self.dtype = torch.bfloat16

        self._sdxl_checkpoint = sdxl_checkpoint
        self._edit_model = edit_model
        self._edit_lora = edit_lora
        self._lora_fused = False
        self._lora_load_s: float | None = None

        self._pipe_sdxl: StableDiffusionXLPipeline | None = None
        self._compel = None
        self._pipe_qwen: QwenImageEditPlusPipeline | None = None

        t0 = time.perf_counter()
        if sdxl_checkpoint:
            self._load_sdxl(sdxl_checkpoint)
        if edit_model:
            self._load_qwen(edit_model, edit_lora)
        self._load_time_s = round(time.perf_counter() - t0, 2)

        loaded = []
        if self._pipe_sdxl:
            loaded.append("SDXL")
        if self._pipe_qwen:
            loaded.append("Qwen")
        if loaded:
            print(f"Session ready ({'+'.join(loaded)}, {self._load_time_s}s)")

    # ── SDXL ────────────────────────────────────────────────────────────

    def _load_sdxl(self, checkpoint_path: str) -> None:
        pipe = StableDiffusionXLPipeline.from_single_file(
            checkpoint_path,
            torch_dtype=self.dtype,
        )
        pipe.to(self.device)
        patch_unet_attention(pipe.unet)
        self._pipe_sdxl = pipe
        self._compel = build_compel(pipe)

    # ── Qwen Image Edit ─────────────────────────────────────────────────

    def _load_qwen(self, model_path: str, lora_path: str | None) -> None:
        transformer = self._load_transformer(model_path)
        vae = AutoencoderKLQwenImage.from_pretrained(
            VAE_ID, subfolder="vae", torch_dtype=self.dtype
        )
        text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            TEXT_ENCODER_ID, torch_dtype=self.dtype
        )
        tokenizer = Qwen2Tokenizer.from_pretrained(TEXT_ENCODER_ID)
        processor = Qwen2VLProcessor.from_pretrained(TEXT_ENCODER_ID)
        scheduler = FlowMatchEulerDiscreteScheduler()

        pipe = QwenImageEditPlusPipeline(
            scheduler=scheduler,
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            processor=processor,
            transformer=transformer,
        )
        pipe.to(self.device)
        patch_qwen_transformer(transformer)

        if lora_path:
            t_lora = time.perf_counter()
            pipe.load_lora_weights(lora_path)
            pipe.fuse_lora(lora_scale=1.0, components=["transformer"])
            self._lora_fused = True
            self._lora_load_s = round(time.perf_counter() - t_lora, 2)
            print(f"LoRA loaded + fused: {self._lora_load_s}s")

        self._pipe_qwen = pipe

    def _load_transformer(self, path: str) -> QwenImageTransformer2DModel:
        """Load Qwen Image Edit transformer from a single FP8 safetensors file.

        Uses init_empty_weights and incremental conversion to keep peak memory
        manageable. The model is 20B parameters (20 GB FP8, 40 GB BF16).
        """
        config = QwenImageTransformer2DModel.load_config(
            "Qwen/Qwen-Image-Edit", subfolder="transformer"
        )
        state_dict = safetensors.torch.load_file(path)
        prefix = "model.diffusion_model."

        # Convert FP8 -> target dtype, freeing FP8 tensors as we go
        cleaned = {}
        for k in list(state_dict.keys()):
            if k.startswith(prefix):
                v = state_dict.pop(k)
                cleaned[k[len(prefix):]] = v.to(self.dtype)
                del v
        del state_dict
        gc.collect()

        with init_empty_weights():
            model = QwenImageTransformer2DModel.from_config(config, torch_dtype=self.dtype)

        model.load_state_dict(cleaned, strict=True, assign=True)
        del cleaned
        gc.collect()
        return model

    # ── Properties ──────────────────────────────────────────────────────

    @property
    def load_time_s(self) -> float:
        """Total time spent loading models at session construction (seconds)."""
        return self._load_time_s

    @property
    def has_sdxl(self) -> bool:
        return self._pipe_sdxl is not None

    @property
    def has_qwen(self) -> bool:
        return self._pipe_qwen is not None

    # ── Generate (SDXL text-to-image) ───────────────────────────────────

    def generate(
        self,
        prompt: str,
        negative_prompt: str = "",
        width: int = 1024,
        height: int = 1024,
        steps: int = 20,
        cfg: float = 4.5,
        seed: int | None = None,
        output_path: str = "output.png",
        raw: bool = False,
    ) -> None:
        """Generate an image from the loaded SDXL checkpoint.

        Args:
            prompt: Positive prompt (ComfyUI weighting syntax unless ``raw``).
            negative_prompt: Negative prompt.
            width, height: Output resolution in pixels.
            steps: Number of inference steps.
            cfg: CFG scale.
            seed: RNG seed (random if None).
            output_path: Where to save the PNG. Metadata is written to
                ``{output_path}.json``. Parent directories are created.
            raw: If True, bypass Compel prompt weighting and use plain
                diffusers encoding.

        Raises:
            RuntimeError: If SDXL was not loaded at session construction.
        """
        if self._pipe_sdxl is None:
            raise RuntimeError(
                "SDXL model not loaded. Provide sdxl_checkpoint when creating VnAssetsSession."
            )

        if seed is None:
            seed = random.randint(0, 2**32 - 1)

        output = Path(output_path)
        output.parent.mkdir(parents=True, exist_ok=True)

        generator = torch.Generator(device=self.device).manual_seed(seed)

        t0 = time.perf_counter()
        if raw:
            image = self._pipe_sdxl(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                num_inference_steps=steps,
                guidance_scale=cfg,
                generator=generator,
            ).images[0]
        else:
            prompt_embeds, pooled_embeds, neg_embeds, neg_pooled = encode_prompts(
                self._compel, prompt, negative_prompt
            )
            image = self._pipe_sdxl(
                prompt_embeds=prompt_embeds,
                pooled_prompt_embeds=pooled_embeds,
                negative_prompt_embeds=neg_embeds,
                negative_pooled_prompt_embeds=neg_pooled,
                width=width,
                height=height,
                num_inference_steps=steps,
                guidance_scale=cfg,
                generator=generator,
            ).images[0]
        t_infer = round(time.perf_counter() - t0, 2)

        image.save(output)
        print(f"Saved {output}")

        meta_path = output.with_suffix(".json")
        meta = {
            "checkpoint": str(Path(self._sdxl_checkpoint).resolve()),
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "width": width,
            "height": height,
            "steps": steps,
            "cfg": cfg,
            "seed": seed,
            "load_time_s": self._load_time_s,
            "inference_time_s": t_infer,
        }
        meta_path.write_text(json.dumps(meta, indent=2))
        print(f"Saved {meta_path}")

    # ── Edit (Qwen Image Edit) ──────────────────────────────────────────

    def edit(
        self,
        input_path: str,
        prompt: str,
        steps: int = 20,
        cfg: float = 4.0,
        seed: int | None = None,
        output_path: str = "output.png",
    ) -> None:
        """Edit an image using the loaded Qwen Image Edit model.

        Args:
            input_path: Path to the input image to edit.
            prompt: Edit instruction (e.g. "make her smile").
            steps: Number of inference steps (4 with Lightning LoRA).
            cfg: CFG scale (1.0 with Lightning LoRA).
            seed: RNG seed (random if None).
            output_path: Where to save the PNG. Metadata is written to
                ``{output_path}.json``. Parent directories are created.

        Raises:
            RuntimeError: If Qwen was not loaded at session construction.
        """
        if self._pipe_qwen is None:
            raise RuntimeError(
                "Qwen model not loaded. Provide edit_model when creating VnAssetsSession."
            )

        if seed is None:
            seed = random.randint(0, 2**32 - 1)

        output = Path(output_path)
        output.parent.mkdir(parents=True, exist_ok=True)

        input_image = Image.open(input_path).convert("RGB")
        generator = torch.Generator(device=self.device).manual_seed(seed)

        t0 = time.perf_counter()
        image = self._pipe_qwen(
            image=input_image,
            prompt=prompt,
            true_cfg_scale=cfg,
            num_inference_steps=steps,
            generator=generator,
        ).images[0]
        t_infer = round(time.perf_counter() - t0, 2)

        image.save(output)
        print(f"Saved {output}")

        meta_path = output.with_suffix(".json")
        meta = {
            "model": str(Path(self._edit_model).resolve()),
            "vae": VAE_ID,
            "text_encoder": TEXT_ENCODER_ID,
            "input_image": str(Path(input_path).resolve()),
            "prompt": prompt,
            "steps": steps,
            "cfg": cfg,
            "seed": seed,
            "lora_path": str(Path(self._edit_lora).resolve()) if self._edit_lora else None,
            "lora_load_s": self._lora_load_s,
            "lora_fused": self._lora_fused,
            "load_time_s": self._load_time_s,
            "inference_time_s": t_infer,
        }
        meta_path.write_text(json.dumps(meta, indent=2))
        print(f"Saved {meta_path}")

    # ── Lifecycle ───────────────────────────────────────────────────────

    def close(self) -> None:
        """Release all models and free GPU memory."""
        if self._pipe_sdxl:
            del self._pipe_sdxl
            self._pipe_sdxl = None
        if self._pipe_qwen:
            del self._pipe_qwen
            self._pipe_qwen = None
        self._compel = None
        if self.device == "cuda":
            torch.cuda.empty_cache()

    def __enter__(self) -> "VnAssetsSession":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
        self.close()
        return False