final_demo.py

import os
import math
import gradio as gr
import numpy as np
import torch
import safetensors.torch as sf
import db_examples
import datetime
from pathlib import Path
from io import BytesIO

from PIL import Image
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler
from diffusers.models.attention_processor import AttnProcessor2_0
from transformers import CLIPTextModel, CLIPTokenizer
from briarmbg import BriaRMBG
from enum import Enum
from torch.hub import download_url_to_file

from torch.hub import download_url_to_file
import cv2

from typing import Optional

from Depth.depth_anything_v2.dpt import DepthAnythingV2



# from FLORENCE
import spaces
import supervision as sv
import torch
from PIL import Image

from utils.sam import load_sam_image_model, run_sam_inference


try:
    import xformers
    import xformers.ops
    XFORMERS_AVAILABLE = True
    print("xformers is available - Using memory efficient attention")
except ImportError:
    XFORMERS_AVAILABLE = False
    print("xformers not available - Using default attention")

# Memory optimizations for RTX 2070
torch.backends.cudnn.benchmark = True
if torch.cuda.is_available():
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    # Set a smaller attention slice size for RTX 2070
    torch.backends.cuda.max_split_size_mb = 512
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

# 'stablediffusionapi/realistic-vision-v51'
# 'runwayml/stable-diffusion-v1-5'
sd15_name = 'stablediffusionapi/realistic-vision-v51'
tokenizer = CLIPTokenizer.from_pretrained(sd15_name, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(sd15_name, subfolder="text_encoder")
vae = AutoencoderKL.from_pretrained(sd15_name, subfolder="vae")
# unet = None
unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")
rmbg = BriaRMBG.from_pretrained("briaai/RMBG-1.4")

model = DepthAnythingV2(encoder='vits', features=64, out_channels=[48, 96, 192, 384])
model.load_state_dict(torch.load('checkpoints/depth_anything_v2_vits.pth', map_location=device))
model = model.to(device)
model.eval()

def initialize_pipeline_fc():
    global unet, unet_original_forward, t2i_pipe, i2i_pipe, current_model
    # Clear previous models
    if unet is not None:
        del unet
        unet = None
    if t2i_pipe is not None:
        del t2i_pipe
        t2i_pipe = None
    if i2i_pipe is not None:
        del i2i_pipe
        i2i_pipe = None
    clear_memory()

    # Load UNet model
    unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")

    # Modify UNet to have 8 input channels
    with torch.no_grad():
        new_conv_in = torch.nn.Conv2d(8, unet.conv_in.out_channels,
                                      unet.conv_in.kernel_size, unet.conv_in.stride,
                                      unet.conv_in.padding)
        new_conv_in.weight.zero_()
        new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
        new_conv_in.bias = unet.conv_in.bias
        unet.conv_in = new_conv_in

    # Load model weights
    sd_offset = sf.load_file(model_path_fc)
    sd_origin = unet.state_dict()
    sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
    unet.load_state_dict(sd_merged, strict=True)
    del sd_offset, sd_origin, sd_merged

    # Set up forward function
    unet_original_forward = unet.forward
    unet.forward = hooked_unet_forward

    # Move UNet to device
    unet = unet.to(device=device, dtype=dtype)
    enable_efficient_attention()

    # Create pipelines
    t2i_pipe = StableDiffusionPipeline(
        vae=vae,
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        unet=unet,
        scheduler=dpmpp_2m_sde_karras_scheduler,
        safety_checker=None,
        feature_extractor=None
    )
    i2i_pipe = StableDiffusionImg2ImgPipeline(
        vae=vae,
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        unet=unet,
        scheduler=dpmpp_2m_sde_karras_scheduler,
        safety_checker=None,
        feature_extractor=None
    )

    current_model = 'fc'

def initialize_pipeline_fbc():
    global unet, unet_original_forward, t2i_pipe, i2i_pipe, current_model
    # Clear previous models
    if unet is not None:
        del unet
        unet = None
    if t2i_pipe is not None:
        del t2i_pipe
        t2i_pipe = None
    if i2i_pipe is not None:
        del i2i_pipe
        i2i_pipe = None
    clear_memory()

    # Load UNet model
    unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")

    # Modify UNet to have 12 input channels
    with torch.no_grad():
        new_conv_in = torch.nn.Conv2d(12, unet.conv_in.out_channels,
                                      unet.conv_in.kernel_size, unet.conv_in.stride,
                                      unet.conv_in.padding)
        new_conv_in.weight.zero_()
        new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
        new_conv_in.bias = unet.conv_in.bias
        unet.conv_in = new_conv_in

    # Load model weights
    sd_offset = sf.load_file(model_path_fbc)
    sd_origin = unet.state_dict()
    sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
    unet.load_state_dict(sd_merged, strict=True)
    del sd_offset, sd_origin, sd_merged

    # Set up forward function
    unet_original_forward = unet.forward
    unet.forward = hooked_unet_forward

    # Move UNet to device
    unet = unet.to(device=device, dtype=dtype)
    enable_efficient_attention()

    # Create pipelines
    t2i_pipe = StableDiffusionPipeline(
        vae=vae,
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        unet=unet,
        scheduler=dpmpp_2m_sde_karras_scheduler,
        safety_checker=None,
        feature_extractor=None
    )
    i2i_pipe = StableDiffusionImg2ImgPipeline(
        vae=vae,
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        unet=unet,
        scheduler=dpmpp_2m_sde_karras_scheduler,
        safety_checker=None,
        feature_extractor=None
    )

    current_model = 'fbc'

# Change UNet

# with torch.no_grad():
#     new_conv_in = torch.nn.Conv2d(8, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding)
#     new_conv_in.weight.zero_()
#     new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
#     new_conv_in.bias = unet.conv_in.bias
#     unet.conv_in = new_conv_in


# unet_original_forward = unet.forward


def enable_efficient_attention():
    if XFORMERS_AVAILABLE:
        try:
            # RTX 2070 specific settings
            unet.set_use_memory_efficient_attention_xformers(True)
            vae.set_use_memory_efficient_attention_xformers(True)
            print("Enabled xformers memory efficient attention")
        except Exception as e:
            print(f"Xformers error: {e}")
            print("Falling back to sliced attention")
            # Use sliced attention for RTX 2070
            unet.set_attention_slice_size(4)
            vae.set_attention_slice_size(4)
            unet.set_attn_processor(AttnProcessor2_0())
            vae.set_attn_processor(AttnProcessor2_0())
    else:
        # Fallback for when xformers is not available
        print("Using sliced attention")
        unet.set_attention_slice_size(4)
        vae.set_attention_slice_size(4)
        unet.set_attn_processor(AttnProcessor2_0())
        vae.set_attn_processor(AttnProcessor2_0())

# Add memory clearing function
def clear_memory():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

# Enable efficient attention
enable_efficient_attention()


def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs):
    c_concat = kwargs['cross_attention_kwargs']['concat_conds'].to(sample)
    c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0)
    new_sample = torch.cat([sample, c_concat], dim=1)

    # Pad new_sample to have 12 channels
    if new_sample.shape[1] < 12:
        padding_channels = 12 - new_sample.shape[1]
        padding = torch.zeros(
            new_sample.shape[0],
            padding_channels,
            new_sample.shape[2],
            new_sample.shape[3],
            device=new_sample.device,
            dtype=new_sample.dtype
        )
        new_sample = torch.cat([new_sample, padding], dim=1)

    kwargs['cross_attention_kwargs'] = {}
    return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)


