Add diffusers backend support for image generation

[ericcurtin]

This commit introduces a new diffusers backend that enables image
generation capabilities in the model runner. The backend integrates with
Hugging Face diffusers library to provide OpenAI-compatible image
generation API endpoints.

The implementation includes a new Python server component that runs as a
subprocess and handles image generation requests through a Unix socket
interface. The backend supports various diffusers models and provides
configuration options for device selection, precision, and memory
optimization features.

Additionally, the Dockerfile now includes a diffusers variant with the
necessary Python dependencies pre-installed. The model distribution
logic has been updated to properly handle diffusers model formats with
their directory structure preserved in the model bundles.

Signed-off-by: Eric Curtin [email protected]

[gemini-code-assist]

Summary of Changes

Hello @ericcurtin, I'm Gemini Code Assist¹! 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 significantly expands the model runner's capabilities by adding support for image generation. It integrates the Hugging Face diffusers library through a new backend, which includes a Python server to process image generation requests via an OpenAI-compatible API. The changes also encompass updates to the Docker build process and model handling to properly package and deploy diffusers models, making the system more versatile for different AI tasks.

Highlights

• New Diffusers Backend: Introduced a new backend that leverages the Hugging Face diffusers library to enable image generation capabilities within the model runner.
• OpenAI-Compatible API: The new backend provides OpenAI-compatible image generation API endpoints, allowing for standardized interaction.
• Python Server Component: A dedicated Python server runs as a subprocess, handling image generation requests and communicating via a Unix socket interface.
• Dockerfile Updates: The Dockerfile now includes a specific 'diffusers' variant, pre-installing necessary Python dependencies like diffusers, torch, transformers, and FastAPI.
• Enhanced Model Distribution: The model distribution logic has been updated to correctly identify and bundle diffusers models, preserving their directory structure for proper loading.
• Configurable Options: The diffusers backend supports various configuration options for device selection (CPU, CUDA, MPS), precision (fp16, bf16, fp32), and memory optimization features like attention slicing.

🧠 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 Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[sourcery-ai]

Hey - I've found 2 security issues, and left some high level feedback:

Security issues:

• Detected non-static command inside Command. Audit the input to 'exec.Command'. If unverified user data can reach this call site, this is a code injection vulnerability. A malicious actor can inject a malicious script to execute arbitrary code. (link)
• Detected non-static command inside Command. Audit the input to 'exec.Command'. If unverified user data can reach this call site, this is a code injection vulnerability. A malicious actor can inject a malicious script to execute arbitrary code. (link)

General comments:

• In buildDiffusersModel, builder.FromSafetensors is called even when only .bin files are found and only the first weight file is used, which is both misleading and likely incorrect for sharded or non-safetensors diffusers models; consider using a more generic builder entry point and handling multi-part weights explicitly.
• The Dockerfile copies the diffusers server into a hard-coded python3.12 site-packages path, which will break if the venv is created with a different Python minor version; consider discovering the venv’s site-packages directory dynamically or using python -m site to construct this path.
• The isContainerInstall helper in the diffusers backend is currently unused; either wire it into behavior (e.g., for conditional logging or paths) or remove it to avoid dead code.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- In `buildDiffusersModel`, `builder.FromSafetensors` is called even when only `.bin` files are found and only the first weight file is used, which is both misleading and likely incorrect for sharded or non-safetensors diffusers models; consider using a more generic builder entry point and handling multi-part weights explicitly.
- The Dockerfile copies the diffusers server into a hard-coded `python3.12` site-packages path, which will break if the venv is created with a different Python minor version; consider discovering the venv’s site-packages directory dynamically or using `python -m site` to construct this path.
- The `isContainerInstall` helper in the diffusers backend is currently unused; either wire it into behavior (e.g., for conditional logging or paths) or remove it to avoid dead code.

## Individual Comments

### Comment 1
<location> `pkg/inference/backends/diffusers/diffusers.go:125` </location>
<code_context>
		cmd := exec.Command(pythonPath, "-c", "import diffusers; print(diffusers.__version__)")
</code_context>

<issue_to_address>
**security (go.lang.security.audit.dangerous-exec-command):** Detected non-static command inside Command. Audit the input to 'exec.Command'. If unverified user data can reach this call site, this is a code injection vulnerability. A malicious actor can inject a malicious script to execute arbitrary code.

*Source: opengrep*
</issue_to_address>

### Comment 2
<location> `pkg/inference/backends/diffusers/diffusers.go:140` </location>
<code_context>
	cmd := exec.Command(d.pythonPath, "-c", "import diffusers; print(diffusers.__version__)")
</code_context>

<issue_to_address>
**security (go.lang.security.audit.dangerous-exec-command):** Detected non-static command inside Command. Audit the input to 'exec.Command'. If unverified user data can reach this call site, this is a code injection vulnerability. A malicious actor can inject a malicious script to execute arbitrary code.

*Source: opengrep*
</issue_to_address>

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[gemini-code-assist]

Code Review

This pull request introduces a new diffusers backend for image generation. The changes are comprehensive, including a new Dockerfile variant, modifications to the model distribution logic to handle diffusers models, and a new Python server for handling image generation requests. The implementation is well-structured. I have a few suggestions to improve robustness and code clarity. Specifically, I've pointed out a brittle path in the Dockerfile, suggested refactoring some duplicated code, proposed a simplification for model path resolution, and recommended aligning an HTTP status code with standard practices.