app.py

import random
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import AutoencoderKL
from mixture_tiling_sdxl import StableDiffusionXLTilingPipeline

MAX_SEED = np.iinfo(np.int32).max
SCHEDULERS = [
    "LMSDiscreteScheduler",
    "DEISMultistepScheduler",
    "HeunDiscreteScheduler",
    "EulerAncestralDiscreteScheduler",
    "EulerDiscreteScheduler",
    "DPMSolverMultistepScheduler",
    "DPMSolverMultistepScheduler-Karras",
    "DPMSolverMultistepScheduler-Karras-SDE",
    "UniPCMultistepScheduler"
]

# 모델 로딩: VAE 및 타일링 파이프라인 초기화
vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
).to("cuda")

model_id = "stablediffusionapi/yamermix-v8-vae"
pipe = StableDiffusionXLTilingPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    vae=vae,
    use_safetensors=False,  # for yammermix
).to("cuda")

pipe.enable_model_cpu_offload()  # VRAM이 제한된 경우 사용
pipe.enable_vae_tiling()
pipe.enable_vae_slicing()

#region functions
def select_scheduler(scheduler_name):
    scheduler_parts = scheduler_name.split("-")
    scheduler_class_name = scheduler_parts[0]
    add_kwargs = {
        "beta_start": 0.00085,
        "beta_end": 0.012,
        "beta_schedule": "scaled_linear",
        "num_train_timesteps": 1000
    }
    if len(scheduler_parts) > 1:
        add_kwargs["use_karras_sigmas"] = True
    if len(scheduler_parts) > 2:
        add_kwargs["algorithm_type"] = "sde-dpmsolver++"
    import diffusers
    scheduler_cls = getattr(diffusers, scheduler_class_name)
    scheduler = scheduler_cls.from_config(pipe.scheduler.config, **add_kwargs)
    return scheduler

@spaces.GPU
def predict(left_prompt, center_prompt, right_prompt, negative_prompt, left_gs, center_gs, right_gs,
            overlap_pixels, steps, generation_seed, scheduler, tile_height, tile_width, target_height, target_width):
    global pipe
    print(f"Using scheduler: {scheduler}...")
    pipe.scheduler = select_scheduler(scheduler)
    generator = torch.Generator("cuda").manual_seed(generation_seed)
    
    target_height = int(target_height)
    target_width = int(target_width)
    tile_height = int(tile_height)
    tile_width = int(tile_width)
    
    image = pipe(
        prompt=[[left_prompt, center_prompt, right_prompt]],
        negative_prompt=negative_prompt,
        tile_height=tile_height,
        tile_width=tile_width,
        tile_row_overlap=0,
        tile_col_overlap=overlap_pixels,
        guidance_scale_tiles=[[left_gs, center_gs, right_gs]],
        height=target_height,
        width=target_width,
        generator=generator,
        num_inference_steps=steps,
    )["images"][0]
    return image

def calc_tile_size(target_height, target_width, overlap_pixels, max_tile_width_size=1280):
    num_cols = 3
    num_rows = 1    
    min_tile_dimension = 8
    reduction_step = 8
    max_tile_height_size = 1024
    best_tile_width = 0
    best_tile_height = 0
    best_adjusted_target_width = 0
    best_adjusted_target_height = 0
    found_valid_solution = False

    tile_width = max_tile_width_size
    tile_height = max_tile_height_size
    while tile_width >= min_tile_dimension:
        horizontal_borders = num_cols - 1
        total_horizontal_overlap = overlap_pixels * horizontal_borders
        adjusted_target_width = tile_width * num_cols - total_horizontal_overlap

        vertical_borders = num_rows - 1
        total_vertical_overlap = overlap_pixels * vertical_borders
        adjusted_target_height = tile_height * num_rows - total_vertical_overlap

        if tile_width <= max_tile_width_size and adjusted_target_width <= target_width:
            if adjusted_target_width > best_adjusted_target_width:
                best_tile_width = tile_width
                best_adjusted_target_width = adjusted_target_width
                found_valid_solution = True
        tile_width -= reduction_step

    if found_valid_solution:
        tile_width = best_tile_width
        tile_height = max_tile_height_size
        while tile_height >= min_tile_dimension:
            horizontal_borders = num_cols - 1
            total_horizontal_overlap = overlap_pixels * horizontal_borders
            adjusted_target_width = tile_width * num_cols - total_horizontal_overlap

            vertical_borders = num_rows - 1
            total_vertical_overlap = overlap_pixels * vertical_borders
            adjusted_target_height = tile_height * num_rows - total_vertical_overlap
        
            if tile_height <= max_tile_height_size and adjusted_target_height <= target_height:
                if adjusted_target_height > best_adjusted_target_height:
                    best_tile_height = tile_height
                    best_adjusted_target_height = adjusted_target_height
            tile_height -= reduction_step

    new_target_height = best_adjusted_target_height
    new_target_width = best_adjusted_target_width
    tile_width = best_tile_width
    tile_height = best_tile_height

    print("--- TILE SIZE CALCULATED VALUES ---")    
    print(f"Requested Overlap Pixels: {overlap_pixels}")
    print(f"Tile Height (max {max_tile_height_size}, divisible by 8): {tile_height}")
    print(f"Tile Width (max {max_tile_width_size}, divisible by 8): {tile_width}")
    print(f"Columns: {num_cols} | Rows: {num_rows}")
    print(f"Original Target: {target_height} x {target_width}")
    print(f"Adjusted Target: {new_target_height} x {new_target_width}\n")

    return new_target_height, new_target_width, tile_height, tile_width

