| import torch |
| import os |
| from diffusers import QwenImage21Pipeline |
|
|
| |
| |
|
|
| 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 |
| |
| return self.pipeline, self.pipeline |
|
|
|
|
