app.py

import gradio as gr
import spaces
import os
import numpy as np
os.environ['SPCONV_ALGO'] = 'native'
from typing import *
import torch
import imageio
import shutil
from PIL import Image, ImageFilter
from easydict import EasyDict as edict
import utils.constants as constants
from haishoku.haishoku import Haishoku
from tqdm import tqdm

from tempfile import NamedTemporaryFile
import atexit
import random
import accelerate 
from transformers import AutoTokenizer, DPTImageProcessor, DPTForDepthEstimation
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import render_utils, postprocessing_utils
from pathlib import Path

import logging
#logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
import gc

# Import functions from modules
from utils.file_utils import cleanup_temp_files

from utils.color_utils import (
    hex_to_rgb,
    detect_color_format,
    update_color_opacity,
)
from utils.misc import (
    get_filename, 
    pause, 
    convert_ratio_to_dimensions,
    update_dimensions_on_ratio,
    get_seed,
    get_output_name
    ) #install_cuda_toolkit,install_torch, _get_output, setup_runtime_env)

from utils.image_utils import (
    change_color,
    open_image,
    upscale_image,
    lerp_imagemath,
    shrink_and_paste_on_blank,
    show_lut,
    apply_lut_to_image_path,
    multiply_and_blend_images,
    alpha_composite_with_control,
    crop_and_resize_image,
    resize_and_crop_image,
    convert_to_rgba_png,
    resize_image_with_aspect_ratio,
    build_prerendered_images_by_quality,
    get_image_from_dict,
    calculate_optimal_fill_dimensions,
    save_image_to_temp_png
)

from utils.hex_grid import (
    generate_hexagon_grid,
    generate_hexagon_grid_interface,
)

from utils.excluded_colors import (
    add_color,
    delete_color,
    build_dataframe,
    on_input,
    excluded_color_list,
    on_color_display_select
)

# from utils.ai_generator import (
#     generate_ai_image,
# )

from utils.lora_details import (
    upd_prompt_notes,
    upd_prompt_notes_by_index,
    split_prompt_precisely,
    approximate_token_count,
    get_trigger_words,
    is_lora_loaded,
    get_lora_models
)
from diffusers import (
    FluxPipeline,
    FluxImg2ImgPipeline,
    FluxControlPipeline,
    FluxControlPipeline,
    DiffusionPipeline, 
    AutoencoderTiny, 
    AutoencoderKL
)

PIPELINE_CLASSES = {
    "FluxPipeline": FluxPipeline,
    "FluxImg2ImgPipeline": FluxImg2ImgPipeline,
    "FluxControlPipeline": FluxControlPipeline,
    "FluxFillPipeline": FluxControlPipeline
}

from utils.version_info import (
    versions_html,
    #initialize_cuda,
    #release_torch_resources,
    #get_torch_info
)
#from utils.depth_estimation import (get_depth_map_from_state)

input_image_palette = []
current_prerendered_image = gr.State("./images/Beeuty-1.png")
user_dir = constants.TMPDIR
lora_models = get_lora_models()
selected_index = gr.State(value=-1)

image_processor: Optional[DPTImageProcessor] = None
depth_model: Optional[DPTForDepthEstimation] = None
TRELLIS_PIPELINE: Optional[TrellisImageTo3DPipeline] = None
pipe: Optional[Union[FluxPipeline, FluxImg2ImgPipeline, FluxControlPipeline]] = None

def start_session(req: gr.Request):
    print(f"Starting session with hash: {req.session_hash}")
    user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    
    
def end_session(req: gr.Request):
    print(f"Ending session with hash: {req.session_hash}")
    user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
    shutil.rmtree(user_dir)

# Register the cleanup function
atexit.register(end_session)

def hex_create(hex_size, border_size, input_image_path, start_x, start_y, end_x, end_y, rotation, background_color_hex, background_opacity, border_color_hex, border_opacity, fill_hex, excluded_colors_var, filter_color, x_spacing, y_spacing, add_hex_text_option=None, custom_text_list=None, custom_text_color_list=None):
    global input_image_palette

    try:
        # Load and process the input image
        input_image = Image.open(input_image_path).convert("RGBA")
    except Exception as e:
        print(f"Failed to convert image to RGBA: {e}")
        # Open the original image without conversion
        input_image = Image.open(input_image_path)
        # Ensure the canvas is at least 1344x768 pixels
        min_width, min_height = 1344, 768
        canvas_width = max(min_width, input_image.width)
        canvas_height = max(min_height, input_image.height)

        # Create a transparent canvas with the required dimensions
        new_canvas = Image.new("RGBA", (canvas_width, canvas_height), (0, 0, 0, 0))

        # Calculate position to center the input image on the canvas
        paste_x = (canvas_width - input_image.width) // 2
        paste_y = (canvas_height - input_image.height) // 2

        # Paste the input image onto the canvas
        new_canvas.paste(input_image, (paste_x, paste_y))

        # Save the 'RGBA' image to a temporary file and update 'input_image_path'
        with NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
            new_canvas.save(tmp_file.name, format="PNG")
            input_image_path = tmp_file.name
            constants.temp_files.append(tmp_file.name)

        # Update 'input_image' with the new image as a file path
        input_image = Image.open(input_image_path)
    
    # Use Haishoku to get the palette from the new image
    input_palette = Haishoku.loadHaishoku(input_image_path)
    input_image_palette = input_palette.palette

    # Update colors with opacity
    background_color = update_color_opacity(
        hex_to_rgb(background_color_hex),
        int(background_opacity * (255 / 100))
    )
    border_color = update_color_opacity(
        hex_to_rgb(border_color_hex),
        int(border_opacity * (255 / 100))
    )

    # Prepare excluded colors list
    excluded_color_list = [tuple(lst) for lst in excluded_colors_var]

    # Generate the hexagon grid images
    grid_image = generate_hexagon_grid_interface(
        hex_size,
        border_size,
        input_image,
        start_x,
        start_y,
        end_x,
        end_y,
        rotation,
        background_color,
        border_color,
        fill_hex,
        excluded_color_list,
        filter_color,
        x_spacing,
        y_spacing,
        add_hex_text_option,
        custom_text_list,
        custom_text_color_list
    )    
    return grid_image

def get_model_and_lora(model_textbox):
    """
    Determines the model and LoRA weights based on the model_textbox input.
    wieghts must be in an array ["Borcherding/FLUX.1-dev-LoRA-FractalLand-v0.1"]
    """
    # If the input is in the list of models, return it with None as LoRA weights
    if model_textbox in constants.MODELS:
        return model_textbox, []
    # If the input is in the list of LoRA weights, get the corresponding model
    elif model_textbox in constants.LORA_WEIGHTS:
        model = constants.LORA_TO_MODEL.get(model_textbox)
        return model, model_textbox.split()
    else:
        # Default to a known model if input is unrecognized
        default_model = model_textbox
        return default_model, []

