refactoring for Flux

app.py

@@ -59,8 +59,8 @@ def start_over(gallery_state, loaded_model_setup):
 
 def setup_model(prompt, model, seed, num_iterations, learning_rate, hps_w, imgrw_w, pcks_w, clip_w, progress=gr.Progress(track_tqdm=True)):
     if prompt is None:
-        raise gr.Error("You forgot to provide a prompt !")
-    
+        raise gr.Error("You forgot the prompt !")
+        
     """Clear CUDA memory before starting the training."""
     torch.cuda.empty_cache()  # Free up cached memory
     
@@ -81,8 +81,13 @@ def setup_model(prompt, model, seed, num_iterations, learning_rate, hps_w, imgrw
     args.pickscore_weighting = pcks_w
     args.clip_weighting = clip_w
 
-    args, trainer, device, dtype, shape, enable_grad, settings = setup(args)
-    loaded_setup = [args, trainer, device, dtype, shape, enable_grad, settings]
+    if model == "flux":
+        args.cpu_offloading = True
+        args.enable_multi_apply= True
+        args.multi_step_model = "flux"
+
+    args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings = setup(args)
+    loaded_setup = [args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings]
 
     return None, loaded_setup
 
@@ -95,11 +100,12 @@ def generate_image(setup_args, num_iterations):
     dtype = setup_args[3]
     shape = setup_args[4]
     enable_grad = setup_args[5]
+    multi_apply_fn = setup_args[6]
 
-    settings = setup_args[6]
+    settings = setup_args[7]
     print(f"SETTINGS: {settings}")
 
-    save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
+    save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt[:150]}"
     clean_dir(save_dir)
     
     try:
@@ -119,7 +125,7 @@ def generate_image(setup_args, num_iterations):
             thread_status["running"] = True  # Mark thread as running
             try:
                 execute_task(
-                    args, trainer, device, dtype, shape, enable_grad, settings, progress_callback
+                    args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings, progress_callback
                 )
             except torch.cuda.OutOfMemoryError as e:
                 print(f"CUDA Out of Memory Error: {e}")
@@ -154,7 +160,7 @@ def generate_image(setup_args, num_iterations):
 
             if error_status["error_occurred"]:
                 torch.cuda.empty_cache()  # Free up cached memory
-                yield (None, "CUDA out of memory.", None)
+                yield (None, "CUDA out of memory. Please reduce your batch size or image resolution.", None)
             else:
                 main_thread.join()  # Ensure thread completion
                 final_image_path = os.path.join(save_dir, "best_image.png")
@@ -213,7 +219,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
             with gr.Column():
                 prompt = gr.Textbox(label="Prompt")
                 with gr.Row():
-                    chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd"], label="Model", value="sd-turbo")
+                    chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd", "flux"], label="Model", value="sd-turbo")
                     seed = gr.Number(label="seed", value=0)
                     model_status = gr.Textbox(label="model status", visible=False)
                 

arguments.py

@@ -38,36 +38,57 @@ def parse_args():
     parser.add_argument("--seed", type=int, help="Seed to use", default=0)
 
     # reward losses
-    parser.add_argument("--disable_hps", default=True, action="store_false",dest="enable_hps")
+    parser.add_argument(
+        "--disable_hps", default=True, action="store_false", dest="enable_hps"
+    )
     parser.add_argument(
         "--hps_weighting", type=float, help="Weighting for HPS", default=5.0
     )
-    parser.add_argument("--disable_imagereward", default=True, action="store_false",dest='enable_imagereward')
+    parser.add_argument(
+        "--disable_imagereward",
+        default=True,
+        action="store_false",
+        dest="enable_imagereward",
+    )
     parser.add_argument(
         "--imagereward_weighting",
         type=float,
         help="Weighting for ImageReward",
         default=1.0,
     )
-    parser.add_argument("--disable_clip", default=True, action="store_false",dest='enable_clip')
+    parser.add_argument(
+        "--disable_clip", default=True, action="store_false", dest="enable_clip"
+    )
     parser.add_argument(
         "--clip_weighting", type=float, help="Weighting for CLIP", default=0.01
     )
