Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Policy Optimized Text-to-Image Pipeline Design

Uri Gadot, Rinon Gal, Yftah Ziser, Gal Chechik, Shie Mannor

TL;DR

This work trains an ensemble of reward models capable of predicting image quality scores directly from prompt-workflow combinations, eliminating the need for costly image generation during training and incorporates a classifier-free guidance based enhancement technique that extrapolates along the path between the initial and GRPO-tuned models, further improving output quality.

Abstract

Text-to-image generation has evolved beyond single monolithic models to complex multi-component pipelines. These combine fine-tuned generators, adapters, upscaling blocks and even editing steps, leading to significant improvements in image quality. However, their effective design requires substantial expertise. Recent approaches have shown promise in automating this process through large language models (LLMs), but they suffer from two critical limitations: extensive computational requirements from generating images with hundreds of predefined pipelines, and poor generalization beyond memorized training examples. We introduce a novel reinforcement learning-based framework that addresses these inefficiencies. Our approach first trains an ensemble of reward models capable of predicting image quality scores directly from prompt-workflow combinations, eliminating the need for costly image generation during training. We then implement a two-phase training strategy: initial workflow vocabulary training followed by GRPO-based optimization that guides the model toward higher-performing regions of the workflow space. Additionally, we incorporate a classifier-free guidance based enhancement technique that extrapolates along the path between the initial and GRPO-tuned models, further improving output quality. We validate our approach through a set of comparisons, showing that it can successfully create new flows with greater diversity and lead to superior image quality compared to existing baselines.

Policy Optimized Text-to-Image Pipeline Design

TL;DR

This work trains an ensemble of reward models capable of predicting image quality scores directly from prompt-workflow combinations, eliminating the need for costly image generation during training and incorporates a classifier-free guidance based enhancement technique that extrapolates along the path between the initial and GRPO-tuned models, further improving output quality.

Abstract

Text-to-image generation has evolved beyond single monolithic models to complex multi-component pipelines. These combine fine-tuned generators, adapters, upscaling blocks and even editing steps, leading to significant improvements in image quality. However, their effective design requires substantial expertise. Recent approaches have shown promise in automating this process through large language models (LLMs), but they suffer from two critical limitations: extensive computational requirements from generating images with hundreds of predefined pipelines, and poor generalization beyond memorized training examples. We introduce a novel reinforcement learning-based framework that addresses these inefficiencies. Our approach first trains an ensemble of reward models capable of predicting image quality scores directly from prompt-workflow combinations, eliminating the need for costly image generation during training. We then implement a two-phase training strategy: initial workflow vocabulary training followed by GRPO-based optimization that guides the model toward higher-performing regions of the workflow space. Additionally, we incorporate a classifier-free guidance based enhancement technique that extrapolates along the path between the initial and GRPO-tuned models, further improving output quality. We validate our approach through a set of comparisons, showing that it can successfully create new flows with greater diversity and lead to superior image quality compared to existing baselines.

Paper Structure

This paper contains 44 sections, 6 equations, 11 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Improving Text-to-Image Generation Quality
  4. Methodology
  5. Training Data
  6. Stage 1: Supervised Fine-Tuning on Flow Dataset
  7. Efficient Flow Representation Scheme
  8. Stage 2: Reward-based policy-optimization
  9. Surrogate Reward Model Training
  10. Component-Aware Hybrid Reward Formulation
  11. Uncertainty-Aware Reinforcement Learning
  12. Dual model guidance
  13. Experiments
  14. Comparisons
  15. Prompt adherence:
  16. ...and 29 more sections

Figures (11)

  • Figure 1: Pipeline overview. Step 1: Finetune LLM for general flow generation (SFT, 500K prompt-flow pairs). Step 2.1: Train reward model (100K prompt-flow-score triplets). Step 2.2: Optimize for quality using GRPO. = learning, ❄ = frozen.
  • Figure 2: An example of a single ComfyUI node tokenized. (a) displays the original JSON input as tokenized by the standard Llama tokenizer. (b) shows our custom encoding, with introducing additional tokens to explicitly represent relevant components within workflow. Colored segment corresponds to a different token. (c) histogram of flows length (in token) of all training-set
  • Figure 3: Example of generations with FlowRL
  • Figure 4: Qualitative results on CivitAI and GenEval prompts.
  • Figure 5: Human study win rate of FlowRL vs other relevant baselines
  • ...and 6 more figures