unet.forward = hooked_unet_forward

# Load
model_path_fc = './models/iclight_sd15_fc.safetensors'   # Foreground Conditioned model
model_path_fbc = './models/iclight_sd15_fbc.safetensors' # Foreground and Background Conditioned model



# if not os.path.exists(model_path):
#     download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fc.safetensors', dst=model_path)

sd_offset = sf.load_file(model_path_fc)
sd_origin = unet.state_dict()
keys = sd_origin.keys()
sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
unet.load_state_dict(sd_merged, strict=True)
del sd_offset, sd_origin, sd_merged, keys

# Device

# device = torch.device('cuda')
# text_encoder = text_encoder.to(device=device, dtype=torch.float16)
# vae = vae.to(device=device, dtype=torch.bfloat16)
# unet = unet.to(device=device, dtype=torch.float16)
# rmbg = rmbg.to(device=device, dtype=torch.float32)


# Device and dtype setup
device = torch.device('cuda')
dtype = torch.float16  # RTX 2070 works well with float16

# Memory optimizations for RTX 2070
torch.backends.cudnn.benchmark = True
if torch.cuda.is_available():
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    # Set a very small attention slice size for RTX 2070 to avoid OOM
    torch.backends.cuda.max_split_size_mb = 128

# Move models to device with consistent dtype
text_encoder = text_encoder.to(device=device, dtype=dtype)
vae = vae.to(device=device, dtype=dtype)  # Changed from bfloat16 to float16
unet = unet.to(device=device, dtype=dtype)
rmbg = rmbg.to(device=device, dtype=torch.float32)  # Keep this as float32


ddim_scheduler = DDIMScheduler(
    num_train_timesteps=1000,
    beta_start=0.00085,
    beta_end=0.012,
    beta_schedule="scaled_linear",
    clip_sample=False,
    set_alpha_to_one=False,
    steps_offset=1,
)

euler_a_scheduler = EulerAncestralDiscreteScheduler(
    num_train_timesteps=1000,
    beta_start=0.00085,
    beta_end=0.012,
    steps_offset=1
)

dpmpp_2m_sde_karras_scheduler = DPMSolverMultistepScheduler(
    num_train_timesteps=1000,
    beta_start=0.00085,
    beta_end=0.012,
    algorithm_type="sde-dpmsolver++",
    use_karras_sigmas=True,
    steps_offset=1
)

# Pipelines

t2i_pipe = StableDiffusionPipeline(
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    unet=unet,
    scheduler=dpmpp_2m_sde_karras_scheduler,
    safety_checker=None,
    requires_safety_checker=False,
    feature_extractor=None,
    image_encoder=None
)

i2i_pipe = StableDiffusionImg2ImgPipeline(
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    unet=unet,
    scheduler=dpmpp_2m_sde_karras_scheduler,
    safety_checker=None,
    requires_safety_checker=False,
    feature_extractor=None,
    image_encoder=None
)


