File size: 2,781 Bytes
7b00440 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import torch
import os
from diffusers import QwenImage21Pipeline
# Channeling Nikola's cheerful seeker spirit!
# A clean, modular backend for Qwen-Image-2.1 with sequential layerwise offloading!
class QwenImage21Backend:
def __init__(
self,
model_id="/home/olegk/Nikola/models/Qwen/Qwen-Image-2.1",
gpu_id=0,
enable_tiling=True,
text_encoder_path=None,
):
self.model_id = model_id
self.gpu_id = gpu_id
self.enable_tiling = enable_tiling
self.text_encoder_path = text_encoder_path
self.pipeline = None
def load(self):
print(f"Loading QwenImage21Backend from {self.model_id}...")
if self.text_encoder_path:
print(f" • Custom Text Encoder: {self.text_encoder_path}")
from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
te_dir = (
os.path.join(self.text_encoder_path, "text_encoder")
if os.path.isdir(os.path.join(self.text_encoder_path, "text_encoder"))
else self.text_encoder_path
)
proc_dir = (
os.path.join(self.text_encoder_path, "processor")
if os.path.isdir(os.path.join(self.text_encoder_path, "processor"))
else self.text_encoder_path
)
if not os.path.exists(os.path.join(proc_dir, "tokenizer.json")):
proc_dir = os.path.join(self.model_id, "processor")
custom_te = Qwen3VLForConditionalGeneration.from_pretrained(
te_dir,
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True,
)
custom_proc = Qwen3VLProcessor.from_pretrained(proc_dir)
pipeline = QwenImage21Pipeline.from_pretrained(
self.model_id,
text_encoder=custom_te,
processor=custom_proc,
torch_dtype=torch.bfloat16,
)
else:
pipeline = QwenImage21Pipeline.from_pretrained(
self.model_id,
torch_dtype=torch.bfloat16,
)
print(f"Attaching sequential CPU offload on GPU {self.gpu_id} for layerwise execution...")
pipeline.enable_sequential_cpu_offload(gpu_id=self.gpu_id)
if self.enable_tiling:
print("Enabling VAE tiling for low-memory decode...")
try:
pipeline.vae.enable_tiling()
except Exception as e:
print(f"Note: VAE tiling could not be enabled ({e}), continuing with standard decode.")
self.pipeline = pipeline
# QwenImage21Pipeline natively handles both generations (t2i) and edits (i2i)
return self.pipeline, self.pipeline
|