def set_pipeline(
    model_name="black-forest-labs/FLUX.1-dev",
    lora_weights=None,
    pipeline_name="FluxPipeline",
    progress=gr.Progress(track_tqdm=True)
):      
    global pipe
    pbar = tqdm(total=7, desc="Pipeline and Model Load")
    current_pipeline_name =pipe.name_or_path if pipe else None
    current_pipeline_class = type(pipe).__name__ if pipe else None
    if (current_pipeline_name != model_name) or (pipeline_name != current_pipeline_class):
        pipe = None
        gc.collect()
        #from torch import cuda, bfloat16, float32, Generator, no_grad, backends
        # Retrieve the pipeline class from the mapping
        pipeline_class = PIPELINE_CLASSES.get(pipeline_name)
        if not pipeline_class:
            raise ValueError(f"Unsupported pipeline type '{pipeline_name}'. "
                            f"Available options: {list(PIPELINE_CLASSES.keys())}")

        #initialize_cuda()
        dvc = "cpu"
        #from src.condition import Condition
        pbar.update(1)
        print(f"device:{device}\nmodel_name:{model_name}\nlora_weights:{lora_weights}\n")
        #print(f"\n {get_torch_info()}\n")
        # Initialize the pipeline inside the context manager
        pipe = pipeline_class.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float16,
            vae=good_vae
        )
        pbar.update(2)
        pipe.to(dvc)
        # Optionally, don't use CPU offload if not necessary
        pbar.update(1)

        # Access the tokenizer from the pipeline
        tokenizer = pipe.tokenizer

        # Check if add_prefix_space is set and convert to slow tokenizer if necessary
        if getattr(tokenizer, 'add_prefix_space', False):
            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, device_map = 'cpu')
            # Update the pipeline's tokenizer
            pipe.tokenizer = tokenizer
            
    pbar.set_description("Loading LoRA weights")
    pbar.update(1)
    pipe.unload_lora_weights()
    
    # Load LoRA weights
    # note: does not yet handle multiple LoRA weights with different names, needs .set_adapters(["depth", "hyper-sd"], adapter_weights=[0.85, 0.125])
    if lora_weights:
        for lora_weight in lora_weights:
            lora_configs = constants.LORA_DETAILS.get(lora_weight, [])
            lora_weight_set = False
            if lora_configs:
                for config in lora_configs:
                    # Load LoRA weights with optional weight_name and adapter_name
                    if 'weight_name' in config:
                        weight_name = config.get("weight_name")
                        adapter_name = config.get("adapter_name")
                        lora_collection = config.get("lora_collection")
                        if weight_name and adapter_name and lora_collection and lora_weight_set == False:
                            pipe.load_lora_weights(
                                lora_collection,
                                weight_name=weight_name,
                                adapter_name=adapter_name,
                                token=constants.HF_API_TOKEN
                            )
                            lora_weight_set = True
                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}, lora_collection={lora_collection}\n")  
                        elif weight_name and adapter_name==None and lora_collection and lora_weight_set == False:
                            pipe.load_lora_weights(
                                lora_collection,
                                weight_name=weight_name,
                                token=constants.HF_API_TOKEN
                            )
                            lora_weight_set = True
                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}, lora_collection={lora_collection}\n")  
                        elif weight_name and adapter_name and lora_weight_set == False:
                            pipe.load_lora_weights(
                                lora_weight,
                                weight_name=weight_name,
                                adapter_name=adapter_name,
                                token=constants.HF_API_TOKEN
                            )
                            lora_weight_set = True
                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
                        elif weight_name and adapter_name==None and lora_weight_set == False:
                            pipe.load_lora_weights(
                                lora_weight,
                                weight_name=weight_name,
                                token=constants.HF_API_TOKEN
                            )
                            lora_weight_set = True
                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
                        elif lora_weight_set == False:
                            pipe.load_lora_weights(
                                lora_weight,
                                token=constants.HF_API_TOKEN
                            )  
                            lora_weight_set = True
                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
                    # Apply 'pipe' configurations if present
                    if 'pipe' in config:
                        pipe_config = config['pipe']
                        for method_name, params in pipe_config.items():
                            method = getattr(pipe, method_name, None)
                            if method:
                                print(f"Applying pipe method: {method_name} with params: {params}")
                                method(**params)
                            else:
                                print(f"Method {method_name} not found in pipe.")
                    if 'condition_type' in config:
                        condition_type = config['condition_type']
                        if condition_type == "coloring":
                            #pipe.enable_coloring()
                            print("\nEnabled coloring.\n")
                        elif condition_type == "deblurring":
                            #pipe.enable_deblurring()
                            print("\nEnabled deblurring.\n")
                        elif condition_type == "fill":
                            #pipe.enable_fill()
                            print("\nEnabled fill.\n")
                        elif condition_type == "depth":
                            #pipe.enable_depth()
                            print("\nEnabled depth.\n")
                        elif condition_type == "canny":
                            #pipe.enable_canny()
                            print("\nEnabled canny.\n")
                        elif condition_type == "subject":
                            #pipe.enable_subject()
                            print("\nEnabled subject.\n")
                        else:
                            print(f"Condition type {condition_type} not implemented.")
            else:
                pipe.load_lora_weights(lora_weight, use_auth_token=constants.HF_API_TOKEN)
    pbar.set_description("Pipe Loaded.")
    pbar.set_postfix({"Status": "Done"})
    pbar.update(1)
    pbar.close()


@spaces.GPU(duration=200, progress=gr.Progress(track_tqdm=True))
def generate_image_lowmem(
    text,
    neg_prompt=None,
    model_name="black-forest-labs/FLUX.1-dev",
    lora_weights=None,
    conditioned_image=None,
    mask_image=None,
    image_width=1368,
    image_height=848,
    guidance_scale=3.5,
    num_inference_steps=30,
    seed=0,
    true_cfg_scale=1.0,
    pipeline_name="FluxPipeline",
    strength=0.75,
    additional_parameters=None,
    progress=gr.Progress(track_tqdm=True)
):  
    with torch.no_grad():
        #global pipe
        global device
        pipe.to(device)
        flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled()
        if flash_attention_enabled == False:
            #Enable xFormers memory-efficient attention (optional)
            #pipe.enable_xformers_memory_efficient_attention()
            print("\nEnabled xFormers memory-efficient attention.\n")
        else:            
            pipe.attn_implementation="flash_attention_2"
            print("\nEnabled flash_attention_2.\n")
        # alternative version that may be more efficient
        # pipe.enable_sequential_cpu_offload()
        if pipeline_name == "FluxPipeline":
            pipe.enable_model_cpu_offload()
            pipe.vae.enable_slicing()
            #pipe.vae.enable_tiling()
        else:
            pipe.enable_model_cpu_offload()


        mask_parameters = {}
        # Load the mask image if provided
        if (pipeline_name == "FluxFillPipeline"):
            mask_image = open_image(mask_image).convert("RGBA")
            mask_condition_type = constants.condition_type[5]
            guidance_scale = 30
            num_inference_steps=50
            max_sequence_length=512
            print(f"\nAdded mask image.\n {mask_image.size}")
            mask_parameters ={
                "mask_image": mask_image,
            }

        # Set the random seed for reproducibility
        generator = torch.Generator(device=device).manual_seed(seed)
        #conditions = []
        if conditioned_image is not None:
            conditioned_image = resize_and_crop_image(conditioned_image, image_width, image_height)
            #condition = Condition(constants.condition_type[2], conditioned_image)
            #conditions.append(condition)
            print(f"\nAdded conditioned image.\n {conditioned_image.size}")
            # Prepare the parameters for image generation
            additional_parameters ={
                "strength": strength,
                "image": conditioned_image,
            }
            additional_parameters.update(mask_parameters)
        else:
            print("\nNo conditioned image provided.")
            if neg_prompt!=None:
                true_cfg_scale=1.1
            additional_parameters ={
                "negative_prompt": neg_prompt,
                "true_cfg_scale": true_cfg_scale,
            }
        # handle long prompts by splitting them
        if approximate_token_count(text) > 76:
            prompt, prompt2 = split_prompt_precisely(text)
            prompt_parameters = {
                "prompt" : prompt,
                "prompt_2": prompt2,
            }
        else:
            prompt_parameters = {
                "prompt" :text,
        }
        additional_parameters.update(prompt_parameters)
        # Combine all parameters
        generate_params = {
            "height": image_height,
            "width": image_width,
            "guidance_scale": guidance_scale,
            "num_inference_steps": num_inference_steps,
            "generator": generator,
        }
        if additional_parameters:
            generate_params.update(additional_parameters)
        generate_params = {k: v for k, v in generate_params.items() if v is not None}
        print(f"generate_params: {generate_params}")
        # Generate the image
        try:
            result = pipe(**generate_params) #generate_image(pipe,generate_params)
            image = result.images[0]
            # Clean up
            del result
        except Exception as e:
            print(f"Error generating image: {e}")
            image = open_image("./images/Beeuty-1.png")
    #del conditions
    del generator
    # Move the pipeline and clear cache
    pipe.to("cpu")
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()
    print(torch.cuda.memory_summary(device=None, abbreviated=False))
    gc.collect()
    return image

