File size: 2,479 Bytes
1e103b7 | 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 | import torch
from nunchaku.utils import get_gpu_memory, get_precision
from nunchaku.models.transformers.transformer_qwenimage import NunchakuQwenImageTransformer2DModel
class QwenImageBackend:
def __init__(self, model_id, optimized_model_path=None):
self.model_id = model_id
self.optimized_model_path = optimized_model_path
self.pipeline = None
self.rank = 32 # default rank as per example
def load(self):
print(f"Loading QwenImageBackend from {self.model_id}...")
# Scheduler config (same as QwenBackend)
import math
from diffusers import FlowMatchEulerDiscreteScheduler
scheduler_config = {
"base_image_seq_len": 256,
"base_shift": math.log(3),
"invert_sigmas": False,
"max_image_seq_len": 8192,
"max_shift": math.log(3),
"num_train_timesteps": 1000,
"shift": 1.0,
"shift_terminal": None,
"stochastic_sampling": False,
"time_shift_type": "exponential",
"use_beta_sigmas": False,
"use_dynamic_shifting": True,
"use_exponential_sigmas": False,
"use_karras_sigmas": False,
}
scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
# Load transformer (optimized model)
print(f"Loading NunchakuQwenImageTransformer2DModel from {self.optimized_model_path}...")
transformer = NunchakuQwenImageTransformer2DModel.from_pretrained(self.optimized_model_path)
# Load T2I pipeline
from diffusers import QwenImagePipeline
pipeline = QwenImagePipeline.from_pretrained(
self.model_id,
transformer=transformer,
scheduler=scheduler,
torch_dtype=torch.bfloat16,
)
# Offloading logic (same as QwenBackend)
if get_gpu_memory() > 18:
print("GPU memory > 18GB, using cpu offload")
pipeline.enable_model_cpu_offload()
else:
print("GPU memory <= 18GB, using per-layer offloading for low VRAM")
transformer.set_offload(True, use_pin_memory=False, num_blocks_on_gpu=1)
pipeline._exclude_from_cpu_offload.append("transformer")
pipeline.enable_sequential_cpu_offload()
self.pipeline = pipeline
# For edit endpoint we reuse the same pipeline (ignores image)
return self.pipeline, self.pipeline
|