Upload 7 files

Dockerfile

@@ -0,0 +1,47 @@
+# Start from NVIDIA CUDA base image
+FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
+
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+ENV TORCH_HOME=/app/models
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    python3.10 \
+    python3-pip \
+    git \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Copy requirements file
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip3 install --no-cache-dir torch==2.4.1+cu121 torchvision==0.19.1+cu121 --index-url https://download.pytorch.org/whl/cu121
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+# Create models directory
+RUN mkdir -p /app/models /app/outputs
+
+# Copy application files
+COPY . .
+
+# Download model weights if needed
+RUN mkdir -p /app/models && \
+    if [ ! -f /app/models/iclight_sd15_fc.safetensors ]; then \
+        wget -P /app/models https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fc.safetensors; \
+    fi && \
+    if [ ! -f /app/models/iclight_sd15_fbc.safetensors ]; then \
+        wget -P /app/models https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fbc.safetensors; \
+    fi
+
+# Expose port for Gradio
+EXPOSE 7860
+
+# Command to run the application
+CMD ["python3", "gradio_demo.py"] 

briarmbg.py

@@ -0,0 +1,462 @@
+# RMBG1.4 (diffusers implementation)
+# Found on huggingface space of several projects
+# Not sure which project is the source of this file
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class REBNCONV(nn.Module):
+    def __init__(self, in_ch=3, out_ch=3, dirate=1, stride=1):
+        super(REBNCONV, self).__init__()
+
+        self.conv_s1 = nn.Conv2d(
+            in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate, stride=stride
+        )
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+
+def _upsample_like(src, tar):
+    src = F.interpolate(src, size=tar.shape[2:], mode="bilinear")
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
+        super(RSU7, self).__init__()
+
+        self.in_ch = in_ch
+        self.mid_ch = mid_ch
+        self.out_ch = out_ch
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)  ## 1 -> 1/2
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d = self.rebnconv6d(torch.cat((hx7, hx6), 1))
+        hx6dup = _upsample_like(hx6d, hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-6 ###
+class RSU6(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-5 ###
+class RSU5(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-4 ###
+class RSU4(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-4F ###
+class RSU4F(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d, hx2), 1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d, hx1), 1))
+
+        return hx1d + hxin
+
+
+class myrebnconv(nn.Module):
+    def __init__(
+        self,
+        in_ch=3,
+        out_ch=1,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        dilation=1,
+        groups=1,
+    ):
+        super(myrebnconv, self).__init__()
+
+        self.conv = nn.Conv2d(
+            in_ch,
+            out_ch,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+        )
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.rl = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        return self.rl(self.bn(self.conv(x)))
+
+
+class BriaRMBG(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, config: dict = {"in_ch": 3, "out_ch": 1}):
+        super(BriaRMBG, self).__init__()
+        in_ch = config["in_ch"]
+        out_ch = config["out_ch"]
+        self.conv_in = nn.Conv2d(in_ch, 64, 3, stride=2, padding=1)
+        self.pool_in = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage1 = RSU7(64, 32, 64)
+        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage2 = RSU6(64, 32, 128)
+        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage3 = RSU5(128, 64, 256)
+        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage4 = RSU4(256, 128, 512)
+        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage5 = RSU4F(512, 256, 512)
+        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage6 = RSU4F(512, 256, 512)
+
+        # decoder
+        self.stage5d = RSU4F(1024, 256, 512)
+        self.stage4d = RSU4(1024, 128, 256)
+        self.stage3d = RSU5(512, 64, 128)
+        self.stage2d = RSU6(256, 32, 64)
+        self.stage1d = RSU7(128, 16, 64)
+
+        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side3 = nn.Conv2d(128, out_ch, 3, padding=1)
+        self.side4 = nn.Conv2d(256, out_ch, 3, padding=1)
+        self.side5 = nn.Conv2d(512, out_ch, 3, padding=1)
+        self.side6 = nn.Conv2d(512, out_ch, 3, padding=1)
+
+        # self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self, x):
+        hx = x
+
+        hxin = self.conv_in(hx)
+        # hx = self.pool_in(hxin)
+
+        # stage 1
+        hx1 = self.stage1(hxin)
+        hx = self.pool12(hx1)
+
+        # stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        # stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        # stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        # stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        # stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6, hx5)
+
+        # -------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        # side output
+        d1 = self.side1(hx1d)
+        d1 = _upsample_like(d1, x)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2, x)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3, x)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4, x)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5, x)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6, x)
+
+        return [
+            F.sigmoid(d1),
+            F.sigmoid(d2),
+            F.sigmoid(d3),
+            F.sigmoid(d4),
+            F.sigmoid(d5),
+            F.sigmoid(d6),
+        ], [hx1d, hx2d, hx3d, hx4d, hx5d, hx6]

