| """ | |
| ONNX Runtime ODE Solver for SmolOmni Flow-Matching Image Generation | |
| Usage: | |
| import onnxruntime as ort | |
| sess_ctx = ort.InferenceSession("smolomni_256M_gen_context.onnx") | |
| sess_flow = ort.InferenceSession("smolomni_256M_flow_head_step.onnx") | |
| def generate_image(prompt_tokens, num_steps=50): | |
| ctx = sess_ctx.run(None, {"input_ids": prompt_tokens})[0] | |
| latents = np.random.randn(1, 4, 32, 32).astype(np.float32) | |
| dt = 1.0 / num_steps | |
| for i in range(num_steps): | |
| t = np.array([i * dt * 1000], dtype=np.float32) | |
| velocity = sess_flow.run(None, { | |
| "noisy_latents": latents, | |
| "timestep": t, | |
| "context": ctx, | |
| })[0] | |
| latents = latents + velocity * dt | |
| return latents # Pass to VAE decoder for final image | |
| """ | |