app.py

import gradio as gr

from absl import flags
from absl import app
from ml_collections import config_flags
import os

import spaces #[uncomment to use ZeroGPU]
import torch


import io
import random
import tempfile

import numpy as np
import torch
import torch.nn.functional as F
from torchvision.utils import save_image
from torchvision.transforms import ToPILImage
from huggingface_hub import hf_hub_download

from absl import logging
import ml_collections

from diffusion.flow_matching import ODEEulerFlowMatchingSolver
import utils
import libs.autoencoder
from libs.clip import FrozenCLIPEmbedder
from configs import t2i_512px_clip_dimr, t2i_256px_clip_dimr


def unpreprocess(x: torch.Tensor) -> torch.Tensor:
    x = 0.5 * (x + 1.0)
    x.clamp_(0.0, 1.0)
    return x
    
def cosine_similarity_torch(latent1: torch.Tensor, latent2: torch.Tensor) -> torch.Tensor:
    latent1_flat = latent1.view(-1)
    latent2_flat = latent2.view(-1)
    cosine_similarity = F.cosine_similarity(
        latent1_flat.unsqueeze(0), latent2_flat.unsqueeze(0), dim=1
    )
    return cosine_similarity

def kl_divergence(latent1: torch.Tensor, latent2: torch.Tensor) -> torch.Tensor:
    latent1_prob = F.softmax(latent1, dim=-1)
    latent2_prob = F.softmax(latent2, dim=-1)
    latent1_log_prob = torch.log(latent1_prob)
    kl_div = F.kl_div(latent1_log_prob, latent2_prob, reduction="batchmean")
    return kl_div

def batch_decode(_z: torch.Tensor, decode, batch_size: int = 5) -> torch.Tensor:
    num_samples = _z.size(0)
    decoded_batches = []

    for i in range(0, num_samples, batch_size):
        batch = _z[i : i + batch_size]
        decoded_batch = decode(batch)
        decoded_batches.append(decoded_batch)

    return torch.cat(decoded_batches, dim=0)

def get_caption(llm: str, text_model, prompt_dict: dict, batch_size: int):
    if batch_size == 3:
        # Only addition or only subtraction mode.
        assert len(prompt_dict) == 2, "Expected 2 prompts for batch_size 3."
        batch_prompts = list(prompt_dict.values()) + [" "]
    elif batch_size == 4:
        # Addition and subtraction mode.
        assert len(prompt_dict) == 3, "Expected 3 prompts for batch_size 4."
        batch_prompts = list(prompt_dict.values()) + [" "]
    elif batch_size >= 5:
        # Linear interpolation mode.
        assert len(prompt_dict) == 2, "Expected 2 prompts for linear interpolation."
        batch_prompts = [prompt_dict["prompt_1"]] + [" "] * (batch_size - 2) + [prompt_dict["prompt_2"]]
    else:
        raise ValueError(f"Unsupported batch_size: {batch_size}")

    if llm == "clip":
        latent, latent_and_others = text_model.encode(batch_prompts)
        context = latent_and_others["token_embedding"].detach()
    elif llm == "t5":
        latent, latent_and_others = text_model.get_text_embeddings(batch_prompts)
        context = (latent_and_others["token_embedding"] * 10.0).detach()
    else:
        raise NotImplementedError(f"Language model {llm} not supported.")

    token_mask = latent_and_others["token_mask"].detach()
    tokens = latent_and_others["tokens"].detach()
    captions = batch_prompts

    return context, token_mask, tokens, captions

# Load configuration and initialize models.
# config_dict = t2i_512px_clip_dimr.get_config()
config_dict = t2i_256px_clip_dimr.get_config()
config_1 = ml_collections.ConfigDict(config_dict)
config_dict = t2i_512px_clip_dimr.get_config()
config_2 = ml_collections.ConfigDict(config_dict)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")

# Freeze configuration.
config_1 = ml_collections.FrozenConfigDict(config_1)
config_2 = ml_collections.FrozenConfigDict(config_2)

torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024  # Currently not used.

# Load the main diffusion model.
repo_id = "QHL067/CrossFlow"
# filename = "pretrained_models/t2i_512px_clip_dimr.pth"
filename = "pretrained_models/t2i_256px_clip_dimr.pth"
checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
nnet_1 = utils.get_nnet(**config_1.nnet)
nnet_1 = nnet_1.to(device)
state_dict = torch.load(checkpoint_path, map_location=device)
nnet_1.load_state_dict(state_dict)
nnet_1.eval()