-    parser.add_argument("--disable_pickscore", default=True, action="store_false",dest='enable_pickscore')
+    parser.add_argument(
+        "--disable_pickscore",
+        default=True,
+        action="store_false",
+        dest="enable_pickscore",
+    )
     parser.add_argument(
         "--pickscore_weighting",
         type=float,
         help="Weighting for PickScore",
         default=0.05,
     )
-    parser.add_argument("--disable_aesthetic", default=False, action="store_false",dest='enable_aesthetic')
+    parser.add_argument(
+        "--disable_aesthetic",
+        default=False,
+        action="store_false",
+        dest="enable_aesthetic",
+    )
     parser.add_argument(
         "--aesthetic_weighting",
         type=float,
         help="Weighting for Aesthetic",
         default=0.0,
     )
-    parser.add_argument("--disable_reg", default=True, action="store_false",dest='enable_reg')
+    parser.add_argument(
+        "--disable_reg", default=True, action="store_false", dest="enable_reg"
+    )
     parser.add_argument(
         "--reg_weight", type=float, help="Regularization weight", default=0.01
     )
@@ -104,8 +125,20 @@ def parse_args():
     parser.add_argument("--no_optim", default=False, action="store_true")
     parser.add_argument("--imageselect", default=False, action="store_true")
     parser.add_argument("--memsave", default=False, action="store_true")
-    parser.add_argument("--device", type=str, help="Device to use", default="cuda")
-    parser.add_argument("--device_id", type=int, help="Device ID to use", default=None)
+    parser.add_argument("--dtype", type=str, help="Data type to use", default="float16")
+    parser.add_argument("--device_id", type=str, help="Device ID to use", default=None)
+    parser.add_argument(
+        "--cpu_offloading",
+        help="Enable CPU offloading",
+        default=False,
+        action="store_true",
+    )
+
+    # optional multi-step model
+    parser.add_argument("--enable_multi_apply", default=False, action="store_true")
+    parser.add_argument(
+        "--multi_step_model", type=str, help="Model to use", default="flux"
+    )
 
     args = parser.parse_args()
-    return args
+    return args

assets/example_prompts.txt

@@ -12,6 +12,7 @@ High quality photo of a monkey astronaut infront of the Eiffel tower
 A bird with 8 legs
 A brain riding a rocketship towards the moon
 A toaster riding a bike
+An upside-down snowman in a winter scene, with the snowman wearing a smile.
 A blue scooter is parked near a curb in front of a green vintage car
 A curious, orange fox and a fluffy, white rabbit, playing together in a lush, green meadow filled with yellow dandelions
 An epic oil painting: a red portal infront of a cityscape, a solitary figure, and a colorful sky over snowy mountains
@@ -19,6 +20,7 @@ A futuristic painting: Red car escapes giant shark's leap, right; ominous mounta
 A majestic, resilient sea ship navigates the icy wilderness in the style of Star Wars
 Dwayne Johnson depicted as a philosopher king in an academic painting by Greg Rutkowski
 Taylor Swift depicted as a prime minister in an academic painting by Kandinsky
+A swirling, multicolored portal emerges from the depths of an ocean of coffee, with waves of the rich liquid gently rippling outward. The portal engulfs a coffee cup, which serves as a gateway to a fantastical dimension. Digital art.
 A watercolor painting: a floating island, multiple animals under a majestic tree with golden leaves, and a vibrant rainbow stretching across a pastel sky
 A Japanese-style ink painting: a traditional wooden bridge, a pagoda, a lone samurai warrior, and cherry blossom petals over a tranquil river
 A retro-futuristic pixel art scene: a flying car, an imperial senate building, a green park, and a purple sunset

main.py

@@ -9,13 +9,13 @@ from pytorch_lightning import seed_everything
 from tqdm import tqdm
 
 from arguments import parse_args
-from models import get_model
+from models import get_model, get_multi_apply_fn
 from rewards import get_reward_losses
 from training import LatentNoiseTrainer, get_optimizer
 
 
 def setup(args):
