|
| 1 | +import math |
| 2 | + |
| 3 | +import torch |
| 4 | +from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageEditPlusPipeline |
| 5 | +from diffusers.utils import load_image |
| 6 | + |
| 7 | +from nunchaku import NunchakuQwenImageTransformer2DModel |
| 8 | +from nunchaku.utils import get_gpu_memory, get_precision |
| 9 | + |
| 10 | +# From https://github.com/ModelTC/Qwen-Image-Lightning/blob/342260e8f5468d2f24d084ce04f55e101007118b/generate_with_diffusers.py#L82C9-L97C10 |
| 11 | +scheduler_config = { |
| 12 | + "base_image_seq_len": 256, |
| 13 | + "base_shift": math.log(3), # We use shift=3 in distillation |
| 14 | + "invert_sigmas": False, |
| 15 | + "max_image_seq_len": 8192, |
| 16 | + "max_shift": math.log(3), # We use shift=3 in distillation |
| 17 | + "num_train_timesteps": 1000, |
| 18 | + "shift": 1.0, |
| 19 | + "shift_terminal": None, # set shift_terminal to None |
| 20 | + "stochastic_sampling": False, |
| 21 | + "time_shift_type": "exponential", |
| 22 | + "use_beta_sigmas": False, |
| 23 | + "use_dynamic_shifting": True, |
| 24 | + "use_exponential_sigmas": False, |
| 25 | + "use_karras_sigmas": False, |
| 26 | +} |
| 27 | +scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config) |
| 28 | + |
| 29 | +num_inference_steps = 4 # you can also use the 8-step model to improve the quality |
| 30 | +rank = 32 # you can also use the rank=128 model to improve the quality |
| 31 | +model_path = f"nunchaku-tech/nunchaku-qwen-image-edit-2509-lightning/svdq-{get_precision()}_r{rank}-qwen-image-edit-2509-lightningv2.0-{num_inference_steps}steps.safetensors" |
| 32 | + |
| 33 | +# Load the model |
| 34 | +transformer = NunchakuQwenImageTransformer2DModel.from_pretrained(model_path) |
| 35 | + |
| 36 | +pipeline = QwenImageEditPlusPipeline.from_pretrained( |
| 37 | + "Qwen/Qwen-Image-Edit-2509", transformer=transformer, torch_dtype=torch.bfloat16 |
| 38 | +) |
| 39 | + |
| 40 | +if get_gpu_memory() > 18: |
| 41 | + pipeline.enable_model_cpu_offload() |
| 42 | +else: |
| 43 | + # use per-layer offloading for low VRAM. This only requires 3-4GB of VRAM. |
| 44 | + transformer.set_offload( |
| 45 | + True, use_pin_memory=False, num_blocks_on_gpu=1 |
| 46 | + ) # increase num_blocks_on_gpu if you have more VRAM |
| 47 | + pipeline._exclude_from_cpu_offload.append("transformer") |
| 48 | + pipeline.enable_sequential_cpu_offload() |
| 49 | + |
| 50 | +image1 = load_image("https://huggingface.co/datasets/nunchaku-tech/test-data/resolve/main/inputs/man.png") |
| 51 | +image1 = image1.convert("RGB") |
| 52 | +image2 = load_image("https://huggingface.co/datasets/nunchaku-tech/test-data/resolve/main/inputs/puppy.png") |
| 53 | +image2 = image2.convert("RGB") |
| 54 | +image3 = load_image("https://huggingface.co/datasets/nunchaku-tech/test-data/resolve/main/inputs/sofa.png") |
| 55 | +image3 = image3.convert("RGB") |
| 56 | + |
| 57 | +prompt = "Let the man in image 1 lie on the sofa in image 3, and let the puppy in image 2 lie on the floor to sleep." |
| 58 | +inputs = { |
| 59 | + "image": [image1, image2, image3], |
| 60 | + "prompt": prompt, |
| 61 | + "true_cfg_scale": 1.0, |
| 62 | + "num_inference_steps": num_inference_steps, |
| 63 | +} |
| 64 | + |
| 65 | +output = pipeline(**inputs) |
| 66 | +output_image = output.images[0] |
| 67 | +output_image.save(f"qwen-image-edit-2509-lightning-r{rank}-{num_inference_steps}steps.png") |
0 commit comments