Skip to content

[fix]:flux2 i2i auto target shape #985

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
llmc-reviewer merged 1 commit into from
Apr 2, 2026
+7 −16
Merged

[fix]:flux2 i2i auto target shape #985

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 7 additions & 16 deletions lightx2v/models/runners/flux2_klein/flux2_klein_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ def _run_input_encoder_local_i2i(self):
input_image = [input_image]

condition_images = []
for img in input_image:
for index, img in enumerate(input_image):
image_processor.check_image_input(img)
image_width, image_height = img.size
if image_width * image_height > 1024 * 1024:
Expand All @@ -122,9 +122,8 @@ def _run_input_encoder_local_i2i(self):
image_height = (image_height // multiple_of) * multiple_of
img = image_processor.preprocess(img, height=image_height, width=image_width, resize_mode="crop")
condition_images.append(img.to(AI_DEVICE))
if not hasattr(self.input_info, "auto_width"):
self.input_info.auto_width = image_width
self.input_info.auto_height = image_height
if index == 0:
self.input_info.target_shape = (image_height, image_width)

torch.cuda.empty_cache()
gc.collect()
Expand Down Expand Up @@ -254,26 +253,18 @@ def get_custom_shape(self):
return (width, height)

def set_target_shape(self):
multiple_of = self.config.get("vae_scale_factor", 8) * 2

task = self.config.get("task", "t2i")
if task == "i2i" and hasattr(self.input_info, "auto_width"):
width = self.input_info.auto_width
height = self.input_info.auto_height
else:
if task == "i2i": # for i2i task, the target shape is already set in _run_input_encoder_local_i2i
height, width = self.input_info.target_shape
else: # for t2i task, calculate the target shape based on the resolution
custom_shape = self.get_custom_shape()
if custom_shape is not None:
width, height = custom_shape
else:
calculated_width, calculated_height, _ = calculate_dimensions(self.resolution * self.resolution, 16 / 9)
width = calculated_width // multiple_of * multiple_of
height = calculated_height // multiple_of * multiple_of

self.input_info.auto_width = width
self.input_info.auto_height = height

self.input_info.target_shape = (height, width)
logger.info(f"Flux2Klein Image Runner set target shape: {width}x{height}")
self.input_info.target_shape = (height, width)

multiple_of = self.config.get("vae_scale_factor", 8) * 2
Comment on lines 256 to 269
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Apr 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The set_target_shape method contains a critical bug and a logic error:

  1. NameError: multiple_of is used on lines 265 and 266 before it is defined on line 269. This will cause a crash during execution for the t2i task.
  2. Logic Error: self.input_info.target_shape is only updated inside the else block (line 267). If a custom_shape is provided (line 261), the target_shape attribute in input_info is never updated, which will lead to incorrect behavior in subsequent pipeline stages that rely on this value.
  3. Missing Logging: The useful log message indicating the final target shape was removed.

I recommend moving the multiple_of definition to the top and ensuring target_shape is always assigned for the t2i task.

Suggested change
task = self.config.get("task", "t2i")
if task == "i2i" and hasattr(self.input_info, "auto_width"):
width = self.input_info.auto_width
height = self.input_info.auto_height
else:
if task == "i2i": # for i2i task, the target shape is already set in _run_input_encoder_local_i2i
height, width = self.input_info.target_shape
else: # for t2i task, calculate the target shape based on the resolution
custom_shape = self.get_custom_shape()
if custom_shape is not None:
width, height = custom_shape
else:
calculated_width, calculated_height, _ = calculate_dimensions(self.resolution * self.resolution, 16 / 9)
width = calculated_width // multiple_of * multiple_of
height = calculated_height // multiple_of * multiple_of
self.input_info.auto_width = width
self.input_info.auto_height = height
self.input_info.target_shape = (height, width)
logger.info(f"Flux2Klein Image Runner set target shape: {width}x{height}")
self.input_info.target_shape = (height, width)
multiple_of = self.config.get("vae_scale_factor", 8) * 2
multiple_of = self.config.get("vae_scale_factor", 8) * 2
task = self.config.get("task", "t2i")
if task == "i2i": # for i2i task, the target shape is already set in _run_input_encoder_local_i2i
height, width = self.input_info.target_shape
else: # for t2i task, calculate the target shape based on the resolution
custom_shape = self.get_custom_shape()
if custom_shape is not None:
width, height = custom_shape
else:
calculated_width, calculated_height, _ = calculate_dimensions(self.resolution * self.resolution, 16 / 9)
width = calculated_width // multiple_of * multiple_of
height = calculated_height // multiple_of * multiple_of
self.input_info.target_shape = (height, width)
logger.info(f"Flux2Klein Image Runner set target shape: {width}x{height}")


Expand Down
Loading