app.py

import spaces
import gradio as gr
import os
import random
import re
import numpy as np
import torch
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer

# Import the custom tiling pipeline and its components
from pipeline_z_image_mod import ZImageMoDTilingPipeline
from diffusers import (
    AutoencoderKL,
    FlowMatchEulerDiscreteScheduler,
    GGUFQuantizationConfig,
)
from diffusers.models import ZImageTransformer2DModel
from huggingface_hub import hf_hub_download

# Import the safety checker function
try:
    from prompt_check import is_unsafe_prompt

    UNSAFE_CHECK_AVAILABLE = True
except ImportError:
    print("Warning: 'prompt_check.py' not found. NSFW prompt check will be disabled.")
    UNSAFE_CHECK_AVAILABLE = False

    def is_unsafe_prompt(*args, **kwargs):
        return False


# 1. Environment Variables for Model Paths
# Base model for VAE, text_encoder, tokenizer
BASE_MODEL_ID = os.getenv("BASE_MODEL_ID", "Tongyi-MAI/Z-Image-Turbo")
# GGUF model for the transformer
GGUF_REPO_ID = os.getenv("GGUF_REPO_ID", "jayn7/Z-Image-Turbo-GGUF")
GGUF_FILENAME = os.getenv("GGUF_FILENAME", "z_image_turbo-Q4_K_M.gguf")
GGUF_LOCAL_DIR = os.getenv(
    "GGUF_LOCAL_DIR", None
)  # Set this path for local use, e.g., "F:\\models\\Z-Image-Turbo"
USE_SPACES_ENV = os.getenv("USE_SPACES", "false").lower()
USE_SPACES = USE_SPACES_ENV not in ("false", "0", "no", "none")

# System prompt for the safety checker
UNSAFE_PROMPT_CHECK = os.getenv("UNSAFE_PROMPT_CHECK")
MAX_SEED = np.iinfo(np.int32).max

# 2. Load Models (GGUF + Standard Components)
print("--- Loading Models ---")
device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading VAE...")
if USE_SPACES:
    vae = AutoencoderKL.from_pretrained(
        BASE_MODEL_ID, subfolder="vae", torch_dtype=torch.bfloat16
    ).to(device)
else:
    vae = AutoencoderKL.from_pretrained(
        BASE_MODEL_ID, subfolder="vae", torch_dtype=torch.bfloat16
    )

print("Loading Text Encoder and Tokenizer...")
if USE_SPACES:
    text_encoder = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID, subfolder="text_encoder", dtype=torch.bfloat16
    ).to(device)
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, subfolder="tokenizer")
else:
    text_encoder = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID, subfolder="text_encoder", dtype=torch.bfloat16
    )
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, subfolder="tokenizer")

print(f"Loading Transformer from GGUF: {GGUF_REPO_ID}/{GGUF_FILENAME}...")
transformer_path = (
    hf_hub_download(
        GGUF_REPO_ID,
        GGUF_FILENAME,
        local_dir=GGUF_LOCAL_DIR,
        local_dir_use_symlinks=False,
    )
    if GGUF_LOCAL_DIR
    else hf_hub_download(GGUF_REPO_ID, GGUF_FILENAME)
)

if USE_SPACES:
    transformer = ZImageTransformer2DModel.from_single_file(
        transformer_path,
        quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
        dtype=torch.bfloat16,
    ).to(device)
else:
    transformer = ZImageTransformer2DModel.from_single_file(
        transformer_path,
        quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
        dtype=torch.bfloat16,
    )


# 3. Assemble the Tiling Pipeline
print("\nAssembling the ZImageMoDTilingPipeline...")
scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0)
if USE_SPACES:
    pipe = ZImageMoDTilingPipeline(
        vae=vae,
        text_encoder=text_encoder,
        tokenizer=tokenizer,  
        scheduler=scheduler,      
        transformer=transformer,
    ).to(device)
else:
    pipe = ZImageMoDTilingPipeline(
        vae=vae,
        text_encoder=text_encoder,
        tokenizer=tokenizer,        
        scheduler=scheduler,
        transformer=transformer,
    )

    print("Enabling model CPU offload...")
    pipe.enable_model_cpu_offload()