-    #args = parse_args()
+    
     seed_everything(args.seed)
     bf.makedirs(f"{args.save_dir}/logs/{args.task}")
     # Set up logging and name settings
@@ -46,21 +46,22 @@ def setup(args):
     if args.device_id is not None:
         logging.info(f"Using CUDA device {args.device_id}")
         os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-        os.environ["CUDA_VISIBLE_DEVICE"] = args.device_id
-    if args.device == "cuda":
-        device = torch.device("cuda")
-    else:
-        device = torch.device("cpu")
-    # Set dtype to fp16
-    dtype = torch.float16
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id
+    device = torch.device("cuda")
+    if args.dtype == "float32":
+        dtype = torch.float32
+    elif args.dtype == "float16":
+        dtype = torch.float16
     # Get reward losses
     reward_losses = get_reward_losses(args, dtype, device, args.cache_dir)
 
     # Get model and noise trainer
-    sd_model = get_model(args.model, dtype, device, args.cache_dir, args.memsave)
+    pipe = get_model(
+        args.model, dtype, device, args.cache_dir, args.memsave, args.cpu_offloading
+    )
     trainer = LatentNoiseTrainer(
         reward_losses=reward_losses,
-        model=sd_model,
+        model=pipe,
         n_iters=args.n_iters,
         n_inference_steps=args.n_inference_steps,
         seed=args.seed,
@@ -75,41 +76,57 @@ def setup(args):
     )
 
     # Create latents
-    if args.model != "pixart":
-        height = sd_model.unet.config.sample_size * sd_model.vae_scale_factor
-        width = sd_model.unet.config.sample_size * sd_model.vae_scale_factor
+    if args.model == "flux":
+        # currently only support 512x512 generation
+        shape = (1, 16 * 64, 64)
+    elif args.model != "pixart":
+        height = pipe.unet.config.sample_size * pipe.vae_scale_factor
+        width = pipe.unet.config.sample_size * pipe.vae_scale_factor
         shape = (
             1,
-            sd_model.unet.in_channels,
-            height // sd_model.vae_scale_factor,
-            width // sd_model.vae_scale_factor,
+            pipe.unet.in_channels,
+            height // pipe.vae_scale_factor,
+            width // pipe.vae_scale_factor,
         )
     else:
-        height = sd_model.transformer.config.sample_size * sd_model.vae_scale_factor
-        width = sd_model.transformer.config.sample_size * sd_model.vae_scale_factor
+        height = pipe.transformer.config.sample_size * pipe.vae_scale_factor
+        width = pipe.transformer.config.sample_size * pipe.vae_scale_factor
         shape = (
             1,
-            sd_model.transformer.config.in_channels,
-            height // sd_model.vae_scale_factor,
-            width // sd_model.vae_scale_factor,
+            pipe.transformer.config.in_channels,
+            height // pipe.vae_scale_factor,
+            width // pipe.vae_scale_factor,
         )
     enable_grad = not args.no_optim
+    
+    if args.enable_multi_apply:
+        multi_apply_fn = get_multi_apply_fn(
+            model_type=args.multi_step_model,
+            seed=args.seed,
+            pipe=pipe,
+            cache_dir=args.cache_dir,
+            device=device,
+            dtype=dtype,
+        )
+    else:
+        multi_apply_fn = None    
 
-    return args, trainer, device, dtype, shape, enable_grad, settings
+    return args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings
 
-def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, progress_callback=None):
-    #args = parse_args()
+def execute_task(args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings, progress_callback=None):
+    
     if args.task == "single":
         init_latents = torch.randn(shape, device=device, dtype=dtype)
         latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
         optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
-        save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
+        save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt[:150]}"
         os.makedirs(f"{save_dir}", exist_ok=True)
