Skip to content

fix bug for Kandinsky5 pipeline load with device_map=balanced#14050

Open
kaixuanliu wants to merge 1 commit into
huggingface:mainfrom
kaixuanliu:kandinsky5_fix
Open

fix bug for Kandinsky5 pipeline load with device_map=balanced#14050
kaixuanliu wants to merge 1 commit into
huggingface:mainfrom
kaixuanliu:kandinsky5_fix

Conversation

@kaixuanliu

Copy link
Copy Markdown
Contributor

Fix Kandinsky5 I2I/I2V failures with multi-GPU device_map.

Under device_map="balanced", VAE-encoded image latents may be produced on a different GPU than the generated noise latents. Aligning image_latents to latents.device and latents.dtype before concatenation/assignment prevents cross-device runtime errors during image-conditioned generation.
Here is an example:

import torch
from diffusers import Kandinsky5I2IPipeline # use duffusers fork: pip install git+https://github.com/leffff/diffusers
from diffusers.utils import load_image 
# Load the pipeline
model_id = "kandinskylab/Kandinsky-5.0-I2I-Lite-sft-Diffusers"
pipe = Kandinsky5I2IPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="balanced") # <--- Load the model with automatic device mapping
                                            # <--- Enable CPU offloading for single GPU inference

# Edit the input image
image = load_image(
    "https://huggingface.co/kandinsky-community/kandinsky-3/resolve/main/assets/title.jpg?download=true"
)

prompt = "Change the background from a winter night scene to a bright summer day. Place the character on a sandy beach with clear blue sky, soft sunlight, and gentle waves in the distance. Replace the winter clothing with a light short-sleeved T-shirt (in soft pastel colors) and casual shorts. Ensure the character’s fur reflects warm daylight instead of cold winter tones. Add small beach details such as seashells, footprints in the sand, and a few scattered beach toys nearby. Keep the oranges in the scene, but place them naturally on the sand."
negative_prompt = ""

output = pipe(
    image=image,
    prompt=prompt,
    negative_prompt=negative_prompt,
    guidance_scale=3.5,
).image[0]

@yiyixuxu @dg845 and @asomoza pls help review, thx!

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>
@github-actions github-actions Bot added pipelines size/S PR with diff < 50 LOC labels Jun 23, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Hi @kaixuanliu, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pipelines size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant