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 = 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 = ml_collections.FrozenConfigDict(config)

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 = utils.get_nnet(**config.nnet)
nnet = nnet.to(device)
state_dict = torch.load(checkpoint_path, map_location=device)
nnet.load_state_dict(state_dict)
nnet.eval()

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

# Load autoencoder.
autoencoder = libs.autoencoder.get_model(**config.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."

    # 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':
                assert config.prompt_s is None, "Only one of prompt_a or prompt_s should be provided."
                z_init_temp = _z_init[0] + _z_init[1]
            elif operation_mode == 'Subtraction':
                assert config.prompt_a is None, "Only one of prompt_a or prompt_s should be provided."
                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 corgi in the park", "red 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("[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():
        # --- 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. 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.)")
            
            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 ('Prompt_1 + Prompt_2' or 'Prompt_1 - 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()