@torch.inference_mode()
def encode_prompt_inner(txt: str):
    max_length = tokenizer.model_max_length
    chunk_length = tokenizer.model_max_length - 2
    id_start = tokenizer.bos_token_id
    id_end = tokenizer.eos_token_id
    id_pad = id_end

    def pad(x, p, i):
        return x[:i] if len(x) >= i else x + [p] * (i - len(x))

    tokens = tokenizer(txt, truncation=False, add_special_tokens=False)["input_ids"]
    chunks = [[id_start] + tokens[i: i + chunk_length] + [id_end] for i in range(0, len(tokens), chunk_length)]
    chunks = [pad(ck, id_pad, max_length) for ck in chunks]

    token_ids = torch.tensor(chunks).to(device=device, dtype=torch.int64)
    conds = text_encoder(token_ids).last_hidden_state

    return conds


@torch.inference_mode()
def encode_prompt_pair(positive_prompt, negative_prompt):
    c = encode_prompt_inner(positive_prompt)
    uc = encode_prompt_inner(negative_prompt)

    c_len = float(len(c))
    uc_len = float(len(uc))
    max_count = max(c_len, uc_len)
    c_repeat = int(math.ceil(max_count / c_len))
    uc_repeat = int(math.ceil(max_count / uc_len))
    max_chunk = max(len(c), len(uc))

    c = torch.cat([c] * c_repeat, dim=0)[:max_chunk]
    uc = torch.cat([uc] * uc_repeat, dim=0)[:max_chunk]

    c = torch.cat([p[None, ...] for p in c], dim=1)
    uc = torch.cat([p[None, ...] for p in uc], dim=1)

    return c, uc


@torch.inference_mode()
def pytorch2numpy(imgs, quant=True):
    results = []
    for x in imgs:
        y = x.movedim(0, -1)

        if quant:
            y = y * 127.5 + 127.5
            y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
        else:
            y = y * 0.5 + 0.5
            y = y.detach().float().cpu().numpy().clip(0, 1).astype(np.float32)

        results.append(y)
    return results


@torch.inference_mode()
def numpy2pytorch(imgs):
    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0  # so that 127 must be strictly 0.0
    h = h.movedim(-1, 1)
    return h


def resize_and_center_crop(image, target_width, target_height):
    pil_image = Image.fromarray(image)
    original_width, original_height = pil_image.size
    scale_factor = max(target_width / original_width, target_height / original_height)
    resized_width = int(round(original_width * scale_factor))
    resized_height = int(round(original_height * scale_factor))
    resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
    left = (resized_width - target_width) / 2
    top = (resized_height - target_height) / 2
    right = (resized_width + target_width) / 2
    bottom = (resized_height + target_height) / 2
    cropped_image = resized_image.crop((left, top, right, bottom))
    return np.array(cropped_image)


def resize_without_crop(image, target_width, target_height):
    pil_image = Image.fromarray(image)
    resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
    return np.array(resized_image)


@torch.inference_mode()
def run_rmbg(img, sigma=0.0):
    # Convert RGBA to RGB if needed
    if img.shape[-1] == 4:
        # Use white background for alpha composition
        alpha = img[..., 3:] / 255.0
        rgb = img[..., :3]
        white_bg = np.ones_like(rgb) * 255
        img = (rgb * alpha + white_bg * (1 - alpha)).astype(np.uint8)
    
    H, W, C = img.shape
    assert C == 3
    k = (256.0 / float(H * W)) ** 0.5
    feed = resize_without_crop(img, int(64 * round(W * k)), int(64 * round(H * k)))
    feed = numpy2pytorch([feed]).to(device=device, dtype=torch.float32)
    alpha = rmbg(feed)[0][0]
    alpha = torch.nn.functional.interpolate(alpha, size=(H, W), mode="bilinear")
    alpha = alpha.movedim(1, -1)[0]
    alpha = alpha.detach().float().cpu().numpy().clip(0, 1)
    
    # Create RGBA image
    rgba = np.dstack((img, alpha * 255)).astype(np.uint8)
    result = 127 + (img.astype(np.float32) - 127 + sigma) * alpha
    return result.clip(0, 255).astype(np.uint8), rgba