-        best_image, total_init_rewards, total_best_rewards = trainer.train(
-            latents, args.prompt, optimizer, save_dir, progress_callback=progress_callback
+        init_image, best_image, total_init_rewards, total_best_rewards = trainer.train(
+            latents, args.prompt, optimizer, save_dir, multi_apply_fn, progress_callback=progress_callback
         )
         best_image.save(f"{save_dir}/best_image.png")
-        return best_image, total_init_rewards, total_best_rewards
+        #init_image.save(f"{save_dir}/init_image.png")
+
     elif args.task == "example-prompts":
         fo = open("assets/example_prompts.txt", "r")
         prompts = fo.readlines()
@@ -121,11 +138,11 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pro
             optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
 
             prompt = prompt.strip()
-            name = f"{i:03d}_{prompt}.png"
+            name = f"{i:03d}_{prompt[:150]}.png"
             save_dir = f"{args.save_dir}/{args.task}/{settings}/{name}"
             os.makedirs(save_dir, exist_ok=True)
-            best_image, init_rewards, best_rewards = trainer.train(
-                latents, prompt, optimizer, save_dir
+            init_image, best_image, init_rewards, best_rewards = trainer.train(
+                latents, prompt, optimizer, save_dir, multi_apply_fn
             )
             if i == 0:
                 total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
@@ -134,6 +151,7 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pro
                 total_best_rewards[k] += best_rewards[k]
                 total_init_rewards[k] += init_rewards[k]
             best_image.save(f"{save_dir}/best_image.png")
+            init_image.save(f"{save_dir}/init_image.png")
             logging.info(f"Initial rewards: {init_rewards}")
             logging.info(f"Best rewards: {best_rewards}")
         for k in total_best_rewards.keys():
@@ -159,8 +177,8 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pro
             optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
 
             prompt = prompt.strip()
-            best_image, init_rewards, best_rewards = trainer.train(
-                latents, prompt, optimizer
+            init_image, best_image, init_rewards, best_rewards = trainer.train(
+                latents, prompt, optimizer, None, multi_apply_fn
             )
             if i == 0:
                 total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
@@ -186,8 +204,8 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pro
                 f"{args.save_dir}/{args.task}/{settings}/{index}", exist_ok=True
             )
             prompt = sample["Prompt"]
-            best_image, init_rewards, best_rewards = trainer.train(
-                latents, prompt, optimizer
+            init_image, best_image, init_rewards, best_rewards = trainer.train(
+                latents, prompt, optimizer, multi_apply_fn
             )
             best_image.save(
                 f"{args.save_dir}/{args.task}/{settings}/{index}/best_image.png"
@@ -252,8 +270,8 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pro
             optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
 
             prompt = metadata["prompt"]
-            best_image, init_rewards, best_rewards = trainer.train(
-                latents, prompt, optimizer
+            init_image, best_image, init_rewards, best_rewards = trainer.train(
+                latents, prompt, optimizer, None, multi_apply_fn
             )
             logging.info(f"Initial rewards: {init_rewards}")
             logging.info(f"Best rewards: {best_rewards}")
@@ -279,9 +297,8 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pro
 
 def main():
     args = parse_args()
-    args, trainer, device, dtype, shape, enable_grad, settings = setup(args)
-    execute_task(args, trainer, device, dtype, shape, enable_grad, settings)
-
+    args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings = setup(args)
+    execute_task(args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings)
 
 if __name__ == "__main__":
     main()

models/RewardFlux.py

@@ -0,0 +1,436 @@
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from diffusers import (
+    FluxPipeline,
+    AutoencoderKL,
+    FluxTransformer2DModel,
+    FlowMatchEulerDiscreteScheduler,
+)
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import FluxPipeline
+
+        >>> pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+        >>> pipe.to("cuda")
+        >>> prompt = "A cat holding a sign that says hello world"
+        >>> # Depending on the variant being used, the pipeline call will slightly vary.
+        >>> # Refer to the pipeline documentation for more details.
+        >>> image = pipe(prompt, num_inference_steps=4, guidance_scale=0.0).images[0]
+        >>> image.save("flux.png")
+        ```
+"""
+
+
+def freeze_params(params):
+    for param in params:
+        param.requires_grad = False
+
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError(
+            "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
+        )
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class RewardFluxPipeline(FluxPipeline):
+    r"""
+    The Flux pipeline for text-to-image generation.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Args:
+        transformer ([`FluxTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        text_encoder_2 ([`T5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`T5TokenizerFast`):
+            Second Tokenizer of class
+            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: T5EncoderModel,
+        tokenizer_2: T5TokenizerFast,
+        transformer: FluxTransformer2DModel,
+        memsave: bool = False,
+    ):
+        # optionally enable memsave_torch
+        if memsave:
+            import memsave_torch.nn
+
+            vae = memsave_torch.nn.convert_to_memory_saving(vae)
+            transformer = memsave_torch.nn.convert_to_memory_saving(transformer)
+            text_encoder = memsave_torch.nn.convert_to_memory_saving(text_encoder)
+            text_encoder_2 = memsave_torch.nn.convert_to_memory_saving(text_encoder_2)
+        # enable checkpointing
+        text_encoder.gradient_checkpointing_enable()
+        text_encoder_2.gradient_checkpointing_enable()
+        transformer.enable_gradient_checkpointing()
+        vae.eval()
+        text_encoder.eval()
+        text_encoder_2.eval()
+        transformer.eval()
+
+        # freeze diffusion parameters
+        freeze_params(vae.parameters())
+        freeze_params(transformer.parameters())
+        freeze_params(text_encoder.parameters())
+        freeze_params(text_encoder_2.parameters())
+        super().__init__(
+            scheduler,
+            vae,
+            text_encoder,
+            tokenizer,
+            text_encoder_2,
+            tokenizer_2,
+            transformer,
+        )
+
+    def apply(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = 512,
+        width: Optional[int] = 512,
+        num_inference_steps: int = 1,
+        timesteps: List[int] = None,
+        guidance_scale: float = 0.0,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 256,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None)
+            if self.joint_attention_kwargs is not None
+            else None
+        )
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas,
+            mu=mu,
+        )
+        self._num_timesteps = len(timesteps)
+
+        # 6. Denoising loop:
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+
+            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+            # handle guidance
+            if self.transformer.config.guidance_embeds:
+                guidance = torch.tensor([guidance_scale], device=device)
+                guidance = guidance.expand(latents.shape[0])
+            else:
+                guidance = None
+
+            noise_pred = self.transformer(
+                hidden_states=latents,
+                # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
+                timestep=timestep / 1000,
+                guidance=guidance,
+                pooled_projections=pooled_prompt_embeds,
+                encoder_hidden_states=prompt_embeds,
+                txt_ids=text_ids,
+                img_ids=latent_image_ids,
+                joint_attention_kwargs=self.joint_attention_kwargs,
+                return_dict=False,
+            )[0]
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents_dtype = latents.dtype
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+            if latents.dtype != latents_dtype:
+                if torch.backends.mps.is_available():
+                    # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                    latents = latents.to(latents_dtype)
+
+            if callback_on_step_end is not None:
+                callback_kwargs = {}
+                for k in callback_on_step_end_tensor_inputs:
+                    callback_kwargs[k] = locals()[k]
+                callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                latents = callback_outputs.pop("latents", latents)
+                prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = self._unpack_latents(
+                latents, height, width, self.vae_scale_factor
+            )
+            latents = (
+                latents / self.vae.config.scaling_factor
+            ) + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = (image / 2 + 0.5).clamp(0, 1)
+            # image = self.image_processor.postprocess(image, output_type=output_type)
+
+        return image

models/RewardPixart.py

@@ -2,8 +2,7 @@ from typing import List, Optional, Union
 
 import torch
 from diffusers import PixArtAlphaPipeline
-from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha import \
-    retrieve_timesteps
+from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha import retrieve_timesteps
 
 
 def freeze_params(params):
@@ -391,4 +390,4 @@ ASPECT_RATIO_512_BIN = {
     "2.5": [800.0, 320.0],
     "3.0": [864.0, 288.0],
     "4.0": [1024.0, 256.0],
-}
+}

models/RewardStableDiffusion.py

@@ -274,4 +274,4 @@ def retrieve_timesteps(
     else:
         scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
         timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
+    return timesteps, num_inference_steps

models/RewardStableDiffusionXL.py

@@ -1,14 +1,18 @@
 from typing import List, Optional, Union
 
 import torch
-from diffusers import (AutoencoderKL, StableDiffusionXLPipeline,
-                       UNet2DConditionModel)
-from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import \
-    retrieve_timesteps
+from diffusers import AutoencoderKL, StableDiffusionXLPipeline, UNet2DConditionModel
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
+    retrieve_timesteps,
+)
 from diffusers.schedulers import KarrasDiffusionSchedulers
-from transformers import (CLIPImageProcessor, CLIPTextModel,
-                          CLIPTextModelWithProjection, CLIPTokenizer,
-                          CLIPVisionModelWithProjection)
+from transformers import (
+    CLIPImageProcessor,
+    CLIPTextModel,
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    CLIPVisionModelWithProjection,
+)
 
 
 def freeze_params(params):
@@ -317,4 +321,4 @@ class RewardStableDiffusionXL(StableDiffusionXLPipeline):
         # Offload all models
         self.maybe_free_model_hooks()
 
-        return image
+        return image

models/__init__.py

@@ -1 +1 @@
-from .utils import get_model
+from .utils import get_model, get_multi_apply_fn

models/utils.py

@@ -1,15 +1,22 @@
 import logging
-
+from typing import Any, Optional
 import torch
-from diffusers import (AutoencoderKL, DDPMScheduler,
-                       EulerAncestralDiscreteScheduler, LCMScheduler,
-                       Transformer2DModel, UNet2DConditionModel)
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    EulerDiscreteScheduler,
+    EulerAncestralDiscreteScheduler,
+    LCMScheduler,
+    Transformer2DModel,
+    UNet2DConditionModel,
+)
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
 from models.RewardPixart import RewardPixartPipeline, freeze_params
 from models.RewardStableDiffusion import RewardStableDiffusion
 from models.RewardStableDiffusionXL import RewardStableDiffusionXL
+from models.RewardFlux import RewardFluxPipeline
 
 
 def get_model(
@@ -18,6 +25,7 @@ def get_model(
     device: torch.device,
     cache_dir: str,
     memsave: bool = False,
+    enable_sequential_cpu_offload: bool = False,
 ):
     logging.info(f"Loading model: {model_name}")
     if model_name == "sd-turbo":
@@ -103,7 +111,80 @@ def get_model(
         pipe = pipe.to(device, dtype)
         # upcast vae
         pipe.vae = pipe.vae.to(dtype=torch.float32)
-        # pipe.enable_sequential_cpu_offload()
+    elif model_name == "flux":
+        pipe = RewardFluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-schnell",
+            torch_dtype=torch.bfloat16,
+            cache_dir=cache_dir,
+        )
+        pipe.to(device, dtype)
     else:
         raise ValueError(f"Unknown model name: {model_name}")
+    if enable_sequential_cpu_offload:
+        pipe.enable_sequential_cpu_offload()
     return pipe
+
+
+def get_multi_apply_fn(
+    model_type: str,
+    seed: int,
+    pipe: Optional[Any] = None,
+    cache_dir: Optional[str] = None,
+    device: Optional[torch.device] = None,
+    dtype: Optional[torch.dtype] = None,
+):
+    generator = torch.Generator("cuda").manual_seed(seed)
+    if model_type == "flux":
+        return lambda latents, prompt: torch.no_grad(pipe.apply)(
+            latents=latents,
+            prompt=prompt,
+            num_inference_steps=4,
+            generator=generator,
+        )
+    elif model_type == "sdxl":
+        vae = AutoencoderKL.from_pretrained(
+            "madebyollin/sdxl-vae-fp16-fix",
+            torch_dtype=torch.float16,
+            cache_dir=cache_dir,
+        )
+        pipe = RewardStableDiffusionXL.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            torch_dtype=torch.float16,
+            variant="fp16",
+            vae=vae,
+            use_safetensors=True,
+            cache_dir=cache_dir,
+        )
+        pipe = pipe.to(device, dtype)
+        pipe.enable_sequential_cpu_offload()
+        return lambda latents, prompt: torch.no_grad(pipe.apply)(
+            latents=latents,
+            prompt=prompt,
+            guidance_scale=5.0,
+            num_inference_steps=50,
+            generator=generator,
+        )
+    elif model_type == "sd2":
+        sd2_base = "stabilityai/stable-diffusion-2-1-base"
+        scheduler = EulerDiscreteScheduler.from_pretrained(
+            sd2_base,
+            subfolder="scheduler",
+            cache_dir=cache_dir,
+        )
+        pipe = RewardStableDiffusion.from_pretrained(
+            sd2_base,
+            torch_dtype=dtype,
+            cache_dir=cache_dir,
+            scheduler=scheduler,
+        )
+        pipe = pipe.to(device, dtype)
+        pipe.enable_sequential_cpu_offload()
+        return lambda latents, prompt: torch.no_grad(pipe.apply)(
+            latents=latents,
+            prompt=prompt,
+            guidance_scale=7.5,
+            num_inference_steps=50,
+            generator=generator,
+        )
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")

requirements.txt

@@ -3,7 +3,7 @@ torchvision==0.18.0
 pytorch-lightning==2.2
 datasets==2.18
 transformers==4.38.2
-diffusers==0.28
+diffusers
 hpsv2==1.2
 hpsv2x==1.2.0
 image-reward==1.5

training/trainer.py

@@ -51,13 +51,16 @@ class LatentNoiseTrainer:
         prompt: str,
         optimizer: torch.optim.Optimizer,
         save_dir: Optional[str] = None,
+        multi_apply_fn=None,
         progress_callback=None,
     ) -> Tuple[PIL.Image.Image, Dict[str, float], Dict[str, float]]:
         logging.info(f"Optimizing latents for prompt '{prompt}'.")
         best_loss = torch.inf
         best_image = None
+        initial_image = None
         initial_rewards = None
         best_rewards = None
+        best_latents = None
         latent_dim = math.prod(latents.shape[1:])
         for iteration in range(self.n_iters):
             to_log = ""
@@ -76,11 +79,17 @@ class LatentNoiseTrainer:
                 )
             else:
                 image = self.model.apply(
-                    latents,
-                    prompt,
+                    latents=latents,
+                    prompt=prompt,
                     generator=generator,
                     num_inference_steps=self.n_inference_steps,
                 )
+            if initial_image is None and multi_apply_fn is not None:
+                multi_step_image = multi_apply_fn(latents.detach(), prompt)
+                image_numpy = (
+                    multi_step_image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
+                )
+                initial_image = DiffusionPipeline.numpy_to_pil(image_numpy)[0]
             if self.no_optim:
                 best_image = image
                 break
@@ -107,12 +116,11 @@ class LatentNoiseTrainer:
                 total_loss += regularization.to(total_loss.dtype)
             if self.log_metrics:
                 logging.info(f"Iteration {iteration}: {to_log}")
-            if initial_rewards is None:
-                initial_rewards = rewards
             if total_reward_loss < best_loss:
                 best_loss = total_reward_loss
                 best_image = image
                 best_rewards = rewards
+                best_latents = latents.detach().cpu()
             if iteration != self.n_iters - 1 and not self.imageselect:
                 total_loss.backward()
                 torch.nn.utils.clip_grad_norm_(latents, self.grad_clip)
@@ -121,8 +129,16 @@ class LatentNoiseTrainer:
                 image_numpy = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
                 image_pil = DiffusionPipeline.numpy_to_pil(image_numpy)[0]
                 image_pil.save(f"{save_dir}/{iteration}.png")
+            if initial_rewards is None:
+                initial_rewards = rewards
             if progress_callback:
                 progress_callback(iteration + 1)
         image_numpy = best_image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
-        image_pil = DiffusionPipeline.numpy_to_pil(image_numpy)[0]
-        return image_pil, initial_rewards, best_rewards
+        best_image_pil = DiffusionPipeline.numpy_to_pil(image_numpy)[0]
+        if multi_apply_fn is not None:
+            multi_step_image = multi_apply_fn(best_latents.to("cuda"), prompt)
+            image_numpy = (
+                multi_step_image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
+            )
+            best_image_pil = DiffusionPipeline.numpy_to_pil(image_numpy)[0]
+        return initial_image, best_image_pil, initial_rewards, best_rewards