1ab1fee15d
Bundles the known-working ROCm torch, triton-rocm, and torchvision wheels
locally so install no longer depends on the PyTorch nightly index being
available. Also fixes model name mismatches between docs and code, and
adds missing dependencies.
What:
• wheels/ — bundled torch 2.11.0+rocm7.2, triton-rocm 3.6.0, and
torchvision 0.26.0+rocm7.2 (wheels/ is in .gitignore; they are
downloaded once and reused)
• README install section — replaced 'pip install torch --index-url ...'
with 'pip install wheels/*.whl' for all three ROCm wheels; added
download instructions for the initial fetch
• README Current State — marked self-contained torch wheel as done,
removed from Future Improvements
• Fixed model name defaults across README examples and session.py
docstring — they used shortened filenames (novaAnimeXL.safetensors,
qwen_image_edit.safetensors) that don't match the actual symlinks
created by the Models section
• pyproject.toml — added torchvision and realesrgan (previously
undeclared; upscale.py imported torchvision.transforms at module
level without the dependency being declared)
Why:
The ROCm nightly index is volatile — versions rotate frequently and the
index itself may not be reachable. Bundling the three ROCm-only wheels
(torch, triton-rocm, torchvision) makes the install reproducible. All
other dependencies resolve from standard PyPI.
Verified by creating a clean Python 3.12 venv, installing all three
wheels from the local files, running 'pip install -e .', and confirming
'vnasset --help' and GPU tensor allocation work without any remote
ROCm index access.
471 lines
17 KiB
Python
471 lines
17 KiB
Python
"""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 .background import remove_background as _remove_bg
|
|
from .background import remove_backgrounds as _remove_bgs
|
|
from .prompt import build_compel, encode_prompts
|
|
from .upscale import _build_upsampler
|
|
from .upscale import upscale as _upscale_fn
|
|
from .upscale import upscales as _upscales_fn
|
|
|
|
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_ilV190.safetensors",
|
|
edit_model="models/qwen_image_edit_2509_fp8_e4m3fn.safetensors",
|
|
edit_lora="models/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.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
|
|
self._rembg_session = None
|
|
self._upsampler = 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}")
|
|
|
|
# ── Background Removal ─────────────────────────────────────────────
|
|
|
|
def remove_background(
|
|
self,
|
|
input_path: str,
|
|
output_path: str,
|
|
model: str = "isnet-anime",
|
|
) -> None:
|
|
"""Remove background from an image. Output is RGBA PNG.
|
|
|
|
Wraps rembg with ``isnet-anime`` (trained on anime images) by
|
|
default. The rembg session is created once and reused across
|
|
calls, so batch use is efficient.
|
|
|
|
Args:
|
|
input_path: Path to the input RGB image.
|
|
output_path: Where to save the RGBA PNG.
|
|
model: rembg model name. Default ``isnet-anime``. Also
|
|
available: ``u2net``, ``u2netp``, ``u2net_human_seg``,
|
|
``isnet-general-use``, ``sam``.
|
|
|
|
Raises:
|
|
FileNotFoundError: If input_path doesn't exist.
|
|
"""
|
|
if self._rembg_session is None:
|
|
from rembg import new_session
|
|
self._rembg_session = new_session(model)
|
|
_remove_bg(input_path, output_path, model=model, session=self._rembg_session)
|
|
|
|
def remove_backgrounds(
|
|
self,
|
|
input_paths: list[str],
|
|
output_dir: str,
|
|
model: str = "isnet-anime",
|
|
) -> None:
|
|
"""Remove backgrounds from multiple images in one batch.
|
|
|
|
Output files are named ``{stem}_nobg.png``.
|
|
|
|
Args:
|
|
input_paths: List of input image paths.
|
|
output_dir: Directory for output RGBA PNGs.
|
|
model: rembg model name. Default ``isnet-anime``.
|
|
"""
|
|
if self._rembg_session is None:
|
|
from rembg import new_session
|
|
self._rembg_session = new_session(model)
|
|
_remove_bgs(input_paths, output_dir, model=model)
|
|
|
|
# ── Upscaling ───────────────────────────────────────────────────
|
|
|
|
def upscale(
|
|
self,
|
|
input_path: str,
|
|
output_path: str,
|
|
scale: int = 2,
|
|
) -> None:
|
|
"""Upscale an image using Real-ESRGAN (anime model).
|
|
|
|
The upsampler is lazy-loaded on first call and reused across
|
|
calls. Uses ``RealESRGAN_x4plus_anime_6B`` (17 MB,
|
|
auto-downloaded).
|
|
|
|
Args:
|
|
input_path: Path to the input RGB image.
|
|
output_path: Where to save the upscaled PNG.
|
|
scale: Output scale factor (2, 3, or 4).
|
|
|
|
Raises:
|
|
FileNotFoundError: If input_path doesn't exist.
|
|
ValueError: If scale is not 2, 3, or 4.
|
|
"""
|
|
if self._upsampler is None:
|
|
self._upsampler = _build_upsampler(self.device)
|
|
# Also use this upsampler for subsequent calls
|
|
_upscale_fn(input_path, output_path, scale=scale, upsampler=self._upsampler)
|
|
|
|
def upscales(
|
|
self,
|
|
input_paths: list[str],
|
|
output_dir: str,
|
|
scale: int = 2,
|
|
) -> None:
|
|
"""Upscale multiple images, reusing one model session.
|
|
|
|
Output files are named ``{stem}_x{scale}.png``.
|
|
|
|
Args:
|
|
input_paths: List of input image paths.
|
|
output_dir: Directory for output PNGs.
|
|
scale: Output scale factor (2, 3, or 4).
|
|
"""
|
|
if self._upsampler is None:
|
|
self._upsampler = _build_upsampler(self.device)
|
|
_upscales_fn(input_paths, output_dir, scale=scale, upsampler=self._upsampler)
|
|
|
|
# ── 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._rembg_session:
|
|
del self._rembg_session
|
|
self._rembg_session = None
|
|
if self._upsampler:
|
|
del self._upsampler
|
|
self._upsampler = 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
|