@torch.inference_mode()
def process(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):
    clear_memory()

    # Get input dimensions
    input_height, input_width = input_fg.shape[:2]

    bg_source = BGSource(bg_source)


    if bg_source == BGSource.UPLOAD:
        pass
    elif bg_source == BGSource.UPLOAD_FLIP:
        input_bg = np.fliplr(input_bg)
    elif bg_source == BGSource.GREY:
        input_bg = np.zeros(shape=(input_height, input_width, 3), dtype=np.uint8) + 64
    elif bg_source == BGSource.LEFT:
        gradient = np.linspace(255, 0, input_width)
        image = np.tile(gradient, (input_height, 1))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.RIGHT:
        gradient = np.linspace(0, 255, input_width)
        image = np.tile(gradient, (input_height, 1))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.TOP:
        gradient = np.linspace(255, 0, input_height)[:, None]
        image = np.tile(gradient, (1, input_width))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.BOTTOM:
        gradient = np.linspace(0, 255, input_height)[:, None]
        image = np.tile(gradient, (1, input_width))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    else:
        raise 'Wrong initial latent!'

    rng = torch.Generator(device=device).manual_seed(int(seed))

    # Use input dimensions directly
    fg = resize_without_crop(input_fg, input_width, input_height)

    concat_conds = numpy2pytorch([fg]).to(device=vae.device, dtype=vae.dtype)
    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor
    conds, unconds = encode_prompt_pair(positive_prompt=prompt + ', ' + a_prompt, negative_prompt=n_prompt)

    if input_bg is None:
        latents = t2i_pipe(
            prompt_embeds=conds,
            negative_prompt_embeds=unconds,
            width=input_width,
            height=input_height,
            num_inference_steps=steps,
            num_images_per_prompt=num_samples,
            generator=rng,
            output_type='latent',
            guidance_scale=cfg,
            cross_attention_kwargs={'concat_conds': concat_conds},
        ).images.to(vae.dtype) / vae.config.scaling_factor
    else:
        bg = resize_without_crop(input_bg, input_width, input_height)
        bg_latent = numpy2pytorch([bg]).to(device=vae.device, dtype=vae.dtype)
        bg_latent = vae.encode(bg_latent).latent_dist.mode() * vae.config.scaling_factor
        latents = i2i_pipe(
            image=bg_latent,
            strength=lowres_denoise,
            prompt_embeds=conds,
            negative_prompt_embeds=unconds,
            width=input_width,
            height=input_height,
            num_inference_steps=int(round(steps / lowres_denoise)),
            num_images_per_prompt=num_samples,
            generator=rng,
            output_type='latent',
            guidance_scale=cfg,
            cross_attention_kwargs={'concat_conds': concat_conds},
        ).images.to(vae.dtype) / vae.config.scaling_factor

    pixels = vae.decode(latents).sample
    pixels = pytorch2numpy(pixels)
    pixels = [resize_without_crop(
        image=p,
        target_width=int(round(input_width * highres_scale / 64.0) * 64),
        target_height=int(round(input_height * highres_scale / 64.0) * 64))
    for p in pixels]

    pixels = numpy2pytorch(pixels).to(device=vae.device, dtype=vae.dtype)
    latents = vae.encode(pixels).latent_dist.mode() * vae.config.scaling_factor
    latents = latents.to(device=unet.device, dtype=unet.dtype)

    highres_height, highres_width = latents.shape[2] * 8, latents.shape[3] * 8

    fg = resize_without_crop(input_fg, highres_width, highres_height)
    concat_conds = numpy2pytorch([fg]).to(device=vae.device, dtype=vae.dtype)
    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor

    latents = i2i_pipe(
        image=latents,
        strength=highres_denoise,
        prompt_embeds=conds,
        negative_prompt_embeds=unconds,
        width=highres_width,
        height=highres_height,
        num_inference_steps=int(round(steps / highres_denoise)),
        num_images_per_prompt=num_samples,
        generator=rng,
        output_type='latent',
        guidance_scale=cfg,
        cross_attention_kwargs={'concat_conds': concat_conds},
    ).images.to(vae.dtype) / vae.config.scaling_factor

    pixels = vae.decode(latents).sample
    pixels = pytorch2numpy(pixels)
    
    # Resize back to input dimensions
    pixels = [resize_without_crop(p, input_width, input_height) for p in pixels]
    pixels = np.stack(pixels)

    return pixels

@torch.inference_mode()
def process_bg(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source):
    clear_memory()
    bg_source = BGSource(bg_source)

    if bg_source == BGSource.UPLOAD:
        pass
    elif bg_source == BGSource.UPLOAD_FLIP:
        input_bg = np.fliplr(input_bg)
    elif bg_source == BGSource.GREY:
        input_bg = np.zeros(shape=(image_height, image_width, 3), dtype=np.uint8) + 64
    elif bg_source == BGSource.LEFT:
        gradient = np.linspace(224, 32, image_width)
        image = np.tile(gradient, (image_height, 1))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.RIGHT:
        gradient = np.linspace(32, 224, image_width)
        image = np.tile(gradient, (image_height, 1))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.TOP:
        gradient = np.linspace(224, 32, image_height)[:, None]
        image = np.tile(gradient, (1, image_width))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.BOTTOM:
        gradient = np.linspace(32, 224, image_height)[:, None]
        image = np.tile(gradient, (1, image_width))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    else:
        raise 'Wrong background source!'

    rng = torch.Generator(device=device).manual_seed(seed)

    fg = resize_and_center_crop(input_fg, image_width, image_height)
    bg = resize_and_center_crop(input_bg, image_width, image_height)
    concat_conds = numpy2pytorch([fg, bg]).to(device=vae.device, dtype=vae.dtype)
    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor
    concat_conds = torch.cat([c[None, ...] for c in concat_conds], dim=1)

    conds, unconds = encode_prompt_pair(positive_prompt=prompt + ', ' + a_prompt, negative_prompt=n_prompt)

    latents = t2i_pipe(
        prompt_embeds=conds,
        negative_prompt_embeds=unconds,
        width=image_width,
        height=image_height,
        num_inference_steps=steps,
        num_images_per_prompt=num_samples,
        generator=rng,
        output_type='latent',
        guidance_scale=cfg,
        cross_attention_kwargs={'concat_conds': concat_conds},
    ).images.to(vae.dtype) / vae.config.scaling_factor

    pixels = vae.decode(latents).sample
    pixels = pytorch2numpy(pixels)
    pixels = [resize_without_crop(
        image=p,
        target_width=int(round(image_width * highres_scale / 64.0) * 64),
        target_height=int(round(image_height * highres_scale / 64.0) * 64))
    for p in pixels]

    pixels = numpy2pytorch(pixels).to(device=vae.device, dtype=vae.dtype)
    latents = vae.encode(pixels).latent_dist.mode() * vae.config.scaling_factor
    latents = latents.to(device=unet.device, dtype=unet.dtype)

    image_height, image_width = latents.shape[2] * 8, latents.shape[3] * 8
    fg = resize_and_center_crop(input_fg, image_width, image_height)
    bg = resize_and_center_crop(input_bg, image_width, image_height)
    concat_conds = numpy2pytorch([fg, bg]).to(device=vae.device, dtype=vae.dtype)
    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor
    concat_conds = torch.cat([c[None, ...] for c in concat_conds], dim=1)

    latents = i2i_pipe(
        image=latents,
        strength=highres_denoise,
        prompt_embeds=conds,
        negative_prompt_embeds=unconds,
        width=image_width,
        height=image_height,
        num_inference_steps=int(round(steps / highres_denoise)),
        num_images_per_prompt=num_samples,
        generator=rng,
        output_type='latent',
        guidance_scale=cfg,
        cross_attention_kwargs={'concat_conds': concat_conds},
    ).images.to(vae.dtype) / vae.config.scaling_factor

    pixels = vae.decode(latents).sample
    pixels = pytorch2numpy(pixels, quant=False)

    clear_memory()
    return pixels, [fg, bg]


