Return the Tensor when input_info.return_result_tensor is true #776
Conversation
Summary of ChangesHello @sitatec, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the flexibility of image generation pipelines by ensuring that the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly implements returning a tensor instead of saving an image to disk when input_info.return_result_tensor is true. The changes are consistently applied across longcat_image, qwen_image, and z_image runners and their corresponding VAE decoders. My review includes a few suggestions to improve code quality by removing redundant f-strings and to fix a potential bug in how file paths are parsed in the qwen_image runner.
| logger.info(f"Image saved: {input_info.save_result_path}") | ||
| if not input_info.return_result_tensor: | ||
| if isinstance(images[0], list) and len(images[0]) > 1: | ||
| image_prefix = f"{input_info.save_result_path}".split(".")[0] |
There was a problem hiding this comment.
Using .split('.') to remove a file extension is not robust as it will fail for file paths that contain dots in their directory names (e.g., /path.with.dots/image.png). It's safer to use rsplit('.', 1) to split only on the last dot. Also, the f-string is redundant.
A more robust way would be to use os.path.splitext, but that would require an import. Using rsplit is a good improvement without adding imports.
| image_prefix = f"{input_info.save_result_path}".split(".")[0] | |
| image_prefix = input_info.save_result_path.rsplit(".", 1)[0] |
| logger.info(f"Image saved: {input_info.save_result_path}") | ||
| if not input_info.return_result_tensor: | ||
| image = images[0] | ||
| image.save(f"{input_info.save_result_path}") |
There was a problem hiding this comment.
The f-string formatting is redundant here as input_info.save_result_path is already a string. You can pass the variable directly to image.save for cleaner code.
| image.save(f"{input_info.save_result_path}") | |
| image.save(input_info.save_result_path) |
| logger.info(f"Image saved: {input_info.save_result_path}") | ||
| if not input_info.return_result_tensor: | ||
| image = images[0] | ||
| image.save(f"{input_info.save_result_path}") |
There was a problem hiding this comment.
The f-string formatting is redundant here since input_info.save_result_path is already a string. You can simplify the code by passing the variable directly to image.save.
| image.save(f"{input_info.save_result_path}") | |
| image.save(input_info.save_result_path) |
1 participant
input_info.return_result_tensorwas ignored for all image generation. Outputting the tensor can be useful for post-processing (such as NSFW checking), without reloading the file from disk.I noticed that for video models, they do not return the tensor directly; they return a map of {"video": tensor} here. I believe this is for compatibility with ComfyUI. If that's the case, we should only return the tensor and move ComfyUI-specific patterns to the ComfyUI wrapper codebase. What do you think?