db_examples.py

@@ -0,0 +1,217 @@
+foreground_conditioned_examples = [
+    [
+        "imgs/i1.webp",
+        "beautiful woman, detailed face, sunshine, outdoor, warm atmosphere",
+        "Right Light",
+        512,
+        960,
+        12345,
+        "imgs/o1.png",
+    ],
+    [
+        "imgs/i1.webp",
+        "beautiful woman, detailed face, sunshine, outdoor, warm atmosphere",
+        "Left Light",
+        512,
+        960,
+        50,
+        "imgs/o2.png",
+    ],
+    [
+        "imgs/i3.png",
+        "beautiful woman, detailed face, neon, Wong Kar-wai, warm",
+        "Left Light",
+        512,
+        768,
+        12345,
+        "imgs/o3.png",
+    ],
+    [
+        "imgs/i3.png",
+        "beautiful woman, detailed face, sunshine from window",
+        "Left Light",
+        512,
+        768,
+        12345,
+        "imgs/o4.png",
+    ],
+    [
+        "imgs/i5.png",
+        "beautiful woman, detailed face, warm atmosphere, at home, bedroom",
+        "Left Light",
+        512,
+        768,
+        123,
+        "imgs/o5.png",
+    ],
+    [
+        "imgs/i6.jpg",
+        "beautiful woman, detailed face, sunshine from window",
+        "Right Light",
+        512,
+        768,
+        42,
+        "imgs/o6.png",
+    ],
+    [
+        "imgs/i7.jpg",
+        "beautiful woman, detailed face, shadow from window",
+        "Left Light",
+        512,
+        768,
+        8888,
+        "imgs/o7.png",
+    ],
+    [
+        "imgs/i8.webp",
+        "beautiful woman, detailed face, sunset over sea",
+        "Right Light",
+        512,
+        640,
+        42,
+        "imgs/o8.png",
+    ],
+    [
+        "imgs/i9.png",
+        "handsome boy, detailed face, neon light, city",
+        "Left Light",
+        512,
+        640,
+        12345,
+        "imgs/o9.png",
+    ],
+    [
+        "imgs/i10.png",
+        "beautiful woman, detailed face, light and shadow",
+        "Left Light",
+        512,
+        960,
+        8888,
+        "imgs/o10.png",
+    ],
+    [
+        "imgs/i11.png",
+        "Buddha, detailed face, sci-fi RGB glowing, cyberpunk",
+        "Left Light",
+        512,
+        768,
+        8888,
+        "imgs/o11.png",
+    ],
+    [
+        "imgs/i11.png",
+        "Buddha, detailed face, natural lighting",
+        "Left Light",
+        512,
+        768,
+        12345,
+        "imgs/o12.png",
+    ],
+    [
+        "imgs/i13.png",
+        "toy, detailed face, shadow from window",
+        "Bottom Light",
+        512,
+        704,
+        12345,
+        "imgs/o13.png",
+    ],
+    [
+        "imgs/i14.png",
+        "toy, detailed face, sunset over sea",
+        "Right Light",
+        512,
+        704,
+        100,
+        "imgs/o14.png",
+    ],
+    [
+        "imgs/i15.png",
+        "dog, magic lit, sci-fi RGB glowing, studio lighting",
+        "Bottom Light",
+        512,
+        768,
+        12345,
+        "imgs/o15.png",
+    ],
+    [
+        "imgs/i16.png",
+        "mysteriou human, warm atmosphere, warm atmosphere, at home, bedroom",
+        "Right Light",
+        512,
+        768,
+        100,
+        "imgs/o16.png",
+    ],
+]
+
+bg_samples = [
+    'imgs/bgs/1.webp',
+    'imgs/bgs/2.webp',
+    'imgs/bgs/3.webp',
+    'imgs/bgs/4.webp',
+    'imgs/bgs/5.webp',
+    'imgs/bgs/6.webp',
+    'imgs/bgs/7.webp',
+    'imgs/bgs/8.webp',
+    'imgs/bgs/9.webp',
+    'imgs/bgs/10.webp',
+    'imgs/bgs/11.png',
+    'imgs/bgs/12.png',
+    'imgs/bgs/13.png',
+    'imgs/bgs/14.png',
+    'imgs/bgs/15.png',
+]
+
+background_conditioned_examples = [
+    [
+        "imgs/alter/i3.png",
+        "imgs/bgs/7.webp",
+        "beautiful woman, cinematic lighting",
+        "Use Background Image",
+        512,
+        768,
+        12345,
+        "imgs/alter/o1.png",
+    ],
+    [
+        "imgs/alter/i2.png",
+        "imgs/bgs/11.png",
+        "statue of an angel, natural lighting",
+        "Use Flipped Background Image",
+        512,
+        768,
+        12345,
+        "imgs/alter/o2.png",
+    ],
+    [
+        "imgs/alter/i1.jpeg",
+        "imgs/bgs/2.webp",
+        "beautiful woman, cinematic lighting",
+        "Use Background Image",
+        512,
+        768,
+        12345,
+        "imgs/alter/o3.png",
+    ],
+    [
+        "imgs/alter/i1.jpeg",
+        "imgs/bgs/3.webp",
+        "beautiful woman, cinematic lighting",
+        "Use Background Image",
+        512,
+        768,
+        12345,
+        "imgs/alter/o4.png",
+    ],
+    [
+        "imgs/alter/i6.webp",
+        "imgs/bgs/15.png",
+        "handsome man, cinematic lighting",
+        "Use Background Image",
+        512,
+        768,
+        12345,
+        "imgs/alter/o5.png",
+    ],
+]