@torch.inference_mode()
def process_relight(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):
    global current_model
    if current_model != 'fc':
        initialize_pipeline_fc()
    input_fg, matting = run_rmbg(input_fg)
    results = process(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source)
    return input_fg, results



@torch.inference_mode()
def process_relight_bg(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source):
    global current_model
    if current_model != 'fbc':
        initialize_pipeline_fbc()
    bg_source = BGSource(bg_source)
    
    # Convert numerical inputs to appropriate types
    image_width = int(image_width)
    image_height = int(image_height)
    num_samples = int(num_samples)
    seed = int(seed)
    steps = int(steps)
    cfg = float(cfg)
    highres_scale = float(highres_scale)
    highres_denoise = float(highres_denoise)

    if bg_source == BGSource.UPLOAD:
        pass
    elif bg_source == BGSource.UPLOAD_FLIP:
        input_bg = np.fliplr(input_bg)
    elif bg_source == BGSource.GREY:
        input_bg = np.zeros(shape=(image_height, image_width, 3), dtype=np.uint8) + 64
    elif bg_source == BGSource.LEFT:
        gradient = np.linspace(224, 32, image_width)
        image = np.tile(gradient, (image_height, 1))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.RIGHT:
        gradient = np.linspace(32, 224, image_width)
        image = np.tile(gradient, (image_height, 1))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.TOP:
        gradient = np.linspace(224, 32, image_height)[:, None]
        image = np.tile(gradient, (1, image_width))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.BOTTOM:
        gradient = np.linspace(32, 224, image_height)[:, None]
        image = np.tile(gradient, (1, image_width))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    else:
        raise ValueError('Wrong background source!')
    
    input_fg, matting = run_rmbg(input_fg)
    results, extra_images = process_bg(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source)
    results = [(x * 255.0).clip(0, 255).astype(np.uint8) for x in results]
    final_results = results + extra_images

    # Save the generated images
    save_images(results, prefix="relight")
    
    return results


quick_prompts = [
    'sunshine from window',
    'neon light, city',
    'sunset over sea',
    'golden time',
    'sci-fi RGB glowing, cyberpunk',
    'natural lighting',
    'warm atmosphere, at home, bedroom',
    'magic lit',
    'evil, gothic, Yharnam',
    'light and shadow',
    'shadow from window',
    'soft studio lighting',
    'home atmosphere, cozy bedroom illumination',
    'neon, Wong Kar-wai, warm'
]
quick_prompts = [[x] for x in quick_prompts]


quick_subjects = [
    'modern sofa, high quality leather',
    'elegant dining table, polished wood',
    'luxurious bed, premium mattress',
    'minimalist office desk, clean design',
    'vintage wooden cabinet, antique finish',
]
quick_subjects = [[x] for x in quick_subjects]


class BGSource(Enum):
    UPLOAD = "Use Background Image"
    UPLOAD_FLIP = "Use Flipped Background Image"
    LEFT = "Left Light"
    RIGHT = "Right Light"
    TOP = "Top Light"
    BOTTOM = "Bottom Light"
    GREY = "Ambient"

# Add save function
def save_images(images, prefix="relight"):
    # Create output directory if it doesn't exist
    output_dir = Path("outputs")
    output_dir.mkdir(exist_ok=True)
    
    # Create timestamp for unique filenames
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    
    saved_paths = []
    for i, img in enumerate(images):
        if isinstance(img, np.ndarray):
            # Convert to PIL Image if numpy array
            img = Image.fromarray(img)
        
        # Create filename with timestamp
        filename = f"{prefix}_{timestamp}_{i+1}.png"
        filepath = output_dir / filename
        
        # Save image
        img.save(filepath)


    # print(f"Saved {len(saved_paths)} images to {output_dir}")
    return saved_paths