def generate_ai_image_local (
    map_option,
    prompt_textbox_value,
    neg_prompt_textbox_value,
    model="black-forest-labs/FLUX.1-dev",
    lora_weights=None,
    conditioned_image=None,
    mask_image=None,
    height=512,
    width=912,
    num_inference_steps=30,
    guidance_scale=3.5,
    seed=777,
    pipeline_name="FluxPipeline",
    strength=0.75,
    progress=gr.Progress(track_tqdm=True)
):    
    print(f"Generating image with lowmem")
    try:
        if map_option != "Prompt":
            prompt = constants.PROMPTS[map_option]
            negative_prompt = constants.NEGATIVE_PROMPTS.get(map_option, "")
        else:
            prompt = prompt_textbox_value
            negative_prompt = neg_prompt_textbox_value or ""
        #full_prompt = f"{prompt} {negative_prompt}"
        additional_parameters = {}
        if lora_weights:
            for lora_weight in lora_weights:
                lora_configs = constants.LORA_DETAILS.get(lora_weight, [])
                for config in lora_configs:
                    if 'parameters' in config:
                        additional_parameters.update(config['parameters'])
                    elif 'trigger_words' in config:
                        trigger_words = get_trigger_words(lora_weight)
                        prompt = f"{trigger_words} {prompt}"
        for key, value in additional_parameters.items():
            if key in ['height', 'width', 'num_inference_steps', 'max_sequence_length']:
                additional_parameters[key] = int(value)
            elif key in ['guidance_scale','true_cfg_scale']:
                additional_parameters[key] = float(value)
        height = additional_parameters.pop('height', height)
        width = additional_parameters.pop('width', width)
        num_inference_steps = additional_parameters.pop('num_inference_steps', num_inference_steps)
        guidance_scale = additional_parameters.pop('guidance_scale', guidance_scale)
        print("Generating image with the following parameters:\n")
        print(f"Model: {model}")
        print(f"LoRA Weights: {lora_weights}")
        print(f"Prompt: {prompt}")
        print(f"Neg Prompt: {negative_prompt}")
        print(f"Height: {height}")
        print(f"Width: {width}")
        print(f"Number of Inference Steps: {num_inference_steps}")
        print(f"Guidance Scale: {guidance_scale}")
        print(f"Seed: {seed}")
        print(f"Additional Parameters: {additional_parameters}")
        print(f"Conditioned Image: {conditioned_image}")
        print(f"Conditioned Image Strength: {strength}")
        print(f"pipeline: {pipeline_name}")
        set_pipeline(
            model_name=model,
            lora_weights=lora_weights,
            pipeline_name=pipeline_name,
            progress=progress
        )
        image = generate_image_lowmem(
            text=prompt,
            model_name=model,
            neg_prompt=negative_prompt,
            lora_weights=lora_weights,
            conditioned_image=conditioned_image,
            mask_image=mask_image,
            image_width=width,
            image_height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            seed=seed,
            pipeline_name=pipeline_name,
            strength=strength,
            additional_parameters=additional_parameters,
            progress=progress
        )
        with NamedTemporaryFile(delete=False, suffix=".png") as tmp:
            image.save(tmp.name, format="PNG")
            constants.temp_files.append(tmp.name)
            print(f"Image saved to {tmp.name}")
            return tmp.name
    except Exception as e:
        print(f"Error generating AI image: {e}")
        gc.collect()
        return None


def generate_input_image_click(image_input, map_option, prompt_textbox_value, negative_prompt_textbox_value, model_textbox_value, randomize_seed=True, seed=None, use_conditioned_image=False, mask_image=None, strength=0.5, image_format="16:9", scale_factor=constants.SCALE_FACTOR, progress=gr.Progress(track_tqdm=True)):
    seed = get_seed(randomize_seed, seed)

    # Get the model and LoRA weights
    model, lora_weights = get_model_and_lora(model_textbox_value)
    global current_prerendered_image
    conditioned_image=None
    formatted_map_option = map_option.lower().replace(' ', '_')

    if use_conditioned_image:
        print(f"Conditioned path: {current_prerendered_image.value}.. converting to RGB\n")
        # ensure the conditioned image is an image and not a string, cannot use RGBA
        if isinstance(current_prerendered_image.value, str):
            conditioned_image = open_image(current_prerendered_image.value).convert("RGB")
            print(f"Conditioned Image: {conditioned_image.size}.. converted to RGB\n")
    # use image_input as the conditioned_image if it is not None
    elif image_input is not None:
        file_path, is_dict = get_image_from_dict(image_input)
        conditioned_image =  open_image(file_path).convert("RGB")
        print(f"Conditioned Image set to modify Input Image!\nClear to generate new image from layered image: {is_dict}\n")
        gr.Info(f"Conditioned Image set to modify Input Image! Clear to generate new image. Layered: {is_dict}",duration=5)

    # Convert image_format from a string split by ":" into two numbers divided
    width_ratio, height_ratio = map(int, image_format.split(":"))
    aspect_ratio = width_ratio / height_ratio
    
    width, height = convert_ratio_to_dimensions(aspect_ratio, constants.BASE_HEIGHT)
    pipeline = "FluxPipeline"
    if conditioned_image is not None:
        pipeline = "FluxImg2ImgPipeline"

    if (model == "black-forest-labs/FLUX.1-Fill-dev"):
        pipeline = "FluxFillPipeline"
        width, height = calculate_optimal_fill_dimensions(conditioned_image)
        mask_image = get_image_from_dict(mask_image)
        print(f"Optimal Dimensions: {width} x {height} \n")
    else:
        mask_image = None

    # Generate the AI image and get the image path
    image_path = generate_ai_image_local(
        map_option,
        prompt_textbox_value,
        negative_prompt_textbox_value,
        model,
        lora_weights,
        conditioned_image,
        mask_image,
        strength=strength,
        height=height,
        width=width,
        seed=seed,
        pipeline_name=pipeline,
    )
       
    # Open the generated image
    try:
        image = Image.open(image_path).convert("RGBA")
    except Exception as e:
        print(f"Failed to open generated image: {e}")
        return image_path, seed  # Return the original image path if opening fails
       
    # Upscale the image
    upscaled_image = upscale_image(image, scale_factor)
       
    # Save the upscaled image to a temporary file
    with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_") as tmp_upscaled:
        upscaled_image.save(tmp_upscaled.name, format="PNG")
        constants.temp_files.append(tmp_upscaled.name)
        print(f"Upscaled image saved to {tmp_upscaled.name}")
    gc.collect()
    # Return the path of the upscaled image
    return tmp_upscaled.name, seed

def update_prompt_visibility(map_option):
      is_visible = (map_option == "Prompt")
      return (
          gr.update(visible=is_visible),
          gr.update(visible=is_visible),
          gr.update(visible=is_visible)
      )

def update_prompt_notes(model_textbox_value):
    return upd_prompt_notes(model_textbox_value)

def update_selection(evt: gr.SelectData, aspect_ratio):
    selected_lora = constants.LORAS[evt.index]
    new_placeholder = f"Type a prompt for {selected_lora['title']}"
    new_aspect_ratio = aspect_ratio  # default to the currently selected aspect ratio
    lora_repo = selected_lora["repo"]
    #updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅"
    # If the selected LoRA model specifies an aspect ratio, use it to update dimensions.
    if "aspect" in selected_lora:
        try:
            new_aspect_ratio = selected_lora["aspect"]
            # Recalculate dimensions using constants.BASE_HEIGHT as the height reference.
            new_width, new_height = update_dimensions_on_ratio(new_aspect_ratio, constants.BASE_HEIGHT)
            # (Optionally, you could log or use new_width/new_height as needed)
        except Exception as e:
            print(f"\nError in update selection aspect ratios: {e}\nSkipping")
    return [gr.update(value=lora_repo), gr.update(visible=(lora_repo == "Manual Entry")), evt.index, new_aspect_ratio, upd_prompt_notes_by_index(evt.index)]

def on_prerendered_gallery_selection(event_data: gr.SelectData):
    global current_prerendered_image
    selected_index = event_data.index
    selected_image = constants.pre_rendered_maps_paths[selected_index]
    print(f"Template Image Selected: {selected_image} ({event_data.index})\n")
    gr.Info(f"Template Image Selected: {selected_image} ({event_data.index})",duration=5)
    current_prerendered_image.value = selected_image
    return current_prerendered_image

def combine_images_with_lerp(input_image, output_image, alpha):
    in_image = open_image(input_image)
    out_image = open_image(output_image)
    print(f"Combining images with alpha: {alpha}")
    return lerp_imagemath(in_image, out_image, alpha)

def add_border(image, mask_width, mask_height, blank_color):
    bordered_image_output = Image.open(image).convert("RGBA")
    margin_color = detect_color_format(blank_color)
    print(f"Adding border to image with width: {mask_width}, height: {mask_height}, color: {margin_color}")
    return shrink_and_paste_on_blank(bordered_image_output, mask_width, mask_height, margin_color)

def on_input_image_change(image_path):
    if image_path is None:
        gr.Warning("Please upload an Input Image to get started.")
        return None, gr.update()
    img, img_path = convert_to_rgba_png(image_path)
    width, height = img.size
    return [img_path, gr.update(width=width, height=height)]

def update_sketch_dimensions(input_image, sketch_image):
    # Load the images using open_image() if they are provided as file paths.
    in_img = open_image(input_image) if isinstance(input_image, str) else input_image
    sk_img_path, _ = get_image_from_dict(sketch_image)
    sk_img = open_image(sk_img_path)
    # Resize sketch image if dimensions don't match input image.
    if (in_img) and (in_img.size != sk_img.size):
        sk_img = sk_img.resize(in_img.size, Image.LANCZOS)
        return [sk_img, gr.update(width=in_img.width, height=in_img.height)]
    else:
        return [sk_img, gr.update()]
    

