app.py

import os
import uuid
import gradio as gr
import spaces
from clip_slider_pipeline import CLIPSliderFlux
from diffusers import FluxPipeline, AutoencoderTiny
import torch
import numpy as np
import cv2
from PIL import Image
from diffusers.utils import load_image
from diffusers.utils import export_to_video
import random
from transformers import pipeline
# Translation model load
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

# English menu labels
english_labels = {
    "Prompt": "Prompt",
    "1st direction to steer": "1st Direction",
    "2nd direction to steer": "2nd Direction",
    "Strength": "Strength",
    "Generate directions": "Generate Directions",
    "Generated Images": "Generated Images",
    "From 1st to 2nd direction": "From 1st to 2nd Direction",
    "Strip": "Image Strip",
    "Looping video": "Looping Video",
    "Advanced options": "Advanced Options",
    "Num of intermediate images": "Number of Intermediate Images",
    "Num iterations for clip directions": "Number of CLIP Direction Iterations",
    "Num inference steps": "Number of Inference Steps",
    "Guidance scale": "Guidance Scale",
    "Randomize seed": "Randomize Seed",
    "Seed": "Seed"
}

# load pipelines
base_model = "black-forest-labs/FLUX.1-schnell"

taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to("cuda")
pipe = FluxPipeline.from_pretrained(base_model,
                                    vae=taef1,
                                    torch_dtype=torch.bfloat16)

pipe.transformer.to(memory_format=torch.channels_last)
clip_slider = CLIPSliderFlux(pipe, device=torch.device("cuda"))

MAX_SEED = 2**32-1

def save_images_with_unique_filenames(image_list, save_directory):
    if not os.path.exists(save_directory):
        os.makedirs(save_directory)

    paths = []
    for image in image_list:
        unique_filename = f"{uuid.uuid4()}.png"
        file_path = os.path.join(save_directory, unique_filename)

        image.save(file_path)
        paths.append(file_path)

    return paths

def convert_to_centered_scale(num):
    if num % 2 == 0:  # even
        start = -(num // 2 - 1)
        end = num // 2
    else:  # odd
        start = -(num // 2)
        end = num // 2
    return tuple(range(start, end + 1))

def translate_if_korean(text):
    if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text):
        return translator(text)[0]['translation_text']
    return text

@spaces.GPU(duration=85)
def generate(prompt,
             concept_1,
             concept_2,
             scale,
             randomize_seed=True,
             seed=42,
             recalc_directions=True,
             iterations=200,
             steps=3,
             interm_steps=33,
             guidance_scale=3.5,
             x_concept_1="", x_concept_2="",
             avg_diff_x=None,
             total_images=[],
             gradio_progress=gr.Progress()
    ):
    # Translate prompt and concepts if Korean
    prompt = translate_if_korean(prompt)
    concept_1 = translate_if_korean(concept_1)
    concept_2 = translate_if_korean(concept_2)

    print(f"Prompt: {prompt}, ← {concept_2}, {concept_1} ➡️ . scale {scale}, interm steps {interm_steps}")
    slider_x = [concept_2, concept_1]
    # check if avg diff for directions need to be re-calculated
    if randomize_seed:
            seed = random.randint(0, MAX_SEED)

    if not sorted(slider_x) == sorted([x_concept_1, x_concept_2]) or recalc_directions:
        gradio_progress(0, desc="Calculating directions...")
        avg_diff = clip_slider.find_latent_direction(slider_x[0], slider_x[1], num_iterations=iterations)
        x_concept_1, x_concept_2 = slider_x[0], slider_x[1]

    images = []
    high_scale = scale
    low_scale = -1 * scale
    for i in gradio_progress.tqdm(range(interm_steps), desc="Generating images"):
        cur_scale = low_scale + (high_scale - low_scale) * i / (interm_steps - 1)
        image = clip_slider.generate(prompt,
                                     width=768,
                                     height=768,
                                     guidance_scale=guidance_scale,
                                     scale=cur_scale,  seed=seed, num_inference_steps=steps, avg_diff=avg_diff)
        images.append(image)
    canvas = Image.new('RGB', (256*interm_steps, 256))
    for i, im in enumerate(images):
        canvas.paste(im.resize((256,256)), (256 * i, 0))

    comma_concepts_x = f"{slider_x[1]}, {slider_x[0]}"

    scale_total = convert_to_centered_scale(interm_steps)
    scale_min = scale_total[0]
    scale_max = scale_total[-1]
    scale_middle = scale_total.index(0)
    post_generation_slider_update = gr.update(label=comma_concepts_x, value=0, minimum=scale_min, maximum=scale_max, interactive=True)
    avg_diff_x = avg_diff.cpu()

    video_path = f"{uuid.uuid4()}.mp4"
    print(video_path)
    return x_concept_1,x_concept_2, avg_diff_x, export_to_video(images, video_path, fps=5), canvas, images, images[scale_middle], post_generation_slider_update, seed

def update_pre_generated_images(slider_value, total_images):
    number_images = 0
    if total_images:
        number_images = len(total_images)
    if(number_images > 0):
        scale_tuple = convert_to_centered_scale(number_images)
        return total_images[scale_tuple.index(slider_value)][0]
    else:
        return None

def reset_recalc_directions():
    return True

# Updated examples with English text
examples = [
    ["flower in mountain", "spring", "winter", 1.5],
    ["a tomato", "super fresh", "rotten", 2.5],
    ["여자", "아기", "노인", 2.5]
]


css = """
footer {
    visibility: hidden;
}

.container {
    max-width: 1200px;
    margin: auto;
}

.main-panel {
    background-color: rgba(255, 255, 255, 0.05);
    border-radius: 12px;
    padding: 20px;
    margin-bottom: 20px;
}

.controls-panel {
    background-color: rgba(255, 255, 255, 0.02);
    border-radius: 8px;
    padding: 16px;
}

.image-display {
    min-height: 400px;
    display: flex;
    flex-direction: column;
    justify-content: center;
}

.slider-container {
    padding: 10px 0;
}

.advanced-panel {
    margin-top: 20px;
    border-top: 1px solid rgba(255, 255, 255, 0.1);
    padding-top: 20px;
}
"""

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
    x_concept_1 = gr.State("")
    x_concept_2 = gr.State("")
    total_images = gr.State([])
    avg_diff_x = gr.State()
    recalc_directions = gr.State(False)

    with gr.Row(elem_classes="container"):
        # Left Column - Controls
        with gr.Column(scale=4):
            with gr.Group(elem_classes="main-panel"):
                gr.Markdown("### Image Generation Controls")
                with gr.Group(elem_classes="controls-panel"):
                    prompt = gr.Textbox(
                        label=english_labels["Prompt"],
                        info="Enter the description",
                        placeholder="A dog in the park",
                        lines=2
                    )
                    with gr.Row():
                        with gr.Column(scale=1):
                            concept_1 = gr.Textbox(
                                label=english_labels["1st direction to steer"],
                                info="Initial state",
                                placeholder="winter"
                            )
                        with gr.Column(scale=1):
                            concept_2 = gr.Textbox(
                                label=english_labels["2nd direction to steer"],
                                info="Final state",
                                placeholder="summer"
                            )

                    with gr.Row(elem_classes="slider-container"):
                        x = gr.Slider(
                            minimum=0,
                            value=1.75,
                            step=0.1,
                            maximum=4.0,
                            label=english_labels["Strength"],
                            info="Maximum strength for each direction (above 2.5 may be unstable)"
                        )

                    submit = gr.Button(english_labels["Generate directions"], size="lg", variant="primary")

            # Advanced Options Panel
            with gr.Accordion(label=english_labels["Advanced options"], open=False, elem_classes="advanced-panel"):
                with gr.Row():
                    with gr.Column(scale=1):
                        interm_steps = gr.Slider(
                            label=english_labels["Num of intermediate images"],
                            minimum=3,
                            value=7,
                            maximum=65,
                            step=2
                        )
                    with gr.Column(scale=1):
                        guidance_scale = gr.Slider(
                            label=english_labels["Guidance scale"],
                            minimum=0.1,
                            maximum=10.0,
                            step=0.1,
                            value=3.5
                        )

                with gr.Row():
                    with gr.Column(scale=1):
                        iterations = gr.Slider(
                            label=english_labels["Num iterations for clip directions"],
                            minimum=0,
                            value=200,
                            maximum=400,
                            step=1
                        )
                    with gr.Column(scale=1):
                        steps = gr.Slider(
                            label=english_labels["Num inference steps"],
                            minimum=1,
                            value=3,
                            maximum=4,
                            step=1
                        )

                with gr.Row():
                    with gr.Column(scale=1):
                        randomize_seed = gr.Checkbox(
                            True,
                            label=english_labels["Randomize seed"]
                        )
                    with gr.Column(scale=1):
                        seed = gr.Slider(
                            minimum=0,
                            maximum=MAX_SEED,
                            step=1,
                            label=english_labels["Seed"],
                            interactive=True,
                            randomize=True
                        )

        # Right Column - Output
        with gr.Column(scale=6):
            with gr.Group(elem_classes="main-panel"):
                gr.Markdown("### Generated Results")
                with gr.Row():
                    with gr.Column(scale=3):
                        image_seq = gr.Image(
                            label=english_labels["Strip"],
                            elem_id="strip",
                            height=100
                        )
                    with gr.Column(scale=2):
                        output_image = gr.Video(
                            label=english_labels["Looping video"],
                            elem_id="video",
                            loop=True,
                            autoplay=True,
                            height=100
                        )
                with gr.Row(): # Moved this block to be after the video
                    with gr.Column():
                        post_generation_image = gr.Image(
                            label=english_labels["Generated Images"],
                            type="filepath",
                            elem_id="interactive",
                            elem_classes="image-display"
                        )
                        post_generation_slider = gr.Slider(
                            minimum=-10,
                            maximum=10,
                            value=0,
                            step=1,
                            label=english_labels["From 1st to 2nd direction"]
                        )


    # Examples Section
    gr.Examples(
        examples=examples,
        inputs=[prompt, concept_1, concept_2, x],
        fn=generate,
        outputs=[x_concept_1, x_concept_2, avg_diff_x, output_image, image_seq, total_images,
                post_generation_image, post_generation_slider, seed],
        cache_examples="lazy"
    )

    # Event Handlers
    submit.click(
        fn=generate,
        inputs=[prompt, concept_1, concept_2, x, randomize_seed, seed, recalc_directions,
                iterations, steps, interm_steps, guidance_scale, x_concept_1, x_concept_2,
                avg_diff_x, total_images],
        outputs=[x_concept_1, x_concept_2, avg_diff_x, output_image, image_seq, total_images,
                post_generation_image, post_generation_slider, seed]
    )

    iterations.change(fn=reset_recalc_directions, outputs=[recalc_directions])
    seed.change(fn=reset_recalc_directions, outputs=[recalc_directions])
    post_generation_slider.change(
        fn=update_pre_generated_images,
        inputs=[post_generation_slider, total_images],
        outputs=[post_generation_image],
        queue=False,
        show_progress="hidden",
        concurrency_limit=None
    )

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