# Load Translation Models
print("Loading translation models...")
try:
    ko_en_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
    zh_en_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en")
except Exception as e:
    ko_en_translator, zh_en_translator = None, None
    print(f"Warning: Could not load translation models: {e}")

print("Pipeline loaded and ready.")


# Helper Functions
def translate_prompt(text: str, language: str) -> str:
    """Translates text to English if the selected language is not English."""
    if language == "English" or not text.strip():
        return text
    translated = text
    if (
        language == "Korean"
        and ko_en_translator
        and any("\uac00" <= char <= "\ud7a3" for char in text)
    ):
        translated = ko_en_translator(text)[0]["translation_text"]
    elif (
        language == "Chinese"
        and zh_en_translator
        and any("\u4e00" <= char <= "\u9fff" for char in text)
    ):
        translated = zh_en_translator(text)[0]["translation_text"]
    return translated


def create_hdr_effect(image, hdr_strength):
    if hdr_strength == 0:
        return image
    from PIL import ImageEnhance, Image

    if isinstance(image, Image.Image):
        image = np.array(image)
    from scipy.ndimage import gaussian_filter

    blurred = gaussian_filter(image, sigma=5)
    sharpened = np.clip(image + hdr_strength * (image - blurred), 0, 255).astype(
        np.uint8
    )
    pil_img = Image.fromarray(sharpened)
    converter = ImageEnhance.Color(pil_img)
    return converter.enhance(1 + hdr_strength)