def composite_with_control_sync(input_image, sketch_image, slider_value):
    # Load the images using open_image() if they are provided as file paths.
    in_img = open_image(input_image) if isinstance(input_image, str) else input_image
    sk_img_path, _ = get_image_from_dict(sketch_image)
    sk_img = open_image(sk_img_path)

    # Resize sketch image if dimensions don't match input image.
    if in_img.size != sk_img.size:
        sk_img = sk_img.resize(in_img.size, Image.LANCZOS)

    # Now composite using the original alpha_composite_with_control function.
    result_img = alpha_composite_with_control(in_img, sk_img, slider_value)
    return result_img

def replace_input_with_sketch_image(sketch_image):
    print(f"Sketch Image: {sketch_image}\n")
    sketch, is_dict = get_image_from_dict(sketch_image)
    return sketch
####################################### DEPTH ESTIMATION #######################################

@spaces.GPU(progress=gr.Progress(track_tqdm=True))
def load_trellis_model():
    global TRELLIS_PIPELINE
    loaded = False
    if TRELLIS_PIPELINE == None:
        try:
            TRELLIS_PIPELINE = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
            TRELLIS_PIPELINE.cuda()
            # Preload with a dummy image to finalize initialization
            try:
                TRELLIS_PIPELINE.preprocess_image(Image.fromarray(np.zeros((512, 512, 4), dtype=np.uint8)))    # Preload rembg
            except:
                pass
            print("TRELLIS_PIPELINE loaded\n")
            gr.Info("TRELLIS_PIPELINE loaded")
            loaded = True
        except Exception as e:
            print(f"Error preloading TRELLIS_PIPELINE: {e}")
            gr.Error(f"Failed to load TRELLIS_PIPELINE: {e}")
            TRELLIS_PIPELINE = None
    

def load_3d_models(is_open: bool = True) -> bool:
    if is_open:
        gr.Info("Loading 3D models...")
        global image_processor, depth_model
        image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
        depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large", ignore_mismatched_sizes=True)
        print("DPT models loaded\n")
        # is_open = load_trellis_model()
        # if is_open:
        print("3D models loaded")
        gr.Info("3D models loaded.")
        #else:
        #     gr.Error("Failed to load TRELLIS_PIPELINE.")
    return gr.update(interactive = is_open)

def unload_3d_models(is_open: bool = False) -> bool:
    if not is_open:
        gr.Info("Unloading 3D models...")
        global image_processor, depth_model, TRELLIS_PIPELINE
        if TRELLIS_PIPELINE:
            TRELLIS_PIPELINE.to("cpu")
            TRELLIS_PIPELINE = None
        if depth_model:
            del image_processor
            del depth_model
        # torch.cuda.empty_cache()
        # torch.cuda.ipc_collect()
        gc.collect()
        print("3D models unloaded and CUDA memory freed")
        gr.Info("3D models unloaded.")
    return gr.update(interactive = is_open)

def preprocess_image(image: Image.Image) -> Image.Image:
    """
    Preprocess the input image.
    Args:
        image (Image.Image): The input image.
    Returns:
        Image.Image: The preprocessed image.
    """
    processed_image = TRELLIS_PIPELINE.preprocess_image(image)
    return processed_image

def pack_state(gs: Gaussian, mesh: MeshExtractResult, name: str) -> dict:
    return {
        'gaussian': {
            **gs.init_params,
            '_xyz': gs._xyz.cpu().numpy(),
            '_features_dc': gs._features_dc.cpu().numpy(),
            '_scaling': gs._scaling.cpu().numpy(),
            '_rotation': gs._rotation.cpu().numpy(),
            '_opacity': gs._opacity.cpu().numpy(),
        },
        'mesh': {
            'vertices': mesh.vertices.cpu().numpy(),
            'faces': mesh.faces.cpu().numpy(),
        },
        'name': name
    }

def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
    gs = Gaussian(
        aabb=state['gaussian']['aabb'],
        sh_degree=state['gaussian']['sh_degree'],
        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
        scaling_bias=state['gaussian']['scaling_bias'],
        opacity_bias=state['gaussian']['opacity_bias'],
        scaling_activation=state['gaussian']['scaling_activation'],
    )
    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
    
    mesh = edict(
        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
    )

    name = state['name']
    
    return gs, mesh, name

@spaces.GPU()
def depth_process_image(image_path, resized_width=800, z_scale=208):
    """
    Processes the input image to generate a depth map.

    Args:
        image_path (str): The file path to the input image.
        resized_width (int, optional): The width to which the image is resized. Defaults to 800.
        z_scale (int, optional): Z-axis scale factor. Defaults to 208.

    Returns:
        list: A list containing the depth image.
    """
    image_path = Path(image_path)
    if not image_path.exists():
        raise ValueError("Image file not found")

    # Load and resize the image
    image_raw = Image.open(image_path).convert("RGB")
    print(f"Original size: {image_raw.size}")
    resized_height = int(resized_width * image_raw.size[1] / image_raw.size[0])
    image = image_raw.resize((resized_width, resized_height), Image.Resampling.LANCZOS)
    print(f"Resized size: {image.size}")

    # Prepare image for the model
    encoding = image_processor(image, return_tensors="pt")

    # Perform depth estimation
    with torch.no_grad():
        outputs = depth_model(**encoding)
        predicted_depth = outputs.predicted_depth

    # Interpolate depth to match the image size
    prediction = torch.nn.functional.interpolate(
        predicted_depth.unsqueeze(1),
        size=(image.height, image.width),
        mode="bicubic",
        align_corners=False,
    ).squeeze()

    # Normalize the depth image to 8-bit
    if torch.cuda.is_available():
        prediction = prediction.numpy()
    else:
        prediction = prediction.cpu().numpy()
    depth_min, depth_max = prediction.min(), prediction.max()
    depth_image = ((prediction - depth_min) / (depth_max - depth_min) * 255).astype("uint8")
    img = Image.fromarray(depth_image)

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
    return img

def generate_3d_asset_part1(depth_image_source, randomize_seed, seed, input_image, output_image, overlay_image, bordered_image_output, req: gr.Request, progress=gr.Progress(track_tqdm=True)):
    # Choose the image based on source
    if depth_image_source == "Input Image":
        image_path = input_image
    elif depth_image_source == "Output Image":
        image_path = output_image
    elif depth_image_source == "Image with Margins":
        image_path = bordered_image_output
    else:  # "Overlay Image"
        image_path = overlay_image

    output_name = get_output_name(input_image, output_image, overlay_image, bordered_image_output)

    # Ensure the file exists
    if not Path(image_path).exists():
        raise ValueError("Image file not found.")

    # Determine the final seed using default MAX_SEED from constants
    final_seed = np.random.randint(0, constants.MAX_SEED) if randomize_seed else seed
    # Process the image for depth estimation
    depth_img = depth_process_image(image_path, resized_width=1536, z_scale=336)
    depth_img = resize_image_with_aspect_ratio(depth_img, 1536, 1536)
    
    user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
    depth_img = save_image_to_temp_png(depth_img, user_dir, f"{output_name}_depth")

    return depth_img, image_path, output_name, final_seed

@spaces.GPU(duration=150,progress=gr.Progress(track_tqdm=True))
def generate_3d_asset_part2(depth_img, image_path, output_name, seed, steps, model_resolution, video_resolution, req: gr.Request, progress=gr.Progress(track_tqdm=True)):
    # Open image using standardized defaults
    image_raw = Image.open(image_path).convert("RGB")
    resized_image = resize_image_with_aspect_ratio(image_raw, model_resolution, model_resolution)
    depth_img = Image.open(depth_img).convert("RGBA")
    # Preprocess and run the Trellis pipeline with fixed sampler settings
    try:
        #TRELLIS_PIPELINE.cuda()
        processed_image = TRELLIS_PIPELINE.preprocess_image(resized_image, max_resolution=model_resolution)
        outputs = TRELLIS_PIPELINE.run(
            processed_image,
            seed=seed,
            formats=["gaussian", "mesh"],
            preprocess_image=False,
            sparse_structure_sampler_params={
                "steps": steps,
                "cfg_strength": 7.5,
            },
            slat_sampler_params={
                "steps": steps,
                "cfg_strength": 3.0,
            },
        )

        # Validate the mesh
        mesh = outputs['mesh'][0]
        meshisdict = isinstance(mesh, dict)
        if meshisdict:
            vertices = mesh['vertices']
            faces = mesh['faces']
        else:
            vertices = mesh.vertices
            faces = mesh.faces

        print(f"Mesh vertices: {vertices.shape}, faces: {faces.shape}")
        if faces.max() >= vertices.shape[0]:
            raise ValueError(f"Invalid mesh: face index {faces.max()} exceeds vertex count {vertices.shape[0]}")
    except Exception as e:
        gr.Warning(f"Error generating 3D asset: {e}")
        print(f"Error generating 3D asset: {e}")
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
        return None,None, depth_img

    # Ensure data is on GPU and has correct type
    if not vertices.is_cuda or not faces.is_cuda:
        raise ValueError("Mesh data must be on GPU")
    if vertices.dtype != torch.float32 or faces.dtype != torch.int32:
        if meshisdict:
            mesh['faces'] = faces.to(torch.int32)
            mesh['vertices'] = vertices.to(torch.float32)
        else:
            mesh.faces = faces.to(torch.int32)
            mesh.vertices = vertices.to(torch.float32)

    user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)

    video = render_utils.render_video(outputs['gaussian'][0], resolution=video_resolution, num_frames=64, r=1, fov=45)['color']
    try:
        video_geo = render_utils.render_video(outputs['mesh'][0], resolution=video_resolution, num_frames=64, r=1, fov=45)['normal']
        video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
    except Exception as e:
        gr.Info(f"Error rendering video: {e}")
        print(f"Error rendering video: {e}")
    video_path = os.path.join(user_dir, f'{output_name}.mp4')
    imageio.mimsave(video_path, video, fps=8)

    #snapshot_results = render_utils.render_snapshot_depth(outputs['mesh'][0], resolution=1280, r=1, fov=80)
    #depth_snapshot = Image.fromarray(snapshot_results['normal'][0]).convert("L")
    depth_snapshot = depth_img

    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], output_name)
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
    return [state, video_path, depth_snapshot]