def do_calc_tile(target_height, target_width, overlap_pixels, max_tile_size):    
    new_target_height, new_target_width, tile_height, tile_width = calc_tile_size(target_height, target_width, overlap_pixels, max_tile_size)    
    return gr.update(value=tile_height), gr.update(value=tile_width), gr.update(value=new_target_height), gr.update(value=new_target_width)

def clear_result():
    return gr.update(value=None)

def randomize_seed_fn(generation_seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        generation_seed = random.randint(0, MAX_SEED)
    return generation_seed
#endregion

# CSS 강화: 그라데이션 배경, 입체적 그림자, 부드러운 전환 효과 등으로 UI를 비주얼하게 개선
css = """
body {
    background: linear-gradient(135deg, #f5f7fa, #c3cfe2);
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}
.gradio-container {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 20px;
    padding: 30px;
    box-shadow: 0 20px 40px rgba(0, 0, 0, 0.2);
    border: 1px solid rgba(255, 255, 255, 0.3);
}
.gradio-container h1 {
    color: #333333;
    text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.2);
}
.fillable {
    width: 95% !important;
    max-width: unset !important;
}
#examples_container {
    margin: auto;
    width: 90%;
}
#examples_row {
    justify-content: center;
}

/* 스타일리시한 버튼 */
button {
    background: linear-gradient(145deg, #6a82fb, #fc5c7d);
    color: #ffffff;
    font-weight: bold;
    border: none;
    border-radius: 10px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.2);
    padding: 10px 20px;
    cursor: pointer;
    transition: transform 0.2s ease, box-shadow 0.2s ease;
}
button:hover {
    transform: translateY(-3px);
    box-shadow: 0 8px 12px rgba(0, 0, 0, 0.3);
}

/* 입력 필드 및 텍스트박스 포커스 효과 */
input[type="text"], textarea {
    border: 2px solid #ddd;
    border-radius: 10px;
    padding: 10px;
    transition: border-color 0.3s ease, box-shadow 0.3s ease;
}
input[type="text"]:focus, textarea:focus {
    border-color: #6a82fb;
    box-shadow: 0 0 10px rgba(106, 130, 251, 0.5);
}
"""

title = """
<h1 align="center" style="margin-bottom: 0.2em;">Panorama X3 IMAGE 🤗</h1>

"""

with gr.Blocks(css=css, title="SDXL Tiling Pipeline") as app:
    gr.Markdown(title)

    with gr.Row():
        # 좌/중앙/우 프롬프트 및 결과 영역
        with gr.Column(scale=7):
            generate_button = gr.Button("Generate", elem_id="generate_btn")
            with gr.Row():
                with gr.Column(variant="panel"):
                    gr.Markdown("### Left Region")
                    left_prompt = gr.Textbox(lines=4, placeholder="예: 울창한 숲과 햇살이 비추는 나무...", label="Left Prompt")
                    left_gs = gr.Slider(minimum=0, maximum=15, value=7, step=1, label="Left CFG Scale")
                with gr.Column(variant="panel"):
                    gr.Markdown("### Center Region")
                    center_prompt = gr.Textbox(lines=4, placeholder="예: 잔잔한 호수와 반짝이는 수면...", label="Center Prompt")
                    center_gs = gr.Slider(minimum=0, maximum=15, value=7, step=1, label="Center CFG Scale")
                with gr.Column(variant="panel"):
                    gr.Markdown("### Right Region")
                    right_prompt = gr.Textbox(lines=4, placeholder="예: 웅장한 산맥과 하늘을 가르는 구름...", label="Right Prompt")
                    right_gs = gr.Slider(minimum=0, maximum=15, value=7, step=1, label="Right CFG Scale")
            with gr.Row():
                negative_prompt = gr.Textbox(
                    lines=2,
                    label="Negative Prompt",
                    placeholder="예: blurry, low resolution, artifacts, poor details",
                    value="blurry, low resolution, artifacts, poor details"
                )
            with gr.Row():
                result = gr.Image(label="Generated Image", show_label=True, format="png", interactive=False, scale=1)

        # 사이드바: 파라미터 및 타일 크기 계산
        with gr.Sidebar(label="Parameters", open=True):
            gr.Markdown("### Generation Parameters")
            with gr.Row():
                height = gr.Slider(label="Target Height", value=1024, step=8, minimum=512, maximum=1024)
                width = gr.Slider(label="Target Width", value=1280, step=8, minimum=512, maximum=3840)
                overlap = gr.Slider(minimum=0, maximum=512, value=128, step=8, label="Tile Overlap")
                max_tile_size = gr.Dropdown(label="Max Tile Size", choices=[1024, 1280], value=1280)
                calc_tile = gr.Button("Calculate Tile Size")
            with gr.Row():
                tile_height = gr.Textbox(label="Tile Height", value=1024, interactive=False)
                tile_width = gr.Textbox(label="Tile Width", value=1024, interactive=False)
            with gr.Row():
                new_target_height = gr.Textbox(label="New Image Height", value=1024, interactive=False)
                new_target_width = gr.Textbox(label="New Image Width", value=1280, interactive=False)
            with gr.Row():
                steps = gr.Slider(minimum=1, maximum=50, value=30, step=1, label="Inference Steps")
                generation_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
            with gr.Row():
                scheduler = gr.Dropdown(label="Scheduler", choices=SCHEDULERS, value=SCHEDULERS[0])

    # 중앙에 배치된 예제 영역
    with gr.Row(elem_id="examples_row"):
        with gr.Column(scale=12, elem_id="examples_container"):
            gr.Markdown("### Example Prompts")
            gr.Examples(
                examples=[
                    [
                        "Iron Man, repulsor rays blasting enemies in destroyed cityscape, sparks, energy trails, crumbling skyscrapers, smoke, debris, cinematic lighting, photorealistic, intense action. Focus: Iron Man.",
                        "Captain America charging forward, vibranium shield deflecting energy blasts in destroyed cityscape, collapsing buildings, rubble streets, battle-damaged suit, determined expression, distant explosions, cinematic composition, realistic rendering. Focus: Captain America.",
                        "Thor wielding Stormbreaker in destroyed cityscape, lightning crackling, powerful strike downwards, shattered buildings, burning debris, ground trembling, Asgardian armor, cinematic photography, realistic details. Focus: Thor.",
                       "blurry, low resolution, artifacts, poor details",
                        7, 7, 7,
                        128,
                        30,
                        123456789,
                        "UniPCMultistepScheduler",
                        1024, 1280,
                        1024, 1920,
                        1280
                    ],
                    [
                        "A charming house in the countryside, by jakub rozalski, sunset lighting, elegant, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",
                        "A dirt road in the countryside crossing pastures, by jakub rozalski, sunset lighting, elegant, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",
                        "An old and rusty giant robot lying on a dirt road, by jakub rozalski, dark sunset lighting, elegant, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",
                       "blurry, poorly rendered, low quality, disfigured",
                        8, 8, 8,
                        100,
                        35,
                        987654321,
                        "EulerDiscreteScheduler",
                        1024, 1280,
                        1024, 1920,
                        1280
                    ],
                    [
                        "Abstract decorative illustration, by joan miro and gustav klimt and marlina vera and loish, elegant, intricate, highly detailed, smooth, sharp focus, vibrant colors, artstation, stunning masterpiece",
                        "Abstract decorative illustration, by joan miro and gustav klimt and marlina vera and loish, elegant, intricate, highly detailed, smooth, sharp focus, vibrant colors, artstation, stunning masterpiece",
                        "Abstract decorative illustration, by joan miro and gustav klimt and marlina vera and loish, elegant, intricate, highly detailed, smooth, sharp focus, vibrant colors, artstation, stunning masterpiece",
                       "text, watermark, signature, distorted",
                        6, 6, 6,
                        80,
                        25,
                        192837465,
                        "DPMSolverMultistepScheduler-Karras",
                        1024, 1280,
                        1024, 1920,
                        1280
                    ],
                    [
                        "Magical diagrams and runes written with chalk on a blackboard, elegant, intricate, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",
                        "Magical diagrams and runes written with chalk on a blackboard, elegant, intricate, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",
                        "Magical diagrams and runes written with chalk on a blackboard, elegant, intricate, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",
                         "low quality, artifact, deformed, sketchy",
                        9, 9, 9,
                        150,
                        40,
                        1029384756,
                        "DPMSolverMultistepScheduler-Karras-SDE",
                        1024, 1280,
                        1024, 1920,
                        1280
                    ]
                ],
                inputs=[left_prompt, center_prompt, right_prompt, negative_prompt, 
                        left_gs, center_gs, right_gs, overlap, steps, generation_seed, 
                        scheduler, tile_height, tile_width, height, width, max_tile_size],
                cache_examples=False
            )

    # 이벤트 연결: 타일 사이즈 계산 및 이미지 생성
    event_calc_tile_size = {
        "fn": do_calc_tile,
        "inputs": [height, width, overlap, max_tile_size],
        "outputs": [tile_height, tile_width, new_target_height, new_target_width]
    }
    calc_tile.click(**event_calc_tile_size)
    
    generate_button.click(
        fn=clear_result,
        inputs=None,
        outputs=result,
    ).then(**event_calc_tile_size).then(
        fn=randomize_seed_fn,
        inputs=[generation_seed, randomize_seed],
        outputs=generation_seed,
        queue=False,
        api_name=False,
    ).then(
        fn=predict,
        inputs=[left_prompt, center_prompt, right_prompt, negative_prompt, 
                left_gs, center_gs, right_gs, overlap, steps, generation_seed, 
                scheduler, tile_height, tile_width, new_target_height, new_target_width],
        outputs=result,
    )

app.launch(share=False)