class MaskMover:
    def __init__(self):
        self.extracted_fg = None
        self.original_fg = None  # Store original foreground
        
    def set_extracted_fg(self, fg_image):
        """Store the extracted foreground with alpha channel"""
        if isinstance(fg_image, np.ndarray):
            self.extracted_fg = fg_image.copy()
            self.original_fg = fg_image.copy()
        else:
            self.extracted_fg = np.array(fg_image)
            self.original_fg = np.array(fg_image)
        return self.extracted_fg
    
    def create_composite(self, background, x_pos, y_pos, scale=1.0):
        """Create composite with foreground at specified position"""
        if self.original_fg is None or background is None:
            return background
        
        # Convert inputs to PIL Images
        if isinstance(background, np.ndarray):
            bg = Image.fromarray(background).convert('RGBA')
        else:
            bg = background.convert('RGBA')
        
        if isinstance(self.original_fg, np.ndarray):
            fg = Image.fromarray(self.original_fg).convert('RGBA')
        else:
            fg = self.original_fg.convert('RGBA')
        
        # Scale the foreground size
        new_width = int(fg.width * scale)
        new_height = int(fg.height * scale)
        fg = fg.resize((new_width, new_height), Image.LANCZOS)
        
        # Center the scaled foreground at the position
        x = int(x_pos - new_width / 2)
        y = int(y_pos - new_height / 2)
        
        # Create composite
        result = bg.copy()
        result.paste(fg, (x, y), fg)  # Use fg as the mask (requires fg to be in 'RGBA' mode)
        
        return np.array(result.convert('RGB'))  # Convert back to 'RGB' if needed
        
def get_depth(image):
    if image is None:
        return None
    # Convert from PIL/gradio format to cv2
    raw_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    # Get depth map
    depth = model.infer_image(raw_img) # HxW raw depth map
    # Normalize depth for visualization
    depth = ((depth - depth.min()) / (depth.max() - depth.min()) * 255).astype(np.uint8)
    # Convert to RGB for display
    depth_colored = cv2.applyColorMap(depth, cv2.COLORMAP_INFERNO)
    depth_colored = cv2.cvtColor(depth_colored, cv2.COLOR_BGR2RGB)
    return Image.fromarray(depth_colored)


from PIL import Image

def compress_image(image):
    # Convert Gradio image (numpy array) to PIL Image
    img = Image.fromarray(image)
    
    # Resize image if dimensions are too large
    max_size = 1024  # Maximum dimension size
    if img.width > max_size or img.height > max_size:
        ratio = min(max_size/img.width, max_size/img.height)
        new_size = (int(img.width * ratio), int(img.height * ratio))
        img = img.resize(new_size, Image.Resampling.LANCZOS)
    
    quality = 95  # Start with high quality
    img.save("compressed_image.jpg", "JPEG", quality=quality)  # Initial save
    
    # Check file size and adjust quality if necessary
    while os.path.getsize("compressed_image.jpg") > 100 * 1024:  # 100KB limit
        quality -= 5  # Decrease quality
        img.save("compressed_image.jpg", "JPEG", quality=quality)
        if quality < 20:  # Prevent quality from going too low
            break
    
    # Convert back to numpy array for Gradio
    compressed_img = np.array(Image.open("compressed_image.jpg"))
    return compressed_img