@spaces.GPU(duration=90,progress=gr.Progress(track_tqdm=True))
def extract_glb(
    state: dict,
    mesh_simplify: float,
    texture_size: int,
    req: gr.Request,progress=gr.Progress(track_tqdm=True)
) -> Tuple[str, str]:
    """
    Extract a GLB file from the 3D model.

    Args:
        state (dict): The state of the generated 3D model.
        mesh_simplify (float): The mesh simplification factor.
        texture_size (int): The texture resolution.

    Returns:
        str: The path to the extracted GLB file.
    """
    user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
    gs, mesh, name = unpack_state(state)
    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
    glb_path = os.path.join(user_dir, f'{name}.glb')
    glb.export(glb_path)

    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()
    return glb_path, glb_path

@spaces.GPU(progress=gr.Progress(track_tqdm=True))
def extract_gaussian(state: dict, req: gr.Request, progress=gr.Progress(track_tqdm=True)) -> Tuple[str, str]:
    """
    Extract a Gaussian file from the 3D model.

    Args:
        state (dict): The state of the generated 3D model.

    Returns:
        str: The path to the extracted Gaussian file.
    """
    user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
    gs, _, name = unpack_state(state)
    gaussian_path = os.path.join(user_dir, f'{name}.ply')
    gs.save_ply(gaussian_path)

    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()
    return gaussian_path, gaussian_path


@spaces.GPU()
def getVersions():
    return versions_html()

#generate_input_image_click.zerogpu = True
#generate_depth_button_click.zerogpu = True
#def main(debug=False):
title = "HexaGrid Creator"
#description = "Customizable Hexagon Grid Image Generator"
examples = [["assets//examples//hex_map_p1.png", 32, 1, 0, 0, 0, 0, 0, "#ede9ac44","#12165380", True]]

gr.set_static_paths(paths=["images/","images/images","images/prerendered","LUT/","fonts/","assets/"])