filename = "pretrained_models/t2i_512px_clip_dimr.pth"
checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
nnet_2 = utils.get_nnet(**config_2.nnet)
nnet_2 = nnet_2.to(device)
state_dict = torch.load(checkpoint_path, map_location=device)
nnet_2.load_state_dict(state_dict)
nnet_2.eval()

# Initialize text model.
llm = "clip"
clip = FrozenCLIPEmbedder()
clip.eval()
clip.to(device)

# Load autoencoder.
autoencoder = libs.autoencoder.get_model(**config_1.autoencoder)
autoencoder.to(device)


@torch.cuda.amp.autocast()
def encode(_batch: torch.Tensor) -> torch.Tensor:
    """Encode a batch of images using the autoencoder."""
    return autoencoder.encode(_batch)


@torch.cuda.amp.autocast()
def decode(_batch: torch.Tensor) -> torch.Tensor:
    """Decode a batch of latent vectors using the autoencoder."""
    return autoencoder.decode(_batch)


@spaces.GPU #[uncomment to use ZeroGPU]
def infer(
    prompt1,
    prompt2,
    seed,
    randomize_seed,
    guidance_scale,
    num_inference_steps,
    num_of_interpolation,
    operation_mode,
    save_gpu_memory=True,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    torch.manual_seed(seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)

    # Only support interpolation in this implementation.
    prompt_dict = {"prompt_1": prompt1, "prompt_2": prompt2}
    for key, value in prompt_dict.items():
        assert value is not None, f"{key} must not be None."
    if operation_mode == 'Interpolation':
        assert num_of_interpolation >= 5, "For linear interpolation, please sample at least five images."
    else:
        assert num_of_interpolation == 3, "For arithmetic, please sample three images."

    if num_of_interpolation == 3:
        nnet = nnet_2
        config = config_2
    else:
        nnet = nnet_1
        config = config_1
    # nnet = nnet_1
    # config = config_1
    
    # Get text embeddings and tokens.
    _context, _token_mask, _token, _caption = get_caption(
        llm, clip, prompt_dict=prompt_dict, batch_size=num_of_interpolation
    )

    with torch.no_grad():
        _z_gaussian = torch.randn(num_of_interpolation, *config.z_shape, device=device)
        _z_x0, _mu, _log_var = nnet(
            _context, text_encoder=True, shape=_z_gaussian.shape, mask=_token_mask
        )
        _z_init = _z_x0.reshape(_z_gaussian.shape)

        # Prepare the initial latent representations based on the number of interpolations.
        if num_of_interpolation == 3:
            # Addition or subtraction mode.
            if operation_mode == 'Addition':
                z_init_temp = _z_init[0] + _z_init[1]
            elif operation_mode == 'Subtraction':
                z_init_temp = _z_init[0] - _z_init[1]
            else:
                raise NotImplementedError("Either prompt_a or prompt_s must be provided for 3-sample mode.")
            mean = z_init_temp.mean()
            std = z_init_temp.std()
            _z_init[2] = (z_init_temp - mean) / std

        elif num_of_interpolation == 4:
            raise ValueError("Unsupported number of interpolations.")

        elif num_of_interpolation >= 5:
            tensor_a = _z_init[0]
            tensor_b = _z_init[-1]
            num_interpolations = num_of_interpolation - 2
            interpolations = [
                tensor_a + (tensor_b - tensor_a) * (i / (num_interpolations + 1))
                for i in range(1, num_interpolations + 1)
            ]
            _z_init = torch.stack([tensor_a] + interpolations + [tensor_b], dim=0)

        else:
            raise ValueError("Unsupported number of interpolations.")

        assert guidance_scale > 1, "Guidance scale must be greater than 1."

        has_null_indicator = hasattr(config.nnet.model_args, "cfg_indicator")
        ode_solver = ODEEulerFlowMatchingSolver(
            nnet,
            bdv_model_fn=None,
            step_size_type="step_in_dsigma",
            guidance_scale=guidance_scale,
        )
        _z, _ = ode_solver.sample(
            x_T=_z_init,
            batch_size=num_of_interpolation,
            sample_steps=num_inference_steps,
            unconditional_guidance_scale=guidance_scale,
            has_null_indicator=has_null_indicator,
        )

        print("+++++"*20)
        print("Now, save images")
        print("+++++"*20)
        
        if save_gpu_memory:
            image_unprocessed = batch_decode(_z, decode)
        else:
            image_unprocessed = decode(_z)

        samples = unpreprocess(image_unprocessed).contiguous()


    to_pil = ToPILImage()
    pil_images = [to_pil(img) for img in samples]

    if num_of_interpolation == 3:
        return pil_images[0], pil_images[1], pil_images[2], seed
    
    else:
        first_image = pil_images[0]
        last_image = pil_images[-1]

        gif_buffer = io.BytesIO()
        pil_images[0].save(gif_buffer, format="GIF", save_all=True, append_images=pil_images[1:], duration=200, loop=0)
        gif_buffer.seek(0)
        gif_bytes = gif_buffer.read()

        # Save the GIF bytes to a temporary file and get its path
        temp_gif = tempfile.NamedTemporaryFile(delete=False, suffix=".gif")
        temp_gif.write(gif_bytes)
        temp_gif.close()
        gif_path = temp_gif.name

        return first_image, last_image, gif_path, seed
    # return first_image, last_image, seed


# examples = [
#     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
#     "An astronaut riding a green horse",
#     "A delicious ceviche cheesecake slice",
# ]

def infer_tab1(prompt1, prompt2, seed, randomize_seed, guidance_scale, num_inference_steps, num_of_interpolation):
    default_op = "Interpolation"
    return infer(prompt1, prompt2, seed, randomize_seed, guidance_scale, num_inference_steps, num_of_interpolation, default_op)

# Wrapper for Tab 2: Uses operation_mode and fixes num_of_interpolation to 3.
def infer_tab2(prompt1, prompt2, seed, randomize_seed, guidance_scale, num_inference_steps, operation_mode):
    default_interpolation = 3
    return infer(prompt1, prompt2, seed, randomize_seed, guidance_scale, num_inference_steps, default_interpolation, operation_mode)

examples_1 = [
    ["A robot cooking dinner in the kitchen", "An orange cat wearing sunglasses on a ship"],
]

examples_2 = [
    ["A dog wearing sunglasses", "a hat"],
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""
# with gr.Blocks(css=css) as demo:
#     with gr.Column(elem_id="col-container"):
#         gr.Markdown("# CrossFlow")
#         gr.Markdown("[CrossFlow](https://cross-flow.github.io/) directly transforms text representations into images for text-to-image generation, without the need for both the noise distribution and conditioning mechanism.")
#         gr.Markdown("This direct mapping enables meaningful 'Linear Interpolation' and 'Arithmetic Operations' in the text latent space, as demonstrated here.")
#         with gr.Tabs():
#             with gr.Tab("Linear Interpolation"):
#                 gr.Markdown("This demo uses 256px images, 25 sampling steps (instead of 50), and 10 interpolations (instead of 50) to conserve GPU memory. For better results, see the original [code](https://github.com/qihao067/CrossFlow). (You may adjust them in Advanced Settings, but doing so may trigger OOM errors.)")
#                 # gr.Markdown("CrossFlow directly transforms text representations into images for text-to-image generation, enabling interpolation in the input text latent space.")

#                 with gr.Row():
#                     prompt1 = gr.Text(
#                         label="Prompt_1",
#                         show_label=False,
#                         max_lines=1,
#                         placeholder="Enter your prompt for the first image",
#                         container=False,
#                     )
                
#                 with gr.Row():
#                     prompt2 = gr.Text(
#                         label="Prompt_2",
#                         show_label=False,
#                         max_lines=1,
#                         placeholder="Enter your prompt for the second image",
#                         container=False,
#                     )

#                 with gr.Row():
#                     run_button = gr.Button("Run", scale=0, variant="primary")

#                 # Create separate outputs for the first image, last image, and the animated GIF
#                 first_image_output = gr.Image(label="Image of the first prompt", show_label=True)
#                 last_image_output = gr.Image(label="Image of the second prompt", show_label=True)
#                 gif_output = gr.Image(label="Linear interpolation", show_label=True)

#                 with gr.Accordion("Advanced Settings", open=False):
#                     seed = gr.Slider(
#                         label="Seed",
#                         minimum=0,
#                         maximum=MAX_SEED,
#                         step=1,
#                         value=0,
#                     )

#                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

#                     with gr.Row():
#                         guidance_scale = gr.Slider(
#                             label="Guidance scale",
#                             minimum=0.0,
#                             maximum=10.0,
#                             step=0.1,
#                             value=7.0,  # Replace with defaults that work for your model
#                         )
#                     with gr.Row():
#                         num_inference_steps = gr.Slider(
#                             label="Number of inference steps - 50 inference steps are recommended; but you can reduce to 20 if the demo fails.",
#                             minimum=1,
#                             maximum=50,
#                             step=1,
#                             value=25,  # Replace with defaults that work for your model
#                         )
#                     with gr.Row():
#                         num_of_interpolation = gr.Slider(
#                             label="Number of images for interpolation - More images yield smoother transitions but require more resources and may fail.",
#                             minimum=5,
#                             maximum=50,
#                             step=1,
#                             value=10,  # Replace with defaults that work for your model
#                         )

#                 gr.Examples(examples=examples, inputs=[prompt1, prompt2])

#             with gr.Tab("Arithmetic Operations"):
#                 # The second tab is currently empty. You can add more components later.
#                 gr.Markdown("This demo only supports addition or subtraction between two text latents ('Prompt_1 + Prompt_2' or 'Prompt_1 - Prompt_2'). For the other arithmetic operations, see the original [code](https://github.com/qihao067/CrossFlow).")
#                 with gr.Row():
#                     prompt1 = gr.Text(
#                         label="Prompt_1",
#                         show_label=False,
#                         max_lines=1,
#                         placeholder="Enter your prompt for the first image",
#                         container=False,
#                     )
                
#                 with gr.Row():
#                     prompt2 = gr.Text(
#                         label="Prompt_2",
#                         show_label=False,
#                         max_lines=1,
#                         placeholder="Enter your prompt for the second image",
#                         container=False,
#                     )

#                 with gr.Row():
#                     operation_mode = gr.Radio(
#                         choices=["Addition", "Subtraction"],
#                         label="Operation Mode",
#                         value="Addition",
#                     )
#                 with gr.Row():
#                     run_button = gr.Button("Run", scale=0, variant="primary")

#                 # Create separate outputs for the first image, last image, and the animated GIF
#                 first_image_output = gr.Image(label="Image of the first prompt", show_label=True)
#                 last_image_output = gr.Image(label="Image of the second prompt", show_label=True)
#                 gif_output = gr.Image(label="Linear interpolation", show_label=True)

#                 with gr.Accordion("Advanced Settings", open=False):
#                     seed = gr.Slider(
#                         label="Seed",
#                         minimum=0,
#                         maximum=MAX_SEED,
#                         step=1,
#                         value=0,
#                     )

#                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

#                     with gr.Row():
#                         guidance_scale = gr.Slider(
#                             label="Guidance scale",
#                             minimum=0.0,
#                             maximum=10.0,
#                             step=0.1,
#                             value=7.0,  # Replace with defaults that work for your model
#                         )
#                     # with gr.Row():
#                     #     num_inference_steps = gr.Slider(
#                     #         label="Number of inference steps - 50 inference steps are recommended; but you can reduce to 20 if the demo fails.",
#                     #         minimum=1,
#                     #         maximum=50,
#                     #         step=1,
#                     #         value=55,  # Replace with defaults that work for your model
#                     #     )

#                     with gr.Row():
#                         num_of_interpolation = gr.Slider(
#                             label="Number of images for interpolation - More images yield smoother transitions but require more resources and may fail.",
#                             minimum=5,
#                             maximum=50,
#                             step=1,
#                             value=50,  # Replace with defaults that work for your model
#                         )

#                 gr.Examples(examples=examples, inputs=[prompt1, prompt2])

#     gr.on(
#         triggers=[run_button.click, prompt1.submit, prompt2.submit],
#         fn=infer,
#         inputs=[
#             prompt1,
#             prompt2,
#             seed,
#             randomize_seed,
#             guidance_scale,
#             num_inference_steps,
#             num_of_interpolation,
#         ],
#         outputs=[first_image_output, last_image_output, gif_output, seed],
#         # outputs=[first_image_output, last_image_output, seed],
#     )

with gr.Blocks(css=css) as demo:
    gr.Markdown("# CrossFlow")
    gr.Markdown("[CVPR 2025] Flowing from Words to Pixels: A Framework for Cross-Modality Evolution")
    gr.Markdown("[CrossFlow](https://cross-flow.github.io/) achieves text-to-image generation by directly mapping text representations (source distribution) to images (target distribution), eliminating the need for a noise distribution or conditioning mechanism.")
    gr.Markdown("This direct mapping enables meaningful 'Linear Interpolation' and 'Arithmetic Operations' in the text latent space, as demonstrated here.")
    with gr.Tabs():
        # --- Tab 1: Interpolation Mode (no operation_mode) ---
        with gr.Tab("[Linear Interpolation]"):
            gr.Markdown("This demo uses 256px images, 25 sampling steps (instead of 50), and 10 interpolations (instead of 50) to conserve GPU memory.")
            gr.Markdown("**You will get much better results with the original [code](https://github.com/qihao067/CrossFlow)**. (You may also adjust the sampling steps and interpolations in Advanced Settings, but doing so may trigger OOM errors.)")
            
            prompt1_tab1 = gr.Text(placeholder="Prompt for first image", label="Prompt 1")
            prompt2_tab1 = gr.Text(placeholder="Prompt for second image", label="Prompt 2")
            seed_tab1 = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=0, label="Seed")
            randomize_seed_tab1 = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Accordion("Advanced Settings", open=False):
                guidance_scale_tab1 = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=7.0, label="Guidance Scale")
                num_inference_steps_tab1 = gr.Slider(minimum=1, maximum=50, step=1, value=25, label="Number of Inference Steps")
                num_of_interpolation_tab1 = gr.Slider(minimum=5, maximum=50, step=1, value=10, label="Number of Images for Interpolation")
            run_button_tab1 = gr.Button("Run")
            
            first_image_output_tab1 = gr.Image(label="Image of the first prompt")
            last_image_output_tab1 = gr.Image(label="Image of the second prompt")
            gif_output_tab1 = gr.Image(label="Linear interpolation")
            
            run_button_tab1.click(
                fn=infer_tab1,
                inputs=[
                    prompt1_tab1,
                    prompt2_tab1,
                    seed_tab1,
                    randomize_seed_tab1,
                    guidance_scale_tab1,
                    num_inference_steps_tab1,
                    num_of_interpolation_tab1
                ],
                outputs=[first_image_output_tab1, last_image_output_tab1, gif_output_tab1, seed_tab1]
            )

            gr.Examples(examples=examples_1, inputs=[prompt1_tab1, prompt2_tab1])
        
        # --- Tab 2: Operation Mode (no num_of_interpolation) ---
        with gr.Tab("[Arithmetic Operations]"):
            gr.Markdown("This demo only supports addition or subtraction between two text latents, i.e., 'VE(Prompt_1) + VE(Prompt_2)' or 'VE(Prompt_1) - VE(Prompt_2)'. For the other arithmetic operations, see the original [code](https://github.com/qihao067/CrossFlow).")
            
            prompt1_tab2 = gr.Text(placeholder="Prompt for first image", label="Prompt 1")
            prompt2_tab2 = gr.Text(placeholder="Prompt for second image", label="Prompt 2")
            seed_tab2 = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=0, label="Seed")
            randomize_seed_tab2 = gr.Checkbox(label="Randomize seed", value=True)
            guidance_scale_tab2 = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=7.0, label="Guidance Scale")
            num_inference_steps_tab2 = gr.Slider(minimum=1, maximum=50, step=1, value=50, label="Number of Inference Steps")
            operation_mode_tab2 = gr.Radio(choices=["Addition", "Subtraction"], label="Operation Mode", value="Addition")
            run_button_tab2 = gr.Button("Run")
            
            first_image_output_tab2 = gr.Image(label="Image of the first prompt")
            last_image_output_tab2 = gr.Image(label="Image of the second prompt")
            gif_output_tab2 = gr.Image(label="Resulting image produced by the arithmetic operations.")
            
            run_button_tab2.click(
                fn=infer_tab2,
                inputs=[
                    prompt1_tab2,
                    prompt2_tab2,
                    seed_tab2,
                    randomize_seed_tab2,
                    guidance_scale_tab2,
                    num_inference_steps_tab2,
                    operation_mode_tab2
                ],
                outputs=[first_image_output_tab2, last_image_output_tab2, gif_output_tab2, seed_tab2]
            )

            gr.Examples(examples=examples_2, inputs=[prompt1_tab2, prompt2_tab2])


if __name__ == "__main__":
    demo.launch()