block = gr.Blocks().queue()
with block:
    with gr.Tab("Text"):
        with gr.Row():
            gr.Markdown("## Product Placement from Text")
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    input_fg = gr.Image(type="numpy", label="Image", height=480)
                    output_bg = gr.Image(type="numpy", label="Preprocessed Foreground", height=480)
                with gr.Group():
                    prompt = gr.Textbox(label="Prompt")
                    bg_source = gr.Radio(choices=[e.value for e in BGSource],
                                        value=BGSource.GREY.value,
                                        label="Lighting Preference (Initial Latent)", type='value')
                example_quick_subjects = gr.Dataset(samples=quick_subjects, label='Subject Quick List', samples_per_page=1000, components=[prompt])
                example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Lighting Quick List', samples_per_page=1000, components=[prompt])
                relight_button = gr.Button(value="Relight")

                with gr.Group():
                    with gr.Row():
                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
                        seed = gr.Number(label="Seed", value=12345, precision=0)

                    with gr.Row():
                        image_width = gr.Slider(label="Image Width", minimum=256, maximum=1024, value=512, step=64)
                        image_height = gr.Slider(label="Image Height", minimum=256, maximum=1024, value=640, step=64)

                with gr.Accordion("Advanced options", open=False):
                    steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=15, step=1)
                    cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=2, step=0.01)
                    lowres_denoise = gr.Slider(label="Lowres Denoise (for initial latent)", minimum=0.1, maximum=1.0, value=0.9, step=0.01)
                    highres_scale = gr.Slider(label="Highres Scale", minimum=1.0, maximum=3.0, value=1.5, step=0.01)
                    highres_denoise = gr.Slider(label="Highres Denoise", minimum=0.1, maximum=1.0, value=0.5, step=0.01)
                    a_prompt = gr.Textbox(label="Added Prompt", value='best quality')
                    n_prompt = gr.Textbox(label="Negative Prompt", value='lowres, bad anatomy, bad hands, cropped, worst quality')
            with gr.Column():
                result_gallery = gr.Gallery(height=832, object_fit='contain', label='Outputs')
        with gr.Row():
            dummy_image_for_outputs = gr.Image(visible=False, label='Result')
            # gr.Examples(
            #     fn=lambda *args: ([args[-1]], None),
            #     examples=db_examples.foreground_conditioned_examples,
            #     inputs=[
            #         input_fg, prompt, bg_source, image_width, image_height, seed, dummy_image_for_outputs
            #     ],
            #     outputs=[result_gallery, output_bg],
            #     run_on_click=True, examples_per_page=1024
            # )
        ips = [input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source]
        relight_button.click(fn=process_relight, inputs=ips, outputs=[output_bg, result_gallery])
        example_quick_prompts.click(lambda x, y: ', '.join(y.split(', ')[:2] + [x[0]]), inputs=[example_quick_prompts, prompt], outputs=prompt, show_progress=False, queue=False)
        example_quick_subjects.click(lambda x: x[0], inputs=example_quick_subjects, outputs=prompt, show_progress=False, queue=False)

    with gr.Tab("Background"):
        mask_mover = MaskMover()
        
        
        with gr.Row():
            gr.Markdown("## IC-Light (Relighting with Foreground and Background Condition)")
            gr.Markdown("💾 Generated images are automatically saved to 'outputs' folder")
        
        with gr.Row():
            with gr.Column():
                # Step 1: Input and Extract
                with gr.Row():
                    with gr.Group():
                        gr.Markdown("### Step 1: Extract Foreground")
                        input_image = gr.Image(type="numpy", label="Input Image", height=480)
                        # find_objects_button = gr.Button(value="Find Objects")
                        extract_button = gr.Button(value="Remove Background")
                    extracted_fg = gr.Image(type="numpy", label="Extracted Foreground", height=480)

                with gr.Row():    
                # Step 2: Background and Position
                    with gr.Group():
                        gr.Markdown("### Step 2: Position on Background")
                        input_bg = gr.Image(type="numpy", label="Background Image", height=480)
                        
                        with gr.Row():
                            x_slider = gr.Slider(
                                minimum=0,
                                maximum=1000,
                                label="X Position",
                                value=500,
                                visible=False
                            )
                            y_slider = gr.Slider(
                                minimum=0,
                                maximum=1000,
                                label="Y Position",
                                value=500,
                                visible=False
                            )
                            fg_scale_slider = gr.Slider(
                                label="Foreground Scale", 
                                minimum=0.01, 
                                maximum=3.0, 
                                value=1.0, 
                                step=0.01
                            )

                        editor = gr.ImageEditor(
                            type="numpy",
                            label="Position Foreground",
                            height=480,
                            visible=False
                        )
                        get_depth_button = gr.Button(value="Get Depth")
                        depth_image = gr.Image(type="numpy", label="Depth Image", height=480)
                
                # Step 3: Relighting Options
                with gr.Group():
                    gr.Markdown("### Step 3: Relighting Settings")
                    prompt = gr.Textbox(label="Prompt")
                    bg_source = gr.Radio(
                        choices=[e.value for e in BGSource],
                        value=BGSource.UPLOAD.value,
                        label="Background Source", 
                        type='value'
                    )

                    example_prompts = gr.Dataset(
                        samples=quick_prompts, 
                        label='Prompt Quick List', 
                        components=[prompt]
                    )
                    # bg_gallery = gr.Gallery(
                    #     height=450, 
                    #     label='Background Quick List', 
                    #     value=db_examples.bg_samples, 
                    #     columns=5, 
                    #     allow_preview=False
                    # )
                relight_button_bg = gr.Button(value="Relight")

                # Additional settings
                with gr.Group():
                    with gr.Row():
                        num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
                        seed = gr.Number(label="Seed", value=12345, precision=0)
                    with gr.Row():
                        image_width = gr.Slider(label="Image Width", minimum=256, maximum=1024, value=512, step=64)
                        image_height = gr.Slider(label="Image Height", minimum=256, maximum=1024, value=640, step=64)

                with gr.Accordion("Advanced options", open=False):
                    steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
                    cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=7.0, step=0.01)
                    highres_scale = gr.Slider(label="Highres Scale", minimum=1.0, maximum=2.0, value=1.2, step=0.01)
                    highres_denoise = gr.Slider(label="Highres Denoise", minimum=0.1, maximum=0.9, value=0.5, step=0.01)
                    a_prompt = gr.Textbox(label="Added Prompt", value='best quality')
                    n_prompt = gr.Textbox(
                        label="Negative Prompt",
                        value='lowres, bad anatomy, bad hands, cropped, worst quality'
                    )
                
            with gr.Column():
                result_gallery = gr.Image(height=832, label='Outputs')

        def extract_foreground(image):
            if image is None:
                return None, gr.update(visible=True), gr.update(visible=True)
            result, rgba = run_rmbg(image)
            mask_mover.set_extracted_fg(rgba)
            
            return result, gr.update(visible=True), gr.update(visible=True)


        original_bg = None

        extract_button.click(
            fn=extract_foreground,
            inputs=[input_image],
            outputs=[extracted_fg, x_slider, y_slider]
        )

        # find_objects_button.click(
        #     fn=find_objects,
        #     inputs=[input_image],
        #     outputs=[extracted_fg]
        # )
        
        get_depth_button.click(
            fn=get_depth,
            inputs=[input_bg],
            outputs=[depth_image]
        )

        # def update_position(background, x_pos, y_pos, scale):
        #     """Update composite when position changes"""
        #     global original_bg
        #     if background is None:
        #         return None
            
        #     if original_bg is None:
        #         original_bg = background.copy()
            
        #     # Convert string values to float
        #     x_pos = float(x_pos)
        #     y_pos = float(y_pos)
        #     scale = float(scale)
            
        #     return mask_mover.create_composite(original_bg, x_pos, y_pos, scale)

        class BackgroundManager:
            def __init__(self):
                self.original_bg = None

            def update_position(self, background, x_pos, y_pos, scale):
                """Update composite when position changes"""
                if background is None:
                    return None
                
                self.original_bg = background.copy()
                
                # Convert string values to float
                x_pos = float(x_pos)
                y_pos = float(y_pos)
                scale = float(scale)
                
                return mask_mover.create_composite(self.original_bg, x_pos, y_pos, scale)

        # Create an instance of BackgroundManager
        bg_manager = BackgroundManager()


        x_slider.change(
            fn=lambda bg, x, y, scale: bg_manager.update_position(bg, x, y, scale),
            inputs=[input_bg, x_slider, y_slider, fg_scale_slider],
            outputs=[input_bg]
        )

        y_slider.change(
            fn=lambda bg, x, y, scale: bg_manager.update_position(bg, x, y, scale),
            inputs=[input_bg, x_slider, y_slider, fg_scale_slider],
            outputs=[input_bg]
        )

        fg_scale_slider.change(
            fn=lambda bg, x, y, scale: bg_manager.update_position(bg, x, y, scale),
            inputs=[input_bg, x_slider, y_slider, fg_scale_slider],
            outputs=[input_bg]
        )

        # Update inputs list to include fg_scale_slider

    def process_relight_with_position(*args):
        if mask_mover.extracted_fg is None:
            gr.Warning("Please extract foreground first")
            return None
            
        background = args[1]  # Get background image
        x_pos = float(args[-3])  # x_slider value
        y_pos = float(args[-2])  # y_slider value
        scale = float(args[-1])  # fg_scale_slider value
        
        # Get original foreground size after scaling
        fg = Image.fromarray(mask_mover.original_fg)
        new_width = int(fg.width * scale)
        new_height = int(fg.height * scale)
        
        # Calculate crop region around foreground position
        crop_x = int(x_pos - new_width/2)
        crop_y = int(y_pos - new_height/2)
        crop_width = new_width
        crop_height = new_height
        
        # Add padding for context (20% extra on each side)
        padding = 0.2
        crop_x = int(crop_x - crop_width * padding)
        crop_y = int(crop_y - crop_height * padding)
        crop_width = int(crop_width * (1 + 2 * padding))
        crop_height = int(crop_height * (1 + 2 * padding))
        
        # Ensure crop dimensions are multiples of 8
        crop_width = ((crop_width + 7) // 8) * 8
        crop_height = ((crop_height + 7) // 8) * 8
        
        # Ensure crop region is within image bounds
        bg_height, bg_width = background.shape[:2]
        crop_x = max(0, min(crop_x, bg_width - crop_width))
        crop_y = max(0, min(crop_y, bg_height - crop_height))
        
        # Get actual crop dimensions after boundary check
        crop_width = min(crop_width, bg_width - crop_x)
        crop_height = min(crop_height, bg_height - crop_y)
        
        # Ensure dimensions are multiples of 8 again
        crop_width = (crop_width // 8) * 8
        crop_height = (crop_height // 8) * 8
        
        # Crop region from background
        crop_region = background[crop_y:crop_y+crop_height, crop_x:crop_x+crop_width]
        
        # Create composite in cropped region
        fg_local_x = int(new_width/2 + crop_width*padding)
        fg_local_y = int(new_height/2 + crop_height*padding)
        cropped_composite = mask_mover.create_composite(crop_region, fg_local_x, fg_local_y, scale)
        
        # Process the cropped region
        crop_args = list(args)
        crop_args[0] = cropped_composite
        crop_args[1] = crop_region
        crop_args[3] = crop_width
        crop_args[4] = crop_height
        crop_args = crop_args[:-3]  # Remove position and scale arguments
        
        # Get relit result
        relit_crop = process_relight_bg(*crop_args)[0]
        
        # Resize relit result to match crop dimensions if needed
        if relit_crop.shape[:2] != (crop_height, crop_width):
            relit_crop = resize_without_crop(relit_crop, crop_width, crop_height)
        
        # Place relit crop back into original background
        result = background.copy()
        result[crop_y:crop_y+crop_height, crop_x:crop_x+crop_width] = relit_crop
        
        return result
    
    ips_bg = [input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source]

    # Update button click events with new inputs list
    relight_button_bg.click(
        fn=process_relight_with_position,
        inputs=ips_bg,
        outputs=[result_gallery]
    )
    

    example_prompts.click(
        fn=lambda x: x[0],
        inputs=example_prompts,
        outputs=prompt,
        show_progress=False,
        queue=False
    )



block.launch(server_name='0.0.0.0', share=True)