@spaces.GPU(duration=120)
def generate_z_image_panorama(
    left_prompt,
    center_prompt,
    right_prompt,
    left_gs,
    center_gs,
    right_gs,
    overlap_pixels,
    steps,
    shift,
    generation_seed,
    tile_weighting_method,
    prompt_language,
    target_height,
    target_width,
    hdr,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generate a panoramic image using the Z-Image Turbo model with tiling and composition.
    Args:
        left_prompt (str): Text prompt for the left section of the panorama.
        center_prompt (str): Text prompt for the center section of the panorama.
        right_prompt (str): Text prompt for the right section of the panorama.
        left_gs (float): Guidance scale for the left tile.
        center_gs (float): Guidance scale for the center tile.
        right_gs (float): Guidance scale for the right tile.
        overlap_pixels (int): Number of pixels to overlap between tiles.
        steps (int): Number of inference steps for generation.
        shift (float): Time Shift.
        generation_seed (int): Random seed for reproducibility.
        tile_weighting_method (str): Method for weighting overlapping tile regions.
        prompt_language (str): Language code for prompt translation.
        target_height (int): Height of the generated panorama in pixels.
        target_width (int): Width of the generated panorama in pixels.
        hdr (float): HDR effect intensity.
        randomize_seed (boolean): Not used.
        progress (gr.Progress): Gradio progress tracker.
    Returns:
        PIL.Image: The generated panoramic image with optional HDR effect applied.
    """
    if not left_prompt or not center_prompt or not right_prompt:
        gr.Info("⚡️ Prompts must be provided!")
        return gr.skip()
    # Safety Check    
    prompts_to_check = [left_prompt, center_prompt, right_prompt]
    for p in prompts_to_check:
        if UNSAFE_CHECK_AVAILABLE and is_unsafe_prompt(
            pipe.text_encoder,
            device,
            pipe.tokenizer,
            system_prompt=UNSAFE_PROMPT_CHECK,
            user_prompt=p,
        ):
            raise gr.Error(
                f"Unsafe prompt detected. Please modify your prompt and try again."
            )

    generator = torch.Generator("cuda").manual_seed(generation_seed)
    pipe.scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift)
    
    final_height, final_width = int(target_height), int(target_width)

    translated_left = translate_prompt(left_prompt, prompt_language)
    translated_center = translate_prompt(center_prompt, prompt_language)
    translated_right = translate_prompt(right_prompt, prompt_language)

    image = pipe(
        prompt=[[translated_left, translated_center, translated_right]],
        height=final_height,
        width=final_width,
        num_inference_steps=steps,
        guidance_scale_tiles=[[left_gs, center_gs, right_gs]],
        tile_overlap=overlap_pixels,
        tile_weighting_method=tile_weighting_method,
        generator=generator,
    ).images[0]

    return create_hdr_effect(image, hdr)


def calculate_tile_size(target_height, target_width, overlap_pixels):
    """
    Calculate tile dimensions for panoramic image generation.

    Args:
        target_height (int): The target height of the final panoramic image in pixels.
        target_width (int): The target width of the final panoramic image in pixels.
        overlap_pixels (int): The number of overlapping pixels between adjacent tiles.

    Returns:
        tuple: A tuple of 2 gr.update objects containing:
            - final_height: Final panorama height after tiling
            - final_width: Final panorama width after tiling
    """
    num_cols = 3
    num_rows = 1
    tile_width = (target_width + (num_cols - 1) * overlap_pixels) // num_cols
    tile_height = (target_height + (num_rows - 1) * overlap_pixels) // num_rows
    tile_width -= tile_width % 16
    tile_height -= tile_height % 16
    final_width = tile_width * num_cols - (num_cols - 1) * overlap_pixels
    final_height = tile_height * num_rows - (num_rows - 1) * overlap_pixels
    return (
        gr.update(value=final_height),
        gr.update(value=final_width),
    )


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


def run_for_examples(
    left_prompt,
    center_prompt,
    right_prompt,
    left_gs,
    center_gs,
    right_gs,
    overlap_pixels,
    steps,
    shift,
    generation_seed,
    tile_weighting_method,
    prompt_language,
    target_height,
    target_width,
    hdr,
    randomize_seed
):
    return generate_z_image_panorama(
        left_prompt,
        center_prompt,
        right_prompt,
        left_gs,
        center_gs,
        right_gs,
        overlap_pixels,
        steps,
        shift,
        generation_seed,
        tile_weighting_method,
        prompt_language,
        target_height,
        target_width,
        hdr,
        randomize_seed
    )


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


def update_dimensions_from_preset(resolution_preset):
    """Updates the width and height sliders based on a preset string."""
    match = re.search(r"(\d+)\s*[x]\s*(\d+)", resolution_preset)
    if match:
        width, height = int(match.group(1)), int(match.group(2))
        return gr.update(value=width), gr.update(value=height)
    return gr.update(), gr.update()


# UI Layout
theme = gr.themes.Default(
    primary_hue="red", secondary_hue="orange", neutral_hue="gray"
).set(
    body_background_fill="*neutral_100",
    body_background_fill_dark="*neutral_900",
    body_text_color="*neutral_900",
    body_text_color_dark="*neutral_100",
    panel_background_fill="*neutral_800",
    panel_background_fill_dark="*neutral_900",
    input_background_fill="white",
    input_background_fill_dark="*neutral_800",
    button_primary_background_fill="*primary_500",
    button_primary_background_fill_dark="*primary_700",
    button_primary_text_color="white",
    button_primary_text_color_dark="white",
    button_secondary_background_fill="*secondary_500",
    button_secondary_background_fill_dark="*secondary_700",
    button_secondary_text_color="white",
    button_secondary_text_color_dark="white",
)

css_code = ""

try:
    with open("./style.css", "r", encoding="utf-8") as f:
        css_code += f.read() + "\n"
except FileNotFoundError:
    pass

title = """<h1 align="center">Z Image Turbo - Panorama 🏞️✨</h1>
           <div style="text-align: center;">
                <span>Create panoramic compositions with Z-Image-Turbo.</span>
           </div>
           """
# Adapted panoramic resolutions
PANORAMIC_RESOLUTIONS = [
    "2048x1024 (2:1)",
    "2560x1024 (2.5:1)",
    "3072x1024 (3:1)",
    "2304x896 ( ~2.5:1)",
    "2880x1152 (~2.5:1)",
    "3840x1536 (2.5:1)",
]

with gr.Blocks(theme=theme, css=css_code, title="Panorama Z-Image") as app:
    gr.Markdown(title)
    with gr.Row():
        with gr.Column(scale=7):
            generate_button = gr.Button("Generate", variant="primary")
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Left Region")
                    left_prompt = gr.Textbox(lines=4, label="Prompt (Left)")
                    left_gs = gr.Slider(
                        minimum=0.0,
                        maximum=10.0,
                        value=0.0,
                        step=0.1,
                        label="Left Guidance",
                        visible=False,
                    )
                with gr.Column(scale=1):
                    gr.Markdown("### Center Region")
                    center_prompt = gr.Textbox(lines=4, label="Prompt (Center)")
                    center_gs = gr.Slider(
                        minimum=0.0,
                        maximum=10.0,
                        value=0.0,
                        step=0.1,
                        label="Center Guidance",
                        visible=False,
                    )
                with gr.Column(scale=1):
                    gr.Markdown("### Right Region")
                    right_prompt = gr.Textbox(lines=4, label="Prompt (Right)")
                    right_gs = gr.Slider(
                        minimum=0.0,
                        maximum=10.0,
                        value=0.0,
                        step=0.1,
                        label="Right Guidance",
                        visible=False,
                    )
            with gr.Row():
                result = gr.Image(
                    label="Generated Image",
                    show_label=True,
                    format="png",
                    interactive=False,
                )
        with gr.Column(scale=1, min_width=200):
            gr.Markdown(
                """
                ### Related Apps
                Check out these other great demos:
             
                - **[Official Z-Image Demo](https://huggingface.co/spaces/Tongyi-MAI/Z-Image-Turbo)**
                    *The official demo for single image generation.*
                
                - **[Panorama FLUX](https://huggingface.co/spaces/elismasilva/flux-1-panorama)**
                    *The Panorama app using the FLUX model.*
                """
            )
    with gr.Sidebar():
        gr.Markdown("### Generation Parameters")
        prompt_language = gr.Radio(
            choices=["English", "Korean", "Chinese"],
            value="English",
            label="Prompt Language",
        )

        resolution_preset = gr.Dropdown(
            choices=PANORAMIC_RESOLUTIONS,
            value=PANORAMIC_RESOLUTIONS[2],
            label="Panoramic Resolution Preset",
        )

        with gr.Row():
            height = gr.Slider(
                label="Target Height", value=1024, step=16, minimum=512, maximum=2048
            )
            width = gr.Slider(
                label="Target Width", value=3072, step=16, minimum=512, maximum=4096
            )
        with gr.Row():
            overlap = gr.Slider(
                minimum=0, maximum=512, value=256, step=16, label="Tile Overlap"
            )
            tile_weighting_method = gr.Dropdown(
                label="Blending Method", choices=["Cosine", "Gaussian"], value="Cosine"
            )
        with gr.Row():
            calc_tile = gr.Button("Calculate Final Dimensions", variant="primary")
        with gr.Row():
            new_target_height = gr.Textbox(
                label="Actual Height", value=1024, interactive=False
            )
            new_target_width = gr.Textbox(
                label="Actual Width", value=3072, interactive=False
            )
        with gr.Row():
            steps = gr.Slider(
                minimum=1, maximum=50, value=9, step=1, label="Inference Steps"
            )
            shift = gr.Slider(label="Time Shift", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
        with gr.Row():
            generation_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():
            hdr = gr.Slider(
                minimum=0.0, maximum=1.0, value=0.1, step=0.00, label="HDR Effect"
            )

    with gr.Row():
        gr.Examples(
            examples=[                
                [
                    "On the left side of a wide, wet stone courtyard at night, a traditional wooden temple building has a soft, warm glow from its illuminated windows, the environment is a deep, dark night sky.",
                    "Young Chinese woman in red Hanfu, standing gracefully in the center of the stone-paved temple courtyard, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights.",
                    "On the right side of the frame, a wooden temple corridor with a row of glowing paper lanterns recedes into the background. The lantern light reflects on the wet stone path below, and the distant pagoda is silhouetted against the same deep, dark night sky.",    
                    0.0,
                    0.0,
                    0.0,
                    256,
                    8,
                    3.0,
                    109320357,
                    "Cosine",
                    "English",
                    1024,
                    3072,
                    0.0,
                    False,
                ],
                [
                    "Street photography of a bustling, modern Chinese city square in the morning. On the left, people in business attire walk past gleaming glass and steel office buildings under a hazy, bright sun.",
                    "The center of the vast, modern plaza, where a large, elegant water fountain sprays water high into the air. In the background, through the mist of the fountain, the same modern glass and steel office buildings continue across the scene. The sun reflects off the wet plaza floor.",
                    "The right side of the plaza, where groups of fashionable teenagers are casually standing near the concrete pillars of a modern building, looking at their smartphones under the same hazy, bright sun.",
                    0.0,
                    0.0,
                    0.0,
                    256,
                    8,
                    6.0,
                    2024,
                    "Cosine",
                    "English",
                    1024,
                    3072,
                    0.0,
                    False
                ],
                [
                    "The edge of the green carpeted arena floor, with the blurred figures of spectators in the stands rising up in the background. A judge's table is partially visible on the far left.",
                    "In the center of the green carpeted arena floor, a proud man in a fine suit stands with his perfectly groomed white poodle, both looking forward. The poodle wears a matching tiny tuxedo.",
                    "Further down the green carpeted arena on the right, a line of other dog handlers and their breeds are waiting their turn in the distance, their forms slightly blurred by the depth of field.",
                    0.0,
                    0.0,
                    0.0,
                    200,
                    8,
                    3.0,
                    1146985307,
                    "Gaussian",
                    "English",
                    1024,
                    2672,
                    0.0,
                    False
                ],
                [
                    "Iron Man, repulsor rays blasting enemies in destroyed cityscape, cinematic lighting, photorealistic. Focus: Iron Man.",
                    "Captain America charging forward, vibranium shield deflecting energy blasts in destroyed cityscape, cinematic composition. Focus: Captain America.",
                    "Thor wielding Stormbreaker in destroyed cityscape, lightning crackling, powerful strike downwards, cinematic photography. Focus: Thor.",
                    0,
                    0,
                    0,
                    200,
                    8,
                    4.0,
                    85512926,
                    "Cosine",
                    "English",
                    1024,
                    3440,
                    0.0,
                    False
                ]
            ],
            inputs=[
                left_prompt,
                center_prompt,
                right_prompt,
                left_gs,
                center_gs,
                right_gs,
                overlap,
                steps,
                shift,
                generation_seed,
                tile_weighting_method,
                prompt_language,
                height,
                width,
                hdr,
                randomize_seed
            ],
            fn=run_for_examples,
            outputs=result,
            cache_examples=False,
            api_name=False,
        )

    # Event Handling
    predict_inputs = [
        left_prompt,
        center_prompt,
        right_prompt,
        left_gs,
        center_gs,
        right_gs,
        overlap,
        steps,
        shift,
        generation_seed,
        tile_weighting_method,        
        prompt_language,
        new_target_height,
        new_target_width,
        hdr,
        randomize_seed
    ]

    resolution_preset.change(
        fn=update_dimensions_from_preset,
        inputs=resolution_preset,
        outputs=[width, height],
        api_name=False,
    )
    calc_tile.click(
        fn=calculate_tile_size,
        inputs=[height, width, overlap],
        outputs=[new_target_height, new_target_width],
        api_name="calculate_tile_size"
    )
    generate_button.click(
        fn=clear_result,
        inputs=None,
        outputs=result,
        queue=False,
        api_name=False
    ).then(
        fn=calculate_tile_size,
        inputs=[height, width, overlap],
        outputs=[new_target_height, new_target_width],
        show_api=False
    ).then(
        fn=randomize_seed_fn,
        inputs=[generation_seed, randomize_seed],
        outputs=generation_seed,
        queue=False,
        api_name=False,    
    ).then(
        fn=generate_z_image_panorama,
        inputs=predict_inputs,
        outputs=result,
        show_progress="full",
        api_name="generate_z_image_panorama",
    )

app.queue().launch(mcp_server=True, share=True)