gradio_demo.py

@@ -0,0 +1,1103 @@
+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 = 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()
+
+# 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)
+    kwargs['cross_attention_kwargs'] = {}
+    return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)
+
+
+unet.forward = hooked_unet_forward
+
+# Load
+
+model_path = './models/iclight_sd15_fc.safetensors'
+# model_path = './models/iclight_sd15_fbc.safetensors'
+
+
+# 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)
+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):
+    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):
+    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", visible=False):
+        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
+                
+                if self.original_bg is 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)

gradio_demo_bg.py

@@ -0,0 +1,1004 @@
+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 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
+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.florence import load_florence_model, run_florence_inference, \
+    FLORENCE_OPEN_VOCABULARY_DETECTION_TASK
+from utils.sam import load_sam_image_model, run_sam_inference
+
+
+import torch
+DEVICE = torch.device("cuda")
+
+torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
+if torch.cuda.get_device_properties(0).major >= 8:
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+
+FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=DEVICE)
+SAM_IMAGE_MODEL = load_sam_image_model(device=DEVICE)
+
+@spaces.GPU(duration=20)
+@torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+def process_image(image_input, text_input) -> Optional[Image.Image]:
+    # if not image_input:
+    #     gr.Info("Please upload an image.")
+    #     return None
+
+    # if not text_input:
+    #     gr.Info("Please enter a text prompt.")
+    #     return None
+
+    _, result = run_florence_inference(
+        model=FLORENCE_MODEL,
+        processor=FLORENCE_PROCESSOR,
+        device=DEVICE,
+        image=image_input,
+        task=FLORENCE_OPEN_VOCABULARY_DETECTION_TASK,
+        text=text_input
+    )
+    detections = sv.Detections.from_lmm(
+        lmm=sv.LMM.FLORENCE_2,
+        result=result,
+        resolution_wh=image_input.size
+    )
+    detections = run_sam_inference(SAM_IMAGE_MODEL, image_input, detections)
+    if len(detections) == 0:
+        gr.Info("No objects detected.")
+        return None
+    return Image.fromarray(detections.mask[0].astype("uint8") * 255)
+
+
+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")
+
+# '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 = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")
+rmbg = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
+
+# Change UNet
+
+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
+
+unet_original_forward = 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)
+
+
+unet.forward = hooked_unet_forward
+
+# Load
+
+model_path = './models/iclight_sd15_fbc.safetensors'
+
+if not os.path.exists(model_path):
+    download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fbc.safetensors', dst=model_path)
+
+# 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 smaller attention slice size for RTX 2070
+    torch.backends.cuda.max_split_size_mb = 512
+
+# 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
+
+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.eval()
+
+
+from utils.florence import load_florence_model, run_florence_inference, \
+    FLORENCE_OPEN_VOCABULARY_DETECTION_TASK
+from utils.sam import load_sam_image_model, run_sam_inference
+
+
+torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
+if torch.cuda.get_device_properties(0).major >= 8:
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+
+#FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=device)
+SAM_IMAGE_MODEL = load_sam_image_model(device=device)
+
+# Update the state dict merging to use correct dtype
+sd_offset = sf.load_file(model_path)
+sd_origin = unet.state_dict()
+sd_merged = {k: sd_origin[k] + sd_offset[k].to(device=device, dtype=dtype) for k in sd_origin.keys()}
+unet.load_state_dict(sd_merged, strict=True)
+del sd_offset, sd_origin, sd_merged
+
+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()
+
+# Samplers
+
+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, 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]
+
+
+# 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
+
+
+# Modify process_relight to save images
+@torch.inference_mode()
+def process_relight(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source):
+    input_fg, matting = run_rmbg(input_fg)
+    # show input_fg in a new image 
+    input_fg_img = Image.fromarray(input_fg)
+    results, extra_images = process(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
+
+
+# Modify process_normal to save images
+@torch.inference_mode()
+def process_normal(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source):
+    input_fg, matting = run_rmbg(input_fg, sigma=16)
+
+    print('left ...')
+    left = process(input_fg, input_bg, prompt, image_width, image_height, 1, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, BGSource.LEFT.value)[0][0]
+
+    print('right ...')
+    right = process(input_fg, input_bg, prompt, image_width, image_height, 1, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, BGSource.RIGHT.value)[0][0]
+
+    print('bottom ...')
+    bottom = process(input_fg, input_bg, prompt, image_width, image_height, 1, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, BGSource.BOTTOM.value)[0][0]
+
+    print('top ...')
+    top = process(input_fg, input_bg, prompt, image_width, image_height, 1, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, BGSource.TOP.value)[0][0]
+
+    inner_results = [left * 2.0 - 1.0, right * 2.0 - 1.0, bottom * 2.0 - 1.0, top * 2.0 - 1.0]
+
+    ambient = (left + right + bottom + top) / 4.0
+    h, w, _ = ambient.shape
+    matting = resize_and_center_crop((matting[..., 0] * 255.0).clip(0, 255).astype(np.uint8), w, h).astype(np.float32)[..., None] / 255.0
+
+    def safa_divide(a, b):
+        e = 1e-5
+        return ((a + e) / (b + e)) - 1.0
+
+    left = safa_divide(left, ambient)
+    right = safa_divide(right, ambient)
+    bottom = safa_divide(bottom, ambient)
+    top = safa_divide(top, ambient)
+
+    u = (right - left) * 0.5
+    v = (top - bottom) * 0.5
+
+    sigma = 10.0
+    u = np.mean(u, axis=2)
+    v = np.mean(v, axis=2)
+    h = (1.0 - u ** 2.0 - v ** 2.0).clip(0, 1e5) ** (0.5 * sigma)
+    z = np.zeros_like(h)
+
+    normal = np.stack([u, v, h], axis=2)
+    normal /= np.sum(normal ** 2.0, axis=2, keepdims=True) ** 0.5
+    normal = normal * matting + np.stack([z, z, 1 - z], axis=2) * (1 - matting)
+
+    results = [normal, left, right, bottom, top] + inner_results
+    results = [(x * 127.5 + 127.5).clip(0, 255).astype(np.uint8) for x in results]
+
+
+    # Save the generated images
+    save_images(results, prefix="normal")
+    
+    return results
+
+
+
+
+quick_prompts = [
+    'modern sofa in living room',
+    'elegant dining table with chairs',
+    'luxurious bed in bedroom, cinematic lighting',
+    'minimalist office desk, natural lighting',
+    'vintage wooden cabinet, warm lighting',
+    'contemporary bookshelf, ambient lighting',
+    'designer armchair, dramatic lighting',
+    'modern kitchen island, bright lighting',
+]
+quick_prompts = [[x] for x in quick_prompts]
+
+
+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"
+
+
+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"""
+        self.extracted_fg = fg_image.copy()
+        self.original_fg = fg_image.copy()  # Keep original
+        return fg_image
+    
+    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)
+        else:
+            bg = background
+        
+        if isinstance(self.original_fg, np.ndarray):
+            fg = Image.fromarray(self.original_fg)
+        else:
+            fg = self.original_fg
+        
+        # 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()
+        if fg.mode == 'RGBA':  # If foreground has alpha channel
+            result.paste(fg, (x, y), fg.split()[3])  # Use alpha channel as mask
+        else:
+            result.paste(fg, (x, y))
+        
+        return np.array(result)
+    
+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)
+
+# def find_objects(image_input):
+#     detections = run_sam_inference(SAM_IMAGE_MODEL, image_input, detections)
+#     if len(detections) == 0:
+#         gr.Info("No objects detected.")
+#         return None
+#     return Image.fromarray(detections.mask[0].astype("uint8") * 255)
+
+
+
+block = gr.Blocks().queue()
+with block:
+    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.Group():
+                gr.Markdown("### Step 1: Extract Foreground")
+                input_image = gr.Image(type="numpy", label="Input Image", height=480)
+                input_text = gr.Textbox(label="Describe target object")
+
+                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)
+
+            
+            # 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
+                    )
+
+                    get_depth_button = gr.Button(value="Get Depth")
+
+                    depth_image = gr.Image(type="numpy", label="Depth Image", height=480)
+
+                
+                editor = gr.ImageEditor(
+                    type="numpy",
+                    label="Position Foreground",
+                    height=480,
+                    visible=False
+                )
+            
+            # 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 = 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'
+                )
+                normal_button = gr.Button(value="Compute Normal (4x Slower)")
+        
+        with gr.Column():
+            result_gallery = gr.Image(height=832, label='Outputs')
+
+    # Event handlers
+    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=process_image,
+        inputs=[input_image, input_text],
+        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)
+
+    x_slider.change(
+        fn=update_position,
+        inputs=[input_bg, x_slider, y_slider, fg_scale_slider],
+        outputs=[input_bg]
+    )
+
+    y_slider.change(
+        fn=update_position,
+        inputs=[input_bg, x_slider, y_slider, fg_scale_slider],
+        outputs=[input_bg]
+    )
+
+    fg_scale_slider.change(
+        fn=update_position,
+        inputs=[input_bg, x_slider, y_slider, fg_scale_slider],
+        outputs=[input_bg]
+    )
+
+    # Update inputs list to include fg_scale_slider
+    ips = [input_bg, input_bg, prompt, image_width, image_height, num_samples, 
+           seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, 
+           bg_source, x_slider, y_slider, fg_scale_slider]  # Added 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(*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
+
+    # Update button click events with new inputs list
+    relight_button.click(
+        fn=process_relight_with_position,
+        inputs=ips,
+        outputs=[result_gallery]
+    )
+    
+    # Update normal_button to use same argument handling
+    def process_normal_with_position(*args):
+        if mask_mover.extracted_fg is None:
+            gr.Warning("Please extract foreground first")
+            return None
+            
+        background = args[1]
+        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))
+        
+        # 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]
+        
+        # Get processed result
+        processed_crop = process_normal(*crop_args)
+        
+        # Place processed crop back into original background
+        result = background.copy()
+        result[crop_y:crop_y+crop_height, crop_x:crop_x+crop_width] = processed_crop
+        
+        return result[0] if result else None
+
+    normal_button.click(
+        fn=process_normal_with_position,
+        inputs=ips,
+        outputs=[result_gallery]
+    )
+
+    example_prompts.click(
+        fn=lambda x: x[0],
+        inputs=example_prompts,
+        outputs=prompt,
+        show_progress=False,
+        queue=False
+    )
+
+    # def bg_gallery_selected(gal, evt: gr.SelectData):
+    #     return gal[evt.index]['name']
+
+    # bg_gallery.select(
+    #     fn=bg_gallery_selected,
+    #     inputs=bg_gallery,
+    #     outputs=input_bg
+    # )
+
+block.launch(server_name='0.0.0.0')

xformers-0.0.28.post3.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f227dabb9841f5235c452c988b459ebae1c00a1c7d2d53d4ad1335318807c982
+size 4842597