import spaces
import gradio as gr
import os
import random
import numpy as np
import torch
from transformers import pipeline
# Import the pipeline
from flux_pipeline_mod import FluxMoDTilingPipeline
# 1. Load Translation Models ---
# These models are small and run efficiently on CPU.
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")
print("Translation models loaded successfully.")
except Exception as e:
print(f"Could not load translation models: {e}")
ko_en_translator = None
zh_en_translator = None
# 2. Conditional MMGP Setup ---
USE_MMGP_ENV = os.getenv("USE_MMGP", "true").lower()
USE_MMGP = USE_MMGP_ENV not in ("false", "0", "no", "none")
offload = None
if USE_MMGP:
print("INFO: Attempting to use MMGP.")
try:
from mmgp import offload, profile_type
print("Successfully imported MMGP.")
except ImportError:
print("WARNING: MMGP import failed. Falling back to standard offload.")
USE_MMGP = False
else:
print("INFO: MMGP is disabled.")
MAX_SEED = np.iinfo(np.int32).max
# 3. Load the Main Pipeline ---
print("Loading the FLUX Tiling pipeline...")
# Use an environment variable for the model path to make it flexible
MODEL_PATH = os.getenv("MODEL_PATH", "black-forest-labs/FLUX.1-schnell")
print(f"Loading model from: {MODEL_PATH}")
pipe = FluxMoDTilingPipeline.from_pretrained(MODEL_PATH, torch_dtype=torch.bfloat16).to(
"cuda"
)
if USE_MMGP and offload:
print("Applying LowRAM_LowVRAM offload profile via MMGP...")
offload.profile(pipe, profile_type.LowRAM_LowVRAM)
else:
print("Attempting to use the standard Diffusers CPU offload...")
try:
pipe.enable_model_cpu_offload()
except Exception as e:
print(f"Could not apply standard offload: {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 = text
if language == "Korean" and ko_en_translator:
if any(
"\uac00" <= char <= "\ud7a3" for char in text
): # Check if Korean characters are present
print(f"Translating Korean to English: '{text}'")
translated_text = ko_en_translator(text)[0]["translation_text"]
print(f" -> Translated: '{translated_text}'")
elif language == "Chinese" and zh_en_translator:
if any(
"\u4e00" <= char <= "\u9fff" for char in text
): # Check if Chinese characters are present
print(f"Translating Chinese to English: '{text}'")
translated_text = zh_en_translator(text)[0]["translation_text"]
print(f" -> Translated: '{translated_text}'")
return translated_text
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_flux_panorama(
left_prompt,
center_prompt,
right_prompt,
left_gs,
center_gs,
right_gs,
overlap_pixels,
steps,
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 FLUX 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.
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()
global pipe
generator = torch.Generator("cuda").manual_seed(generation_seed)
final_height, final_width = int(target_height), int(target_width)
# Translate prompts if necessary
translated_left = translate_prompt(left_prompt, prompt_language)
translated_center = translate_prompt(center_prompt, prompt_language)
translated_right = translate_prompt(right_prompt, prompt_language)
print("Starting generation with Tiling Pipeline (Composition Mode)...")
image = pipe(
prompt=[[translated_left, translated_center, translated_right]],
height=final_height,
width=final_width,
tile_overlap=overlap_pixels,
guidance_scale_tiles=[[left_gs, center_gs, right_gs]],
tile_weighting_method=tile_weighting_method,
generator=generator,
num_inference_steps=steps,
max_sequence_length=512,
).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 4 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 clear_result():
return gr.update(value=None)
def run_for_examples(
left_prompt,
center_prompt,
right_prompt,
left_gs,
center_gs,
right_gs,
overlap_pixels,
steps,
generation_seed,
tile_weighting_method,
prompt_language,
target_height,
target_width,
hdr,
randomize_seed
):
return generate_flux_panorama(
left_prompt,
center_prompt,
right_prompt,
left_gs,
center_gs,
right_gs,
overlap_pixels,
steps,
generation_seed,
tile_weighting_method,
prompt_language,
target_height,
target_width,
hdr,
randomize_seed
)
def randomize_seed_fn(generation_seed: int, randomize_seed: bool) -> int:
if randomize_seed:
generation_seed = random.randint(0, MAX_SEED)
return generation_seed
# UI Layout
theme = gr.themes.Default(
primary_hue="blue", secondary_hue="teal", neutral_hue="neutral"
).set(
body_background_fill="*neutral_100",
body_background_fill_dark="*neutral_900",
body_text_color="*neutral_700",
body_text_color_dark="*neutral_200",
body_text_weight="400",
link_text_color="*primary_500",
link_text_color_dark="*primary_400",
code_background_fill="*neutral_100",
code_background_fill_dark="*neutral_800",
shadow_drop="0 1px 3px rgba(0,0,0,0.1)",
shadow_inset="inset 0 2px 4px rgba(0,0,0,0.05)",
block_background_fill="*neutral_50",
block_background_fill_dark="*neutral_700",
block_border_color="*neutral_200",
block_border_color_dark="*neutral_600",
block_border_width="1px",
block_border_width_dark="1px",
block_label_background_fill="*primary_50",
block_label_background_fill_dark="*primary_600",
block_label_text_color="*primary_600",
block_label_text_color_dark="*primary_50",
panel_background_fill="white",
panel_background_fill_dark="*neutral_800",
panel_border_color="*neutral_200",
panel_border_color_dark="*neutral_700",
panel_border_width="1px",
panel_border_width_dark="1px",
input_background_fill="white",
input_background_fill_dark="*neutral_800",
input_border_color="*neutral_300",
input_border_color_dark="*neutral_700",
slider_color="*primary_500",
slider_color_dark="*primary_400",
button_primary_background_fill="*primary_600",
button_primary_background_fill_dark="*primary_500",
button_primary_background_fill_hover="*primary_700",
button_primary_background_fill_hover_dark="*primary_400",
button_primary_border_color="transparent",
button_primary_border_color_dark="transparent",
button_primary_text_color="white",
button_primary_text_color_dark="white",
button_secondary_background_fill="*neutral_200",
button_secondary_background_fill_dark="*neutral_600",
button_secondary_background_fill_hover="*neutral_300",
button_secondary_background_fill_hover_dark="*neutral_500",
button_secondary_border_color="transparent",
button_secondary_border_color_dark="transparent",
button_secondary_text_color="*neutral_700",
button_secondary_text_color_dark="*neutral_200",
)
css_code = ""
try:
with open("./style.css", "r", encoding="utf-8") as f:
css_code += f.read() + "\n"
except FileNotFoundError:
pass
title = """Panorama FLUX 🏞️✨
An advanced tiling pipeline for creative composition and large-scale image generation with the FLUX.1-schnell model.
"""
with gr.Blocks(css=css_code, theme=theme, title="Panorama FLUX") 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 for left side")
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 for the 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 for right side")
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 FLUX.1-schnell Image Demo](https://huggingface.co/spaces/black-forest-labs/FLUX.1-schnell)**
*The official demo for single image generation.*
- **[Z Image Turbo - Panorama](https://huggingface.co/spaces/elismasilva/z-image-panorama)**
*The Panorama app using the Z-Image Turbo model.*
"""
)
with gr.Sidebar():
gr.Markdown("### Tiling & Generation Parameters")
# New Language Selector
prompt_language = gr.Radio(
choices=["English", "Korean", "Chinese"],
value="English",
label="Prompt Language",
info="Select the language you will type your prompts in.",
)
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 Image Height", value=1024, interactive=False
)
new_target_width = gr.Textbox(
label="Actual Image Width", value=3072, interactive=False
)
with gr.Row():
steps = gr.Slider(
minimum=1, maximum=10, value=4, step=1, label="Inference Steps"
)
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=[
[
"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.0,
0.0,
200,
8,
619517442,
"Cosine",
"English",
1024,
3440,
0.0,
False
],
[
"A charming house in the countryside, by jakub rozalski, sunset lighting, elegant, highly detailed, sharp focus, artstation, stunning masterpiece",
"A dirt road in the countryside crossing pastures, by jakub rozalski, sunset lighting, elegant, highly detailed, 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, sharp focus, artstation, stunning masterpiece",
0.0,
0.0,
0.0,
256,
4,
358867853,
"Gaussian",
"English",
1024,
3808,
0.0,
False
],
[
"A vibrant mountain slope in full spring bloom, covered in colorful wildflowers and lush green grass, a small stream meandering down, cinematic photo, bright morning light.",
"The majestic, rocky peak of the same mountain under a clear summer sky, patches of green tundra, eagles soaring high above, strong midday sun. cinematic photo.",
"The other side of the mountain descending into a valley ablaze with autumn colors, forests of red, orange, and yellow trees, a gentle haze in the air. cinematic photo, golden hour light.",
0.0,
0.0,
0.0,
280,
4,
20240521,
"Cosine",
"English",
1024,
2896,
0.0,
False
],
[
"Enchanted forest with glowing bioluminescent plants and mystical fog, fantasy art.",
"Ancient elven castle with towering white spires bathed in moonlight, fantasy art.",
"Majestic griffin soaring above a starry night sky, fantasy art.",
0.0,
0.0,
0.0,
224,
6,
1029384756,
"Gaussian",
"English",
1024,
3392,
0.0,
False
],
[
"A serene, beautiful sandy beach with turquoise water and gentle waves, palm trees swaying. a vibrant, colorful parrot is flying. photorealistic, golden hour.",
"The calm ocean stretching to the horizon, a majestic sailboat with white sails in the distance, a few clouds in the sky. photorealistic, golden hour.",
"A lush tropical jungle with dense foliage, exotic flowers, and a waterfall cascading into a clear pool. a playful monkey is swinging from a vine. photorealistic, golden hour.",
0.0,
0.0,
0.0,
200,
4,
54321,
"Gaussian",
"English",
1024,
2672,
0.0,
False
],
[
"A futuristic neon-lit city street at night, flying cars zooming past, towering holographic billboards. cyberpunk aesthetic, blade runner style.",
"The entrance to a grimy, underground nightclub, a mysterious figure in a trench coat standing in the doorway, steam rising from a vent. cyberpunk aesthetic, blade runner style.",
"A dark alleyway off the main street, piles of discarded electronics, graffiti glowing on the walls, rain glistening on the pavement. cyberpunk aesthetic, blade runner style.",
0.0,
0.0,
0.0,
300,
8,
98765,
"Cosine",
"English",
1024,
3240,
0.0,
False
],
],
inputs=[
left_prompt,
center_prompt,
right_prompt,
left_gs,
center_gs,
right_gs,
overlap,
steps,
generation_seed,
tile_weighting_method,
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,
generation_seed,
tile_weighting_method,
prompt_language,
new_target_height,
new_target_width,
hdr,
randomize_seed
]
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_flux_panorama,
inputs=predict_inputs,
outputs=result,
show_progress="full",
api_name="generate_flux_panorama",
)
app.queue().launch(mcp_server=True, share=True)