# Gradio Blocks Interface
with gr.Blocks(css_paths="style_20250314.css", title=title, theme='Surn/beeuty',delete_cache=(21600,86400)) as hexaGrid:
    with gr.Row():
        gr.Markdown("""
        # HexaGrid Creator
        ## Transform Your Images into Mesmerizing Hexagon Grid Masterpieces! ⬢""", elem_classes="intro")
    with gr.Row():
        with gr.Accordion("Welcome to HexaGrid Creator, the ultimate tool for transforming your images into stunning hexagon grid artworks. Whether you're a tabletop game enthusiast, a digital artist, or someone who loves unique patterns, HexaGrid Creator has something for you.", open=False, elem_classes="intro"):
            gr.Markdown ("""

            ## Drop an image into the Input Image and get started!

        

            ## What is HexaGrid Creator?
            HexaGrid Creator is a web-based application that allows you to apply a hexagon grid overlay to any image. You can customize the size, color, and opacity of the hexagons, as well as the background and border colors. The result is a visually striking image that looks like it was made from hexagonal tiles!

            ### What Can You Do?
            - **Generate Hexagon Grids:** Create beautiful hexagon grid overlays on any image with fully customizable parameters.
            - **AI-Powered Image Generation:** Use advanced AI models to generate images based on your prompts and apply hexagon grids to them.
            - **Color Exclusion:** Select and exclude specific colors from your hexagon grid for a cleaner and more refined look.
            - **Interactive Customization:** Adjust hexagon size, border size, rotation, background color, and more in real-time.
            - **Depth and 3D Model Generation:** Generate depth maps and 3D models from your images for enhanced visualization.
            - **Image Filter [Look-Up Table (LUT)] Application:** Apply filters (LUTs) to your images for color grading and enhancement.
            - **Pre-rendered Maps:** Access a library of pre-rendered hexagon maps for quick and easy customization.
            - **Add Margins:** Add customizable margins around your images for a polished finish.

            ### Why You'll Love It
            - **Fun and Easy to Use:** With an intuitive interface and real-time previews, creating hexagon grids has never been this fun!
            - **Endless Creativity:** Unleash your creativity with endless customization options and see your images transform in unique ways.
            - **Hexagon-Inspired Theme:** Enjoy a delightful yellow and purple theme inspired by hexagons! ⬢
            - **Advanced AI Models:** Leverage advanced AI models and LoRA weights for high-quality image generation and customization.

            ### Get Started
            1. **Upload or Generate an Image:** Start by uploading your own image or generate one using our AI-powered tool.
            2. **Customize Your Grid:** Play around with the settings to create the perfect hexagon grid overlay.
            3. **Download and Share:** Once you're happy with your creation, download it and share it with the world!

            ### Advanced Features
            - **Generative AI Integration:** Utilize models like `black-forest-labs/FLUX.1-dev` and various LoRA weights for generating unique images.
            - **Pre-rendered Maps:** Access a library of pre-rendered hexagon maps for quick and easy customization.
            - **Image Filter [Look-Up Table (LUT)] Application:** Apply filters (LUTs) to your images for color grading and enhancement.
            - **Depth and 3D Model Generation:** Create depth maps and 3D models from your images for enhanced visualization.
            - **Add Margins:** Customize margins around your images for a polished finish.

            Join the hive and start creating with HexaGrid Creator today!
        
            """, elem_classes="intro")
    with gr.Row():
        with gr.Column(scale=2):
            input_image = gr.Image(
                label="Input Image",
                type="filepath",
                interactive=True,
                elem_classes="centered solid imgcontainer",
                key="imgInput",
                image_mode=None,
                format="PNG",
                height=450,
                width=800
            )
            with gr.Accordion("Sketch Pad", open = False, elem_id="sketchpd"):
                with gr.Row():
                    sketch_image = gr.Sketchpad(
                        label="Sketch Image",
                        type="filepath",
                        #invert_colors=True,
                        #sources=['upload','canvas'],
                        #tool=['editor','select','color-sketch'],
                        placeholder="Draw a sketch or upload an image.",
                        interactive=True,
                        elem_classes="centered solid imgcontainer",
                        key="imgSketch",
                        image_mode="RGBA",
                        format="PNG",
                        brush=gr.Brush(),
                        canvas_size=(640,360)
                    )
                with gr.Row():
                    with gr.Column(scale=1):
                        sketch_replace_input_image_button = gr.Button("Replace Input Image with sketch", elem_id="sketch_replace_input_image_button", elem_classes="solid")
                        sketch_alpha_composite_slider = gr.Slider(0,100,50,0.5, label="Sketch Transparancy", elem_id="alpha_composite_slider")
                        btn_sketch_alpha_composite = gr.Button("Overlay Sketch on Input Image", elem_id="btn_sketchninput", elem_classes="solid")
                        gr.Markdown("### Do Not add to image if using a fill model")

        with gr.Column():
            with gr.Accordion("Hex Coloring and Exclusion", open = False):
                with gr.Row():
                    with gr.Column():
                        color_picker = gr.ColorPicker(label="Pick a color to exclude",value="#505050")
                    with gr.Column():
                        filter_color = gr.Checkbox(label="Filter Excluded Colors from Sampling", value=False,)
                        fill_hex = gr.Checkbox(label="Fill Hex with color from Image", value=True)
                exclude_color_button = gr.Button("Exclude Color", elem_id="exlude_color_button", elem_classes="solid")
                color_display = gr.DataFrame(label="List of Excluded RGBA Colors", headers=["R", "G", "B", "A"], elem_id="excluded_colors", type="array", value=build_dataframe(excluded_color_list), interactive=True, elem_classes="solid centered")
                selected_row = gr.Number(0, label="Selected Row", visible=False)
                delete_button = gr.Button("Delete Row", elem_id="delete_exclusion_button", elem_classes="solid")                
            with gr.Accordion("Image Filters", open = False):
                with gr.Row():
                    with gr.Column():
                        lut_filename = gr.Textbox(
                            value="", 
                            label="Look Up Table (LUT) File Name",
                            elem_id="lutFileName")
                    with gr.Column():
                        lut_file = gr.File(
                            value=None,
                            file_count="single",
                            file_types=[".cube"],
                            type="filepath",
                            label="LUT cube File",
                            height=120)
                with gr.Row():
                    apply_lut_button = gr.Button("Apply Filter (LUT)", elem_classes="solid", elem_id="apply_lut_button")
                with gr.Row():
                    lut_example_image = gr.Image(type="pil", label="Filter (LUT) Example Image", value=constants.default_lut_example_img)
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("""
                        ### Included Filters (LUTs)
                        There are several included Filters:

                        Try them on the example image before applying to your Input Image.
                        """, elem_id="lut_markdown")
                    with gr.Column():
                        gr.Examples(elem_id="lut_examples",
                            examples=[[f] for f in constants.lut_files],
                            inputs=[lut_filename],
                            outputs=[lut_filename],
                            label="Select a Filter (LUT) file. Populate the LUT File Name field",
                            examples_per_page = 15,
                        )
                    
                lut_file.change(get_filename, inputs=[lut_file], outputs=[lut_filename])
                lut_filename.change(show_lut, inputs=[lut_filename, lut_example_image], outputs=[lut_example_image], scroll_to_output=True)
                apply_lut_button.click(
                    lambda lut_filename, input_image: gr.Warning("Please upload an Input Image to get started.") if input_image is None else apply_lut_to_image_path(lut_filename, input_image)[1], 
                    inputs=[lut_filename, input_image], 
                    outputs=[input_image], 
                    scroll_to_output=True
                )
            with gr.Accordion("Color Composite", open = False):
                with gr.Row():
                    composite_color = gr.ColorPicker(label="Color", value="#ede9ac44")
                    composite_opacity = gr.Slider(label="Opacity %", minimum=0, maximum=100, value=50, interactive=True)
                with gr.Row():
                    composite_button = gr.Button("Composite", elem_classes="solid")
    with gr.Row():
        with gr.Accordion("Generate AI Image (click here for options)", open = False):
            with gr.Row():
                with gr.Column():
                    # model_options = gr.Dropdown(
                    #     label="Choose an AI Model*",
                    #     choices=constants.MODELS + constants.LORA_WEIGHTS + ["Manual Entry"],
                    #     value="Cossale/Frames2-Flex.1",
                    #     elem_classes="solid"
                    # )
                    model_textbox = gr.Textbox(
                        label="LORA/Model",
                        value="Cossale/Frames2-Flex.1",
                        elem_classes="solid",
                        elem_id="inference_model",
                        lines=2,
                        visible=False
                    )
                    with gr.Accordion("Choose Image Style*", open=True):
                        lora_gallery = gr.Gallery(
                            [(open_image(image_path), title) for image_path, title in lora_models],
                            label="Styles",
                            allow_preview=False, preview=False ,
                            columns=2,
                            elem_id="lora_gallery",
                            show_share_button=False,
                            elem_classes="solid", type="filepath", 
                            object_fit="contain", height="auto", format="png",
                        )
                    # Update map_options to a Dropdown with choices from constants.PROMPTS keys
                    with gr.Row():
                        with gr.Column():
                            map_options = gr.Dropdown(
                            label="Map Options*",
                            choices=list(constants.PROMPTS.keys()),
                            value="Alien Landscape",
                            elem_classes="solid",
                            scale=0
                            )
                            # Add Dropdown for sizing of Images, height and width based on selection. Options are 16x9, 16x10, 4x5, 1x1
                            # The values of height and width are based on common resolutions for each aspect ratio
                            # Default to 16x9, 912x512
                            image_size_ratio = gr.Dropdown(label="Image Aspect Ratio", choices=["16:9", "16:10", "4:5", "4:3", "2:1","3:2","1:1", "9:16", "10:16", "5:4", "3:4","1:2", "2:3"], value="16:9", elem_classes="solid", type="value", scale=0, interactive=True)
                        with gr.Column():
                            seed_slider = gr.Slider(
                                label="Seed",
                                minimum=0,
                                maximum=constants.MAX_SEED,
                                step=1,
                                value=0,
                                scale=0, randomize=True, elem_id="rnd_seed"
                            )
                            randomize_seed = gr.Checkbox(label="Randomize seed", value=False, scale=0, interactive=True)
                    prompt_textbox = gr.Textbox(
                        label="Prompt",
                        visible=False,
                        elem_classes="solid",
                        value="top-down, (rectangular tabletop_map) alien planet map, Battletech_boardgame scifi world with forests, lakes, oceans, continents and snow at the top and bottom, (middle is dark, no_reflections, no_shadows), from directly above. From 100,000 feet looking straight down",
                        lines=4
                    )
                    negative_prompt_textbox = gr.Textbox(
                        label="Negative Prompt",
                        visible=False,
                        elem_classes="solid",
                        value="Earth, low quality, bad anatomy, blurry, cropped, worst quality, shadows, people, humans, reflections, shadows, realistic map of the Earth, isometric, text"
                    )
                    prompt_notes_label = gr.Label(
                        "You should use FRM$ as trigger words. @1.5 minutes",
                        elem_classes="solid centered small",
                        show_label=False,
                        visible=False
                    )
                    # Keep the change event to maintain functionality
                    map_options.change(
                        fn=update_prompt_visibility,
                        inputs=[map_options],
                        outputs=[prompt_textbox, negative_prompt_textbox, prompt_notes_label]
                    )
                    with gr.Row():
                        generate_input_image = gr.Button(
                            "Generate from Input Image & Options ",
                            elem_id="generate_input_image",
                            elem_classes="solid"
                        )
                with gr.Column(scale=2):
                        with gr.Accordion("Template Images", open = False):
                            with gr.Row():
                                with gr.Column(scale=2):
                                    # Gallery from PRE_RENDERED_IMAGES GOES HERE
                                    prerendered_image_gallery = gr.Gallery(label="Image Gallery", show_label=True, value=build_prerendered_images_by_quality(3,'thumbnail'), elem_id="gallery", 
                                        elem_classes="solid", type="filepath", columns=[3], rows=[3], preview=False ,object_fit="contain", height="auto", format="png",allow_preview=False)
                                with gr.Column():                                    
                                    replace_input_image_button = gr.Button(
                                        "Replace Input Image",
                                        elem_id="prerendered_replace_input_image_button",
                                        elem_classes="solid"
                                    )
                                    generate_input_image_from_gallery = gr.Button(
                                        "Generate AI Image from Template Image & Options",
                                        elem_id="generate_input_image_from_gallery",
                                        elem_classes="solid"
                                    )
                        image_guidance_stength = gr.Slider(label="Image Guidance Strength (prompt percentage)", minimum=0, maximum=1.0, value=0.85, step=0.01, interactive=True)

    with gr.Accordion("Advanced Hexagon Settings", open = False):
        with gr.Row():
            start_x = gr.Number(label="Start X", value=20, minimum=-512, maximum= 512, precision=0)
            start_y = gr.Number(label="Start Y", value=20, minimum=-512, maximum= 512, precision=0)
            end_x = gr.Number(label="End X", value=-20, minimum=-512, maximum= 512, precision=0)
            end_y = gr.Number(label="End Y", value=-20, minimum=-512, maximum= 512, precision=0)
        with gr.Row():
            x_spacing = gr.Number(label="Adjust Horizontal spacing", value=-8, minimum=-200, maximum=200, precision=1)
            y_spacing = gr.Number(label="Adjust Vertical spacing", value=3, minimum=-200, maximum=200, precision=1)
        with gr.Row():
            rotation = gr.Slider(-90, 180, 0.0, 0.1, label="Hexagon Rotation (degree)")
            add_hex_text = gr.Dropdown(label="Add Text to Hexagons", choices=[None, "Column-Row Coordinates", "Column(Letter)-Row Coordinates", "Column-Row(Letter) Coordinates", "Sequential Numbers", "Playing Cards Sequential", "Playing Cards Alternate Red and Black", "Custom List"], value=None)
        with gr.Row():
            custom_text_list = gr.TextArea(label="Custom Text List", value=constants.cards_alternating, visible=False,)
            custom_text_color_list = gr.TextArea(label="Custom Text Color List", value=constants.card_colors_alternating, visible=False)
        with gr.Row():
            hex_text_info = gr.Markdown("""
            ### Text Color uses the Border Color and Border Opacity, unless you use a custom list.
            ### The Custom Text List and Custom Text Color List are repeating comma separated lists.
            ### The custom color list is a comma separated list of hex colors.
            #### Example: "A,2,3,4,5,6,7,8,9,10,J,Q,K", "red,#0000FF,#00FF00,red,#FFFF00,#00FFFF,#FF8000,#FF00FF,#FF0080,#FF8000,#FF0080,lightblue"
            """, elem_id="hex_text_info", visible=False)
        add_hex_text.change(
        fn=lambda x: (
            gr.update(visible=(x == "Custom List")), 
            gr.update(visible=(x == "Custom List")),
            gr.update(visible=(x != None))
            ),
            inputs=add_hex_text,
            outputs=[custom_text_list, custom_text_color_list, hex_text_info]
        )
    with gr.Row():
        hex_size = gr.Number(label="Hexagon Size", value=90, minimum=1, maximum=768)
        border_size = gr.Slider(-5,25,value=2,step=1,label="Border Size")
    with gr.Row():        
        background_color = gr.ColorPicker(label="Background Color", value="#000000", interactive=True)
        background_opacity = gr.Slider(0,100,0,1,label="Background Opacity %")
        border_color = gr.ColorPicker(label="Border Color", value="#7b7b7b", interactive=True)
        border_opacity = gr.Slider(0,100,50,1,label="Border Opacity %")
    with gr.Row():
        hex_button = gr.Button("Generate Hex Grid!", elem_classes="solid", elem_id="btn-generate")
    with gr.Row():
        output_image = gr.Image(label="Hexagon Grid Image", image_mode = "RGBA", elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgOutput",interactive=True)
        overlay_image = gr.Image(label="Hexagon Overlay Image", image_mode = "RGBA", elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgOverlay",interactive=True)
    with gr.Row():
        output_overlay_composite = gr.Slider(0,100,50,0.5, label="Interpolate Intensity")
        output_blend_multiply_composite = gr.Slider(0,100,50,0.5, label="Overlay Intensity")
        output_alpha_composite = gr.Slider(0,100,50,0.5, label="Alpha Composite Intensity")
    with gr.Accordion("Add Margins (bleed)", open=False):
        with gr.Row():
            border_image_source = gr.Radio(label="Add Margins around which Image", choices=["Input Image", "Overlay Image"], value="Overlay Image")
        with gr.Row():
            mask_width = gr.Number(label="Margins Width", value=10, minimum=0, maximum=100, precision=0)
            mask_height = gr.Number(label="Margins Height", value=10, minimum=0, maximum=100, precision=0)
        with gr.Row():
                margin_color = gr.ColorPicker(label="Margin Color", value="#333333FF", interactive=True)
                margin_opacity = gr.Slider(0,100,95,0.5,label="Margin Opacity %")
        with gr.Row():
            add_border_button = gr.Button("Add Margins", elem_classes="solid", variant="secondary")
        with gr.Row():
            bordered_image_output = gr.Image(label="Image with Margins", image_mode="RGBA",  elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgBordered",interactive=False, show_download_button=True, show_fullscreen_button=True, show_share_button=True)
    accordian_3d = gr.Accordion("Height Maps and 3D", open=False, elem_id="accordian_3d")
    with accordian_3d:
        with gr.Row():
            depth_image_source = gr.Radio(
                label="Depth Image Source",
                choices=["Input Image", "Hexagon Grid Image", "Overlay Image", "Image with Margins"],
                value="Input Image"
            )
        with gr.Accordion("Advanced 3D Generation Settings", open=False):
            with gr.Row():
                with gr.Column():
                    # Use standard seed settings only
                    seed_3d = gr.Slider(0, constants.MAX_SEED, label="Seed (3D Generation)", value=0, step=1, randomize=True)
                    randomize_seed_3d = gr.Checkbox(label="Randomize Seed (3D Generation)", value=True)
                with gr.Column():
                    steps = gr.Slider(6, 36, value=12, step=1, label="Image Sampling Steps", interactive=True)
                    video_resolution = gr.Slider(384, 768, value=480, step=32, label="Video Resolution (*danger*)", interactive=True)
                    model_resolution = gr.Slider(512, 2304, value=1024, step=64, label="3D Model Resolution", interactive=True)
        with gr.Row():
            generate_3d_asset_button = gr.Button("Generate 3D Asset", elem_classes="solid", variant="secondary", interactive=False)
        with gr.Row():
            depth_output = gr.Image(label="Depth Map", image_mode="L", elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="DepthOutput",interactive=False, show_download_button=True, show_fullscreen_button=True, show_share_button=True, height=400)
        with gr.Row():
            # For display: video output and 3D model preview (GLTF)
            video_output = gr.Video(label="3D Asset Video", autoplay=True, loop=True, height=400)
        with gr.Accordion("GLB Extraction Settings", open=False):
            with gr.Row():
                mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
            with gr.Row():
                extract_glb_btn = gr.Button("Extract GLB", interactive=False)
                extract_gaussian_btn = gr.Button("Extract Gaussian", interactive=False)
            with gr.Row():
                with gr.Column(scale=2):
                    model_output = gr.Model3D(label="Extracted 3D Model", clear_color=[1.0, 1.0, 1.0, 1.0], 
                                            elem_classes="centered solid imgcontainer", interactive=True)
                with gr.Column(scale=1):
                    glb_file = gr.File(label="3D GLTF", elem_classes="solid small centered", height=250)
                    gaussian_file = gr.File(label="Gaussian", elem_classes="solid small centered", height=250)
                    gr.Markdown("""
                    ### Files over 10 MB may not display in the 3D model viewer
                    """, elem_id="file_size_info", elem_classes="intro" )

    is_multiimage = gr.State(False)
    output_buf = gr.State()
    ddd_image_path = gr.State("./images/images/Beeuty-1.png")
    ddd_file_name = gr.State("Hexagon_file")
    with gr.Row():
        gr.Examples(examples=[
            ["assets//examples//hex_map_p1.png", False, True, -32,-31,80,80,-1.8,0,35,0,1,"#FFD0D0", 15],
            ["assets//examples//hex_map_p1_overlayed.png", False, False, -32,-31,80,80,-1.8,0,35,0,1,"#FFD0D0", 75],
            ["assets//examples//hex_flower_logo.png", False, True, -95,-95,100,100,-24,-2,190,30,2,"#FF8951", 50],
            ["assets//examples//hexed_fract_1.png", False, True, 0,0,0,0,0,0,10,0,0,"#000000", 5],
            ["assets//examples//tmpzt3mblvk.png", False, True, -20,10,0,0,-6,-2,35,30,1,"#ffffff", 0],
            ],
            inputs=[input_image, filter_color, fill_hex, start_x, start_y, end_x, end_y, x_spacing, y_spacing, hex_size, rotation, border_size, border_color, border_opacity],
            elem_id="examples")
    # with gr.Row():
    #     login_button = gr.LoginButton(size="sm", elem_classes="solid centered", elem_id="hf_login_btn")
    with gr.Row():
        gr.HTML(value=getVersions(), visible=True, elem_id="versions")

    # Handlers
    hexaGrid.load(start_session)
    hexaGrid.unload(end_session)

    color_display.select(on_color_display_select,inputs=[color_display], outputs=[selected_row])
    color_display.input(on_input,inputs=[color_display], outputs=[color_display, gr.State(excluded_color_list)])

    delete_button.click(fn=delete_color, inputs=[selected_row, color_display], outputs=[color_display])
    exclude_color_button.click(fn=add_color, inputs=[color_picker, gr.State(excluded_color_list)], outputs=[color_display, gr.State(excluded_color_list)])
    hex_button.click(
        fn=lambda hex_size, border_size, input_image, start_x, start_y, end_x, end_y, rotation, background_color, background_opacity, border_color, border_opacity, fill_hex, color_display, filter_color, x_spacing, y_spacing, add_hex_text, custom_text_list, custom_text_color_list: 
            gr.Warning("Please upload an Input Image to get started.") if input_image is None else hex_create(hex_size, border_size, input_image, start_x, start_y, end_x, end_y, rotation, background_color, background_opacity, border_color, border_opacity, fill_hex, color_display, filter_color, x_spacing, y_spacing, add_hex_text, custom_text_list, custom_text_color_list),
        inputs=[hex_size, border_size, input_image, start_x, start_y, end_x, end_y, rotation, background_color, background_opacity, border_color, border_opacity, fill_hex, color_display, filter_color, x_spacing, y_spacing, add_hex_text, custom_text_list, custom_text_color_list],
        outputs=[output_image, overlay_image],
        scroll_to_output=True
    )
    generate_input_image.click(
        fn=unload_3d_models,
        trigger_mode="always_last",
        outputs=[generate_3d_asset_button]
    ).then(
        fn=generate_input_image_click,
        inputs=[input_image,map_options, prompt_textbox, negative_prompt_textbox, model_textbox, randomize_seed, seed_slider, gr.State(False), sketch_image, image_guidance_stength, image_size_ratio],
        outputs=[input_image, seed_slider], scroll_to_output=True
    ).then(
        fn=update_sketch_dimensions,
        inputs=[input_image, sketch_image],
        outputs=[sketch_image, sketch_image]
    )
    input_image.input(
        fn=on_input_image_change,
        inputs=[input_image],
        outputs=[input_image,sketch_image], scroll_to_output=True,
    )
    ###################### sketchpad ############################
    btn_sketch_alpha_composite.click(
        fn=composite_with_control_sync,
        inputs=[input_image, sketch_image, sketch_alpha_composite_slider],
        outputs=[input_image],
        scroll_to_output=True
    )
    sketch_replace_input_image_button.click(
        lambda sketch_image: replace_input_with_sketch_image(sketch_image),
        inputs=[sketch_image],
        outputs=[input_image], scroll_to_output=True
    )
    ##################### model #######################################
    model_textbox.change(
        fn=update_prompt_notes,
        inputs=model_textbox,
        outputs=prompt_notes_label,preprocess=False
    )
    # model_options.change(
    #     fn=lambda x: (gr.update(visible=(x == "Manual Entry")), gr.update(value=x) if x != "Manual Entry" else gr.update()),
    #     inputs=model_options,
    #     outputs=[model_textbox, model_textbox]
    # )
    # model_options.change(
    #     fn=update_prompt_notes,
    #     inputs=model_options,
    #     outputs=prompt_notes_label
    # )
    lora_gallery.select(
        fn=update_selection,
        inputs=[image_size_ratio],
        outputs=[model_textbox, model_textbox, gr.State(selected_index), image_size_ratio, prompt_notes_label]
    )

    #################### model end ########################################

    composite_button.click(
        fn=lambda input_image, composite_color, composite_opacity: gr.Warning("Please upload an Input Image to get started.") if input_image is None else change_color(input_image, composite_color, composite_opacity),
        inputs=[input_image, composite_color, composite_opacity],
        outputs=[input_image]
    )

    #use conditioned_image as the input_image for generate_input_image_click
    generate_input_image_from_gallery.click(
        fn=unload_3d_models,
        trigger_mode="always_last",
        outputs=[generate_3d_asset_button]
    ).then(
        fn=generate_input_image_click,
        inputs=[input_image, map_options, prompt_textbox, negative_prompt_textbox, model_textbox,randomize_seed, seed_slider, gr.State(True), sketch_image , image_guidance_stength, image_size_ratio],
        outputs=[input_image, seed_slider], scroll_to_output=True
    ).then(
        fn=update_sketch_dimensions,
        inputs=[input_image, sketch_image],
        outputs=[sketch_image, sketch_image]
    )

    # Update the state variable with the prerendered image filepath when an image is selected
    prerendered_image_gallery.select(
        fn=on_prerendered_gallery_selection, 
        inputs=None, 
        outputs=[gr.State(current_prerendered_image)],  # Update the state with the selected image
        show_api=False
    )
    # replace input image with selected gallery image
    replace_input_image_button.click(
        lambda: current_prerendered_image.value,
        inputs=None,
        outputs=[input_image], scroll_to_output=True
    ).then(
        fn=update_sketch_dimensions,
        inputs=[input_image, sketch_image],
        outputs=[sketch_image, sketch_image]
    )
    output_overlay_composite.change(
        fn=combine_images_with_lerp,
        inputs=[input_image, output_image, output_overlay_composite],
        outputs=[overlay_image], scroll_to_output=True
    )
    output_blend_multiply_composite.change(
        fn=multiply_and_blend_images,
        inputs=[input_image, output_image, output_blend_multiply_composite],
        outputs=[overlay_image],
        scroll_to_output=True
    )
    output_alpha_composite.change(
        fn=alpha_composite_with_control,
        inputs=[input_image, output_image, output_alpha_composite],
        outputs=[overlay_image],
        scroll_to_output=True
    )
    add_border_button.click(
        fn=lambda image_source, mask_w, mask_h, color, opacity, input_img, overlay_img: add_border(input_img if image_source == "Input Image" else overlay_img, mask_w, mask_h, update_color_opacity(detect_color_format(color), opacity * 2.55)),
        inputs=[border_image_source, mask_width, mask_height, margin_color, margin_opacity, input_image, overlay_image],
        outputs=[bordered_image_output],
        scroll_to_output=True
    )
    # 3D Generation

    # generate_depth_button.click(
    #     fn=generate_depth_button_click,
    #     inputs=[depth_image_source, resized_width_slider, z_scale_slider, input_image, output_image, overlay_image, bordered_image_output],
    #     outputs=[depth_map_output, model_output, model_file], scroll_to_output=True
    # )
    accordian_3d.expand(
        fn=load_trellis_model,
        trigger_mode="always_last"
    ).then(
        fn=load_3d_models,
        trigger_mode="always_last",
        outputs=[generate_3d_asset_button],
        show_api=False
    )
    accordian_3d.collapse(
        fn=unload_3d_models,
        trigger_mode="always_last",
        outputs=[generate_3d_asset_button],
        show_api=False
    )
    # Chain the buttons
    generate_3d_asset_button.click(
        fn=generate_3d_asset_part1,
        inputs=[depth_image_source, randomize_seed_3d, seed_3d, input_image, output_image, overlay_image, bordered_image_output],
        outputs=[depth_output, ddd_image_path, ddd_file_name, seed_3d ],
        scroll_to_output=True
    ).then(
        fn=generate_3d_asset_part2,
        inputs=[depth_output, ddd_image_path, ddd_file_name, seed_3d, steps, model_resolution, video_resolution ],
        outputs=[output_buf, video_output, depth_output],
        scroll_to_output=True
    ).then(
        lambda: (gr.Button(interactive=True), gr.Button(interactive=True)),
        outputs=[extract_glb_btn, extract_gaussian_btn]
    )

    # Extraction callbacks remain unchanged from previous behavior
    extract_glb_btn.click(
        fn=extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size],
        outputs=[model_output, glb_file]
    ).then(
        lambda: gr.Button(interactive=True),
        outputs=[glb_file]
    )

    extract_gaussian_btn.click(
        fn=extract_gaussian,
        inputs=[output_buf],
        outputs=[model_output, gaussian_file]
    ).then(
        lambda: gr.Button(interactive=True),
        outputs=[gaussian_file]
    )

if __name__ == "__main__":
    constants.load_env_vars(constants.dotenv_path)
    logging.basicConfig(
        format="[%(levelname)s] %(asctime)s %(message)s", level=logging.INFO
    )
    logging.info("Environment Variables: %s" % os.environ)
    # if _get_output(["nvcc", "--version"]) is None:
    #     logging.info("Installing CUDA toolkit...")
    #     install_cuda_toolkit()
    # else:
    #     logging.info("Detected CUDA: %s" % _get_output(["nvcc", "--version"]))

    # logging.info("Installing CUDA extensions...")
    # setup_runtime_env()
    #main(os.getenv("DEBUG") == "1")
    #main()


    #-------------- ------------------------------------------------MODEL INITIALIZATION------------------------------------------------------------#
    # Load models once during module import
    dtype = torch.bfloat16
    device = "cuda" if torch.cuda.is_available() else "cpu"
    base_model = "black-forest-labs/FLUX.1-dev"
    good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
    #pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=good_vae).to(device)
    #pipe.enable_model_cpu_offload()
    #pipe.vae.enable_slicing()
    #pipe.attn_implementation="flash_attention_2"

    # image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
    # depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large", ignore_mismatched_sizes=True)
    # if constants.IS_SHARED_SPACE:
    #     TRELLIS_PIPELINE = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
    #     TRELLIS_PIPELINE.to(device)
    #     try:
    #        TRELLIS_PIPELINE.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
    #     except:
    #         pass
    hexaGrid.queue(default_concurrency_limit=1,max_size=12,api_open=False)
    hexaGrid.launch(allowed_paths=["assets","/","./assets","images","./images", "./images/prerendered", 'e:/TMP'], favicon_path="./assets/favicon.ico", max_file_size="10mb")