app.py

import math
import random
from enum import Enum

import gradio as gr
import numpy as np
import safetensors.torch as sf
import torch
from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    StableDiffusionImg2ImgPipeline,
    StableDiffusionPipeline,
    UNet2DConditionModel,
)
from diffusers.models.attention_processor import AttnProcessor2_0
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import CLIPTextModel, CLIPTokenizer

from briarmbg import BriaRMBG

try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(duration=30):
            def decorator(fn):
                return fn

            return decorator


BASE_MODEL = "stablediffusionapi/realistic-vision-v51"
ICLIGHT_REPO = "lllyasviel/ic-light"
MODEL_FILE = "iclight_sd15_fc.safetensors"
NEGATIVE_PROMPT = "lowres, bad anatomy, bad hands, cropped, worst quality"
ADDED_PROMPT = "best quality"

_ENGINE = None


class BGSource(Enum):
    NONE = "None"
    LEFT = "Left Light"
    RIGHT = "Right Light"
    TOP = "Top Light"
    BOTTOM = "Bottom Light"


def ensure_rgb(image):
    if image is None:
        raise gr.Error("Upload an image first.")

    if isinstance(image, Image.Image):
        return np.array(image.convert("RGB"))

    if image.ndim == 2:
        image = np.stack([image, image, image], axis=-1)

    if image.shape[-1] == 4:
        image = np.array(Image.fromarray(image).convert("RGB"))

    return image[:, :, :3].astype(np.uint8)


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
    return np.array(resized_image.crop((left, top, right, bottom)))


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


def numpy2pytorch(imgs):
    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0
    return h.movedim(-1, 1)


def pytorch2numpy(imgs):
    results = []
    for x in imgs:
        y = x.movedim(0, -1)
        y = y * 127.5 + 127.5
        y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
        results.append(y)
    return results


class ICLightEngine:
    def __init__(self):
        if not torch.cuda.is_available():
            raise gr.Error("IC-Light inference requires a CUDA GPU. On Hugging Face, enable ZeroGPU hardware.")

        self.device = torch.device("cuda")
        self.tokenizer = CLIPTokenizer.from_pretrained(BASE_MODEL, subfolder="tokenizer")
        self.text_encoder = CLIPTextModel.from_pretrained(BASE_MODEL, subfolder="text_encoder")
        self.vae = AutoencoderKL.from_pretrained(BASE_MODEL, subfolder="vae")
        self.unet = UNet2DConditionModel.from_pretrained(BASE_MODEL, subfolder="unet")
        self.rmbg = BriaRMBG.from_pretrained("briaai/RMBG-1.4")

        self._patch_unet_input()
        self._load_iclight_weights()
        self._move_to_gpu()
        self._build_pipelines()

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

        unet_original_forward = self.unet.forward

        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)
            kwargs["cross_attention_kwargs"] = {}
            return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)

        self.unet.forward = hooked_unet_forward

    def _load_iclight_weights(self):
        model_path = hf_hub_download(ICLIGHT_REPO, MODEL_FILE)
        sd_offset = sf.load_file(model_path, device="cpu")
        sd_origin = self.unet.state_dict()
        sd_merged = {
            key: sd_origin[key] + sd_offset[key].to(dtype=sd_origin[key].dtype)
            for key in sd_origin.keys()
        }
        self.unet.load_state_dict(sd_merged, strict=True)
        del sd_offset, sd_origin, sd_merged

    def _move_to_gpu(self):
        self.text_encoder = self.text_encoder.to(device=self.device, dtype=torch.float16)
        self.vae = self.vae.to(device=self.device, dtype=torch.bfloat16)
        self.unet = self.unet.to(device=self.device, dtype=torch.float16)
        self.rmbg = self.rmbg.to(device=self.device, dtype=torch.float32)
        self.unet.set_attn_processor(AttnProcessor2_0())
        self.vae.set_attn_processor(AttnProcessor2_0())

    def _build_pipelines(self):
        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,
        )
        pipe_kwargs = dict(
            vae=self.vae,
            text_encoder=self.text_encoder,
            tokenizer=self.tokenizer,
            unet=self.unet,
            scheduler=scheduler,
            safety_checker=None,
            requires_safety_checker=False,
            feature_extractor=None,
            image_encoder=None,
        )
        self.t2i_pipe = StableDiffusionPipeline(**pipe_kwargs)
        self.i2i_pipe = StableDiffusionImg2ImgPipeline(**pipe_kwargs)

    @torch.inference_mode()
    def encode_prompt_inner(self, txt):
        max_length = self.tokenizer.model_max_length
        chunk_length = self.tokenizer.model_max_length - 2
        id_start = self.tokenizer.bos_token_id
        id_end = self.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 = self.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(chunk, id_pad, max_length) for chunk in chunks]

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

    @torch.inference_mode()
    def encode_prompt_pair(self, positive_prompt, negative_prompt):
        c = self.encode_prompt_inner(positive_prompt)
        uc = self.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 run_rmbg(self, img):
        height, width, channels = img.shape
        if channels != 3:
            raise gr.Error("Input image must be RGB.")

        k = (256.0 / float(height * width)) ** 0.5
        feed = resize_without_crop(img, int(64 * round(width * k)), int(64 * round(height * k)))
        feed = numpy2pytorch([feed]).to(device=self.device, dtype=torch.float32)
        alpha = self.rmbg(feed)[0][0]
        alpha = torch.nn.functional.interpolate(alpha, size=(height, width), mode="bilinear")
        alpha = alpha.movedim(1, -1)[0]
        alpha = alpha.detach().float().cpu().numpy().clip(0, 1)
        result = 127 + (img.astype(np.float32) - 127) * alpha
        return result.clip(0, 255).astype(np.uint8)

    def make_initial_background(self, bg_source, image_width, image_height):
        bg_source = BGSource(bg_source)
        if bg_source == BGSource.NONE:
            return None
        if bg_source == BGSource.LEFT:
            gradient = np.linspace(255, 0, image_width)
            image = np.tile(gradient, (image_height, 1))
        elif bg_source == BGSource.RIGHT:
            gradient = np.linspace(0, 255, image_width)
            image = np.tile(gradient, (image_height, 1))
        elif bg_source == BGSource.TOP:
            gradient = np.linspace(255, 0, image_height)[:, None]
            image = np.tile(gradient, (1, image_width))
        elif bg_source == BGSource.BOTTOM:
            gradient = np.linspace(0, 255, image_height)[:, None]
            image = np.tile(gradient, (1, image_width))
        else:
            raise gr.Error("Invalid lighting preference.")

        return np.stack((image,) * 3, axis=-1).astype(np.uint8)

    @torch.inference_mode()
    def relight(
        self,
        input_fg,
        prompt,
        image_width,
        image_height,
        num_samples,
        seed,
        steps,
        cfg,
        highres_scale,
        highres_denoise,
        lowres_denoise,
        bg_source,
    ):
        input_fg = ensure_rgb(input_fg)
        input_fg = self.run_rmbg(input_fg)
        input_bg = self.make_initial_background(bg_source, image_width, image_height)

        if seed is None or int(seed) < 0:
            seed = random.randint(0, 2**31 - 1)

        rng = torch.Generator(device=self.device).manual_seed(int(seed))
        fg = resize_and_center_crop(input_fg, image_width, image_height)

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

        conds, unconds = self.encode_prompt_pair(
            positive_prompt=f"{prompt}, {ADDED_PROMPT}",
            negative_prompt=NEGATIVE_PROMPT,
        )

        if input_bg is None:
            latents = self.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(self.vae.dtype) / self.vae.config.scaling_factor
        else:
            bg = resize_and_center_crop(input_bg, image_width, image_height)
            bg_latent = numpy2pytorch([bg]).to(device=self.vae.device, dtype=self.vae.dtype)
            bg_latent = self.vae.encode(bg_latent).latent_dist.mode() * self.vae.config.scaling_factor
            latents = self.i2i_pipe(
                image=bg_latent,
                strength=lowres_denoise,
                prompt_embeds=conds,
                negative_prompt_embeds=unconds,
                width=image_width,
                height=image_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(self.vae.dtype) / self.vae.config.scaling_factor

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

        pixels = numpy2pytorch(pixels).to(device=self.vae.device, dtype=self.vae.dtype)
        latents = self.vae.encode(pixels).latent_dist.mode() * self.vae.config.scaling_factor
        latents = latents.to(device=self.unet.device, dtype=self.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)
        concat_conds = numpy2pytorch([fg]).to(device=self.vae.device, dtype=self.vae.dtype)
        concat_conds = self.vae.encode(concat_conds).latent_dist.mode() * self.vae.config.scaling_factor

        latents = self.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(self.vae.dtype) / self.vae.config.scaling_factor

        pixels = self.vae.decode(latents).sample
        return input_fg, pytorch2numpy(pixels)


def get_engine():
    global _ENGINE
    if _ENGINE is None:
        _ENGINE = ICLightEngine()
    return _ENGINE


@spaces.GPU(duration=180)
def generate(
    image,
    prompt,
    lighting,
    width,
    height,
    samples,
    seed,
    steps,
    cfg,
    highres_scale,
    highres_denoise,
    lowres_denoise,
):
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt.")

    # Random seed always on: every generation gets a fresh seed regardless of
    # whatever is in the Seed field (which is display-only -- see its
    # `interactive=False` below).
    del seed
    seed = random.randint(0, 2**31 - 1)

    engine = get_engine()
    return engine.relight(
        image,
        prompt.strip(),
        int(width),
        int(height),
        int(samples),
        int(seed),
        int(steps),
        float(cfg),
        float(highres_scale),
        float(highres_denoise),
        float(lowres_denoise),
        lighting,
    )


quick_prompts = [

    # ------------------------------------------------------------------
    # Portrait Photography
    # ------------------------------------------------------------------
    ["Professional beauty portrait illuminated by a large octabox, soft wraparound lighting, smooth skin highlights, gentle shadow transitions, luxury cosmetic advertisement"],
    ["Editorial fashion portrait with dramatic Rembrandt lighting, crisp facial definition, rich contrast, subtle catchlights, magazine cover aesthetic"],
    ["Moody cinematic portrait lit by a single practical lamp, deep shadows, warm tungsten glow, atmospheric interior, shallow depth of field"],
    ["High-end studio portrait with clamshell lighting, flawless complexion, symmetrical illumination, premium beauty photography"],
    ["Classic Hollywood glamour lighting, elegant facial highlights, controlled shadow falloff, timeless black tie atmosphere"],
    ["Natural portrait beside a large window, soft daylight wrapping around the face, realistic ambient bounce, calm morning atmosphere"],
    ["Strong rim lighting defining facial contours against a dark background, dramatic cinematic portrait"],
    ["Soft butterfly lighting with subtle catchlights, elegant beauty editorial, natural skin rendering"],
    ["Split lighting creating bold contrast across the face, dramatic character portrait"],
    ["Warm candlelit portrait with flickering highlights and intimate shadows"],

    # ------------------------------------------------------------------
    # Cinematic
    # ------------------------------------------------------------------
    ["Epic cinematic lighting with powerful volumetric sunlight piercing through dust-filled air, dramatic atmosphere, rich dynamic range"],
    ["Film noir detective lighting, hard venetian blind shadows stretching across the room, deep blacks, dramatic side illumination"],
    ["Late-night motel illuminated only by flickering neon signage, colourful reflections across wet pavement, Wong Kar-wai inspired atmosphere"],
    ["Single shaft of sunlight breaking through abandoned warehouse windows, floating dust particles, dramatic industrial mood"],
    ["Strong backlight creating a glowing silhouette with subtle rim lighting separating the subject from the environment"],
    ["Large practical lights motivating realistic cinematic illumination with natural bounce light and believable shadow behaviour"],
    ["Sunbeams cutting through cathedral windows with visible volumetric rays and suspended dust"],
    ["Movie set lighting with practical lamps creating believable cinematic depth"],
    ["Moody alley illuminated by a single overhead streetlight with deep cinematic shadows"],
    ["Orange and teal blockbuster lighting with cinematic colour separation"],

    # ------------------------------------------------------------------
    # Natural Light
    # ------------------------------------------------------------------
    ["Warm sunrise illuminating rolling morning mist, soft golden rays filtering through trees, peaceful tranquil atmosphere"],
    ["Bright overcast daylight producing perfectly diffused shadows and even natural illumination, ideal outdoor portrait lighting"],
    ["Afternoon sunlight filtered through dense forest canopy, beautiful dappled light patterns dancing across surfaces"],
    ["Soft coastal sunset reflecting warm orange light across calm ocean waves, gentle ambient fill from the sky"],
    ["Fresh winter morning with cool blue ambient light reflecting from snow, crisp clean highlights and subtle contrast"],
    ["Mountain sunrise casting long dramatic shadows across rocky terrain with warm atmospheric haze"],
    ["Golden hour light spilling through tall grass with warm glowing highlights"],
    ["Soft spring daylight beneath blossoming trees with natural colour balance"],
    ["Clear summer noon with crisp sunlight and hard defined shadows"],
    ["Blue hour landscape illuminated by soft twilight gradients"],

    # ------------------------------------------------------------------
    # Interior Lighting
    # ------------------------------------------------------------------
    ["Cozy bedroom illuminated by warm bedside lamps, gentle pools of amber light, relaxing evening atmosphere"],
    ["Luxury hotel suite lit with sophisticated architectural lighting, warm indirect illumination, premium interior design"],
    ["Rustic cabin glowing from a roaring fireplace, flickering orange light dancing naturally across wooden walls"],
    ["Modern coffee shop illuminated by hanging Edison bulbs, warm inviting ambience, realistic practical lighting"],
    ["Minimalist Scandinavian living room with huge floor-to-ceiling windows flooding the interior with soft daylight"],
    ["Rainy evening inside a quiet apartment with cool window light mixing beautifully with warm indoor lamps"],
    ["Elegant dining room illuminated by candlelight with soft ambient shadows"],
    ["Industrial loft with skylight casting dramatic shafts of daylight"],
    ["Contemporary office lit by cool daylight through expansive glass walls"],
    ["Reading nook illuminated by a single warm floor lamp creating a peaceful atmosphere"],

    # ------------------------------------------------------------------
    # Studio Photography
    # ------------------------------------------------------------------
    ["Large octabox providing ultra-soft wraparound illumination with subtle shadow transitions"],
    ["Professional beauty dish creating crisp facial highlights and natural skin tones"],
    ["Three-point studio lighting with balanced key, fill and rim illumination"],
    ["High-key studio lighting with bright seamless white background and soft shadows"],
    ["Low-key portrait lighting with controlled highlights and deep dramatic blacks"],
    ["Fashion editorial lighting using strip softboxes and clean rim lighting"],
    ["Luxury commercial photography lighting with perfect exposure and premium finish"],
    ["Classic product photography with large diffused softboxes and controlled reflections"],
    ["Ring light portrait with symmetrical catchlights and flawless illumination"],
    ["Soft overhead studio lighting with natural falloff and premium commercial quality"],

    # ------------------------------------------------------------------
    # Sci-Fi
    # ------------------------------------------------------------------
    ["Cyberpunk alley illuminated by vibrant cyan and magenta neon signs reflecting across rain-soaked streets, futuristic atmosphere"],
    ["Spaceship bridge glowing with clean white LED panels, holographic displays casting subtle blue illumination"],
    ["Advanced laboratory filled with cool sterile lighting, illuminated control panels and glowing scientific equipment"],
    ["Alien reactor core emitting brilliant blue energy, volumetric light scattering through metallic architecture"],
    ["High-tech command centre featuring layered RGB accent lighting and premium futuristic industrial design"],
    ["Futuristic AI chamber illuminated entirely by floating holographic interfaces and glowing translucent surfaces"],
    ["Glowing cybernetic implants illuminating the surrounding environment with electric blue light"],
    ["Sci-fi corridor illuminated by animated LED floor lighting and holographic signage"],
    ["Energy shield casting shimmering blue reflections across futuristic armour"],
    ["Massive orbital station interior illuminated by clean architectural lighting"],

    # ------------------------------------------------------------------
    # Fantasy
    # ------------------------------------------------------------------
    ["Ancient enchanted forest illuminated by thousands of softly glowing magical fireflies, mystical dreamlike atmosphere"],
    ["Powerful wizard holding a radiant magical orb casting dynamic blue light across stone ruins"],
    ["Crystal cave illuminated by naturally glowing gemstones, beautiful refracted rainbow light scattering throughout"],
    ["Ancient temple flooded by divine heavenly light streaming through massive stone columns"],
    ["Mystical portal radiating swirling magical energy with realistic coloured illumination affecting surrounding terrain"],
    ["Moonlit fairy grove glowing with ethereal blue ambient light and floating enchanted particles"],
    ["Dragon's fiery breath illuminating an ancient cavern with glowing embers"],
    ["Sacred relic emitting radiant golden light within forgotten ruins"],
    ["Mystical waterfall glowing with enchanted moonlight and magical mist"],
    ["Ancient spell circle casting colourful magical reflections across stone floors"],

    # ------------------------------------------------------------------
    # Horror
    # ------------------------------------------------------------------
    ["Dark abandoned hospital corridor illuminated only by flickering fluorescent lights, unsettling atmosphere, long eerie shadows"],
    ["Single flashlight beam cutting through thick darkness, dramatic high contrast with realistic light falloff"],
    ["Haunted Victorian house lit by slowly flickering candles, moving shadows creating an unsettling mood"],
    ["Emergency red lighting filling an underground bunker with ominous tension and dramatic contrast"],
    ["Dense fog illuminated by pale moonlight creating mysterious silhouettes and eerie atmospheric depth"],
    ["Old basement with a single exposed light bulb swinging gently, constantly shifting shadows"],
    ["Abandoned church lit by dim candlelight and cold moonlight spilling through broken windows"],
    ["Creepy forest illuminated only by a weak lantern with heavy surrounding darkness"],
    ["Dark corridor disappearing into shadow beneath flickering emergency lights"],
    ["Unnatural green glow emerging from beneath an ancient doorway"],

    # ------------------------------------------------------------------
    # Product Photography
    # ------------------------------------------------------------------
    ["Luxury wristwatch photographed with precision studio lighting, elegant specular highlights, premium commercial advertising quality"],
    ["Cosmetic bottle illuminated by enormous soft diffused light sources producing beautiful glass reflections"],
    ["Premium perfume bottle on polished black acrylic with dramatic spotlight and subtle rim lighting"],
    ["High-end jewellery photography with carefully controlled reflections and sparkling gemstone highlights"],
    ["Minimalist technology product floating in clean white studio lighting with soft shadow gradients"],
    ["Professional food photography illuminated by realistic restaurant window light with rich appetising colour"],
    ["Luxury automotive detail highlighted with dramatic studio reflections"],
    ["Premium leather product illuminated by warm commercial lighting with rich texture"],
    ["Elegant wine bottle lit with soft gradients and controlled specular reflections"],
    ["Luxury cosmetic packaging illuminated with clean premium advertising lighting"],

    # ------------------------------------------------------------------
    # Environmental Effects
    # ------------------------------------------------------------------
    ["Heavy rain illuminated by warm streetlights creating glowing droplets and shimmering reflections"],
    ["Powerful lightning flash briefly illuminating an entire landscape with dramatic storm clouds"],
    ["Dense fog softly scattering distant headlights creating cinematic atmospheric perspective"],
    ["Snow-covered landscape reflecting cool blue twilight with subtle warm cabin lights in the distance"],
    ["Campfire casting dynamic flickering orange light onto nearby faces while stars fill the night sky"],
    ["Dust-filled abandoned church with magnificent god rays pouring through stained glass windows"],
    ["Soft mist diffusing sunrise light into glowing atmospheric beams"],
    ["Sandstorm illuminated by low afternoon sunlight creating dramatic golden haze"],
    ["Ocean spray catching warm sunset light with sparkling reflections"],
    ["Autumn leaves illuminated by warm evening sunlight with rich seasonal colours"],

    # ------------------------------------------------------------------
    # Creative Lighting
    # ------------------------------------------------------------------
    ["Warm golden hour sunlight streaming through trees, long soft shadows, rich amber highlights, gentle atmospheric haze, cinematic depth"],
    ["Large north-facing window casting soft diffused daylight across the subject, subtle falloff into shadow, realistic ambient bounce"],
    ["Brilliant cyan neon signs reflecting across rain-soaked pavement, strong magenta rim lighting, cinematic cyberpunk atmosphere"],
    ["Soft moonlight filtering through clouds, subtle silver highlights, calm nocturnal ambience"],
    ["Warm fireplace glow contrasting against cool blue evening window light"],
    ["Glowing stained glass illuminating the scene with vibrant coloured reflections"],
    ["Volumetric sunlight pouring through ancient ruins with dramatic dust-filled atmosphere"],
    ["Soft reflected light bouncing naturally between white walls creating realistic interior illumination"],
    ["Strong overhead spotlight with crisp shadows creating dramatic theatrical staging"],
    ["Beautiful mixed lighting combining cool twilight with warm practical interior lamps"]
]


with gr.Blocks(title="IC-Light Relighting") as demo:
    gr.Markdown("## IC-Light Relighting")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(sources=["upload"], type="numpy", label="Image", height=440)
            prompt = gr.Textbox(label="Prompt", value="portrait, cinematic lighting")
            lighting = gr.Radio(
                choices=[e.value for e in BGSource],
                value=BGSource.NONE.value,
                label="Lighting Preference",
            )
            prompt_examples = gr.Dataset(
                samples=quick_prompts,
                label="Prompt Quick List",
                components=[prompt],
                samples_per_page=20,
            )
            prompt_examples.click(
                lambda x: x[0],
                inputs=prompt_examples,
                outputs=prompt,
                show_progress=False,
                queue=False,
            )
            run_button = gr.Button("Relight", variant="primary")

            with gr.Row():
                samples = gr.Slider(label="Images", minimum=1, maximum=4, value=1, step=1)
                seed = gr.Number(label="Seed (randomized every run)", value=12345, precision=0, interactive=False)

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

            with gr.Accordion("Advanced", open=False):
                steps = gr.Slider(label="Steps", minimum=1, maximum=80, value=25, step=1)
                cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=16.0, value=2.0, step=0.1)
                lowres_denoise = gr.Slider(
                    label="Lowres Denoise",
                    minimum=0.1,
                    maximum=1.0,
                    value=0.9,
                    step=0.01,
                )
                highres_scale = gr.Slider(
                    label="Highres Scale",
                    minimum=1.0,
                    maximum=2.0,
                    value=1.5,
                    step=0.05,
                )
                highres_denoise = gr.Slider(
                    label="Highres Denoise",
                    minimum=0.1,
                    maximum=1.0,
                    value=0.5,
                    step=0.01,
                )
        with gr.Column():
            foreground = gr.Image(type="numpy", label="Preprocessed Foreground", height=360)
            gallery = gr.Gallery(label="Outputs", height=720, object_fit="contain")

    inputs = [
        input_image,
        prompt,
        lighting,
        width,
        height,
        samples,
        seed,
        steps,
        cfg,
        highres_scale,
        highres_denoise,
        lowres_denoise,
    ]
    run_button.click(fn=generate, inputs=inputs, outputs=[foreground, gallery])


if __name__ == "__main__":
    demo.queue(max_size=20).launch(server_name="0.0.0.0")