app.py

import gradio as gr
import numpy as np
import torch
import spaces
from diffusers import FluxPipeline, FluxTransformer2DModel
from diffusers.utils import export_to_gif
from huggingface_hub import hf_hub_download
from PIL import Image
import uuid
import random

device = "cuda" if torch.cuda.is_available() else "cpu"

if torch.cuda.is_available():
    torch_dtype = torch.bfloat16
else:
    torch_dtype = torch.float32

def split_image(input_image, num_splits=4):
    output_images = []
    for i in range(num_splits):
        left = i * 320
        right = (i + 1) * 320
        box = (left, 0, right, 320)
        output_images.append(input_image.crop(box))
    return output_images

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch_dtype).to(device)

MAX_SEED = np.iinfo(np.int32).max

@spaces.GPU 
def infer(prompt, seed=1, randomize_seed=False, num_inference_steps=28):
    print('entered the function')
    prompt_template = f"A side by side 4 frame image showing high quality consecutive stills from a looped gif animation moving from left to right. The scene has motion. The stills are of {prompt}"
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
        
    generator = torch.Generator().manual_seed(seed)
    
    image = pipe(
        prompt=prompt_template,
        num_inference_steps=num_inference_steps,
        num_images_per_prompt=1,
        generator=generator,
        height=320,
        width=1280
    ).images[0]
    
    gif_name = f"{uuid.uuid4().hex}-flux.gif"
    export_to_gif(split_image(image, 4), gif_name, fps=4)
    
    return gif_name, image, seed

examples = [
    "a red panda playing with a bamboo stick in the snow",
    "an astronaut breakdancing on the moon",
    "a magical butterfly transforming into sparkles",
    "a robot learning to paint like Van Gogh",
    "a dragon hatching from a crystal egg",
    "a time traveler stepping through a portal",
    "a mermaid playing with bioluminescent fish",
    "a steampunk clock with moving gears",
    "a flower blooming in timelapse",
    "a wizard casting a colorful spell"
]

css = """
footer {visibility: hidden}

#col-container {
    max-width: 1200px;
    margin: auto;
    padding: 20px;
    background-color: #ffffff;
    border-radius: 20px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}

.gr-button {
    background: linear-gradient(45deg, #FF6B6B, #4ECDC4);
    border: none;
    color: white;
    font-weight: 600;
    border-radius: 8px;
}

.gr-button:hover {
    background: linear-gradient(45deg, #4ECDC4, #FF6B6B);
    transform: translateY(-2px);
    transition: all 0.3s ease;
}

.gr-input {
    border-radius: 8px;
    border: 2px solid #e0e0e0;
}

.gr-accordion {
    border-radius: 8px;
    background-color: #f8f9fa;
}

/* Examples 텍스트 색상 관련 스타일 수정 */
.gr-examples-text {
    color: black !important;
}

.gr-examples button {
    color: black !important;
}

.gr-examples span {
    color: black !important;
}

.gr-examples div {
    color: black !important;
}

.gr-examples p {
    color: black !important;
}

.gr-examples h3 {
    color: black !important;
}

.gr-sample-text {
    color: black !important;
}

#main-output {
    height: 500px !important;
    width: 100% !important;
}

#preview-output {
    height: 200px !important;
    width: auto !important;
}

#strip-output {
    height: 150px !important;
    width: auto !important;
}

/* 추가적인 Examples 관련 스타일 */
.example-text {
    color: black !important;
}

.example-label {
    color: black !important;
}
"""

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
    gr.HTML("""
        <div style="text-align: center; max-width: 800px; margin: 0 auto;">
            <h1 style="font-size: 3rem; font-weight: 700; margin-bottom: 1rem;">
                FLUX Animation Creator
            </h1>
            <p style="font-size: 1.2rem; color: #666; margin-bottom: 2rem;">
                Create amazing animated GIFs with AI - Just describe what you want to see!
            </p>
        </div>
    """)
    
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            prompt = gr.Text(
                label="Your Animation Prompt",
                show_label=True,
                max_lines=1,
                placeholder="Describe the animation you want to create...",
                container=True,
                elem_id="prompt-input"
            )
            run_button = gr.Button("✨ Generate", scale=0, variant="primary")
        
        # GIF 결과 출력 (큰 크기)
        result = gr.Image(
            label="Generated Animation", 
            show_label=True,
            elem_id="main-output",
            height=500
        )
        
        with gr.Row():
            # 미리보기 이미지들 (작은 크기)
            result_full = gr.Image(
                label="Preview", 
                elem_id="preview-output",
                height=200
            )
            strip_image = gr.Image(
                label="Animation Strip", 
                elem_id="strip-output",
                height=150
            )
        
        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)
            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=1,
                maximum=32,
                step=1,
                value=28, 
            )    
        
        gr.Examples(
            examples=examples,
            inputs=[prompt],
            outputs=[result, result_full, seed],
            fn=infer,
            cache_examples=True,
            label="Click on any example to try it out"
        )
    
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, seed, randomize_seed, num_inference_steps],
        outputs=[result, result_full, seed]
    )

demo.queue().launch()