Upload 4 files

README.md

@@ -1,12 +1,118 @@
 ---
-title: IClight Demo
-emoji: ⚡
+title: IC-Light Background Conditional Relighting Demo
+emoji: ✨
 colorFrom: blue
-colorTo: blue
+colorTo: purple
 sdk: gradio
-sdk_version: 5.31.0
+sdk_version: 3.41.2
 app_file: app.py
 pinned: false
+license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# IC-Light Background Conditional Relighting Demo
+
+这是一个基于IC-Light的背景条件重新打光演示应用，适用于Hugging Face Spaces。
+
+## 功能特性
+
+- **背景条件重新打光**: 使用前景图像和背景图像进行智能重新打光
+- **多种光照模式**: 支持上传背景图像或选择预设光照方向（左、右、上、下、环境光）
+- **自动背景移除**: 自动移除前景图像的背景
+- **高分辨率支持**: 支持高分辨率图像生成和细化
+- **实时预览**: 快速生成预览结果
+
+## 使用方法
+
+### 基本使用
+
+1. **上传前景图像**: 上传包含人物或物体的图像
+2. **选择背景源**:
+   - `Use Background Image`: 上传自定义背景图像
+   - `Left Light/Right Light/Top Light/Bottom Light`: 使用预设的方向性光照
+   - `Ambient`: 使用环境光照
+3. **输入提示词**: 描述期望的光照效果，如 "beautiful woman, cinematic lighting"
+4. **点击"✨ Relight Image"**: 生成重新打光的结果
+
+### 高级设置
+
+- **图像尺寸**: 调整输出图像的宽度和高度
+- **生成步数**: 控制生成质量（更多步数 = 更高质量，但更慢）
+- **CFG Scale**: 控制提示词的影响强度
+- **高分辨率缩放**: 启用高分辨率细化
+- **种子值**: 控制随机性，相同种子产生相同结果
+
+## 技术实现
+
+### 模型架构
+
+- **基础模型**: Stable Diffusion 1.5 (realistic-vision-v51)
+- **IC-Light权重**: iclight_sd15_fbc.safetensors (背景条件模型)
+- **背景移除**: RMBG-1.4 或简化的边缘检测算法
+
+### 核心功能
+
+1. **UNet修改**: 将输入通道从4扩展到12，以支持前景和背景条件
+2. **条件编码**: 将前景和背景图像编码为潜在空间条件
+3. **提示词编码**: 支持长提示词的分块编码
+4. **多阶段生成**: 支持低分辨率生成 + 高分辨率细化
+
+## 依赖项
+
+```
+diffusers==0.27.2
+transformers==4.36.2
+torch
+gradio==3.41.2
+pillow==10.2.0
+safetensors
+numpy
+scipy
+```
+
+## 部署到Hugging Face Spaces
+
+1. 创建新的Gradio Space
+2. 上传所有文件到Space
+3. 确保requirements.txt包含所有依赖
+4. Space会自动构建和部署
+
+## 示例用法
+
+### 人像重新打光
+```
+前景: 人像照片
+背景: 选择"Left Light"
+提示词: "beautiful woman, cinematic lighting"
+```
+
+### 产品摄影
+```
+前景: 产品图片
+背景: 上传工作室背景
+提示词: "product photography, professional lighting"
+```
+
+### 艺术创作
+```
+前景: 任意物体
+背景: 选择"Top Light"
+提示词: "dramatic lighting, artistic photography"
+```
+
+## 注意事项
+
+- 首次运行会下载IC-Light模型文件（约2GB）
+- GPU环境下运行效果最佳
+- 背景移除功能可能需要手动调整
+- 生成时间取决于图像尺寸和步数设置
+
+## 许可证
+
+本项目基于IC-Light官方实现，遵循相应的开源许可证。
+
+## 相关链接
+
+- [IC-Light官方仓库](https://github.com/lllyasviel/IC-Light)
+- [IC-Light论文](https://openreview.net/forum?id=u1cQYxRI1H)
+- [Hugging Face Spaces](https://huggingface.co/spaces) 

app.py

@@ -1,28 +1,436 @@
-
+import os
+import math
 import gradio as gr
-from PIL import Image
+import numpy as np
 import torch
-from ic_light.infer import infer_ic_light  # 假设IC-Light提供一个infer方法
+import safetensors.torch as sf
 
-def run_ic_light(foreground_img, background_img, model_type="fbc+fc"):
-    result = infer_ic_light(
-        foreground_path=foreground_img.name,
-        background_path=background_img.name,
-        model_type=model_type
-    )
-    return Image.open(result)
-
-iface = gr.Interface(
-    fn=run_ic_light,
-    inputs=[
-        gr.Image(type="file", label="Foreground Image"),
-        gr.Image(type="file", label="Background Image"),
-        gr.Radio(choices=["fc", "fbc", "fbc+fc"], label="Model Type", value="fbc+fc")
-    ],
-    outputs=gr.Image(label="Relit Image"),
-    title="IC-Light Background Conditional Relighting",
-    description="Upload a foreground and a background image to perform IC-Light based relighting using background conditions."
+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 enum import Enum
+from torch.hub import download_url_to_file
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
+
+# Try to import RMBG, fallback to local implementation
+try:
+    from transformers import pipeline
+    rmbg_pipeline = pipeline("image-segmentation", model="briaai/RMBG-1.4")
+    USE_RMBG_PIPELINE = True
+except Exception as e:
+    print(f"Failed to load RMBG pipeline: {e}")
+    USE_RMBG_PIPELINE = False
+    try:
+        from briarmbg import BriaRMBG, simple_background_removal
+    except:
+        # Inline simple background removal
+        def simple_background_removal(image):
+            if isinstance(image, np.ndarray):
+                img = image
+            else:
+                img = np.array(image)
+            
+            # Simple fallback - return full mask
+            gray = np.mean(img, axis=2)
+            mask = np.ones_like(gray)
+            return mask
+
+# Model setup
+sd15_name = 'stablediffusionapi/realistic-vision-v51'
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+print(f"Using device: {device}")
+print("Loading models...")
+
+# Initialize models
+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")
+
+# Modify UNet for IC-Light
+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 IC-Light weights
+model_path = './iclight_sd15_fbc.safetensors'
+if not os.path.exists(model_path):
+    print("Downloading IC-Light model...")
+    download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fbc.safetensors', dst=model_path)
+
+sd_offset = sf.load_file(model_path)
+sd_origin = unet.state_dict()
+sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
+unet.load_state_dict(sd_merged, strict=True)
+del sd_offset, sd_origin, sd_merged
+
+# Move models to device
+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)
+
+# Set attention processors
+unet.set_attn_processor(AttnProcessor2_0())
+vae.set_attn_processor(AttnProcessor2_0())
+
+# Scheduler
+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=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=scheduler,
+    safety_checker=None,
+    requires_safety_checker=False,
+    feature_extractor=None,
+    image_encoder=None
 )
 
+print("Models loaded successfully!")
+
+@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
+    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)
+    new_width = int(original_width * scale_factor)
+    new_height = int(original_height * scale_factor)
+    pil_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
+    left = (new_width - target_width) / 2
+    top = (new_height - target_height) / 2
+    right = (new_width + target_width) / 2
+    bottom = (new_height + target_height) / 2
+    pil_image = pil_image.crop((left, top, right, bottom))
+    return np.array(pil_image)
+
+def resize_without_crop(image, target_width, target_height):
+    pil_image = Image.fromarray(image)
+    pil_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
+    return np.array(pil_image)
+
+@torch.inference_mode()
+def run_rmbg(img, sigma=0.0):
+    # Simplified background removal
+    if USE_RMBG_PIPELINE:
+        # Using transformers pipeline
+        try:
+            result = rmbg_pipeline(Image.fromarray(img))
+            mask = np.array(result['mask'])
+            if len(mask.shape) == 3:
+                mask = mask[:, :, 0]
+            mask = mask.astype(np.float32) / 255.0
+        except Exception as e:
+            print(f"RMBG pipeline failed: {e}, using fallback")
+            mask = simple_background_removal(img)
+    else:
+        # Using simple background removal
+        mask = simple_background_removal(img)
+    
+    # Apply sigma smoothing
+    if sigma > 0:
+        try:
+            from scipy import ndimage
+            mask = ndimage.gaussian_filter(mask, sigma=sigma)
+        except ImportError:
+            # Fallback if scipy is not available
+            pass
+    
+    # Create RGBA output
+    result = np.dstack((img, (mask * 255).astype(np.uint8)))
+    return img, mask
+
+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"
+
+@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):
+    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 ValueError('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)
+    
+    if highres_scale > 1.0:
+        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)
+
+    return pixels, [fg, bg]
+
+@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)
+    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]
+    return results + extra_images
+
+# Quick prompts for easy testing
+quick_prompts = [
+    'beautiful woman, cinematic lighting',
+    'handsome man, cinematic lighting',
+    'beautiful woman, natural lighting',
+    'handsome man, natural lighting',
+    'beautiful woman, neo punk lighting, cyberpunk',
+    'handsome man, neo punk lighting, cyberpunk',
+]
+quick_prompts = [[x] for x in quick_prompts]
+
+# Gradio Interface
+def create_demo():
+    with gr.Blocks(title="IC-Light Background Conditional Relighting") as demo:
+        gr.Markdown("## IC-Light: Relighting with Foreground and Background Condition")
+        gr.Markdown("Upload a foreground image and background image (or choose lighting direction) to perform relighting.")
+        
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    input_fg = gr.Image(source='upload', type="numpy", label="Foreground Image", height=400)
+                    input_bg = gr.Image(source='upload', type="numpy", label="Background Image", height=400)
+                
+                prompt = gr.Textbox(label="Prompt", value="beautiful woman, cinematic lighting")
+                bg_source = gr.Radio(
+                    choices=[e.value for e in BGSource],
+                    value=BGSource.UPLOAD.value,
+                    label="Background Source"
+                )
+                
+                example_prompts = gr.Dataset(
+                    samples=quick_prompts, 
+                    label='Quick Prompts', 
+                    components=[prompt]
+                )
+                
+                relight_button = gr.Button(value="✨ Relight Image", variant="primary")
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    with gr.Row():
+                        num_samples = gr.Slider(label="Number of Images", minimum=1, maximum=4, value=1, step=1)
+                        seed = gr.Number(label="Seed", value=12345, precision=0)
+                    with gr.Row():
+                        image_width = gr.Slider(label="Width", minimum=256, maximum=1024, value=512, step=64)
+                        image_height = gr.Slider(label="Height", minimum=256, maximum=1024, value=640, step=64)
+                    with gr.Row():
+                        steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1)
+                        cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, value=7.0, step=0.1)
+                    with gr.Row():
+                        highres_scale = gr.Slider(label="Highres Scale", minimum=1.0, maximum=2.0, value=1.5, step=0.1)
+                        highres_denoise = gr.Slider(label="Highres Denoise", minimum=0.1, maximum=0.9, value=0.5, step=0.1)
+                    
+                    a_prompt = gr.Textbox(label="Additional 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(label='Results', height=600, object_fit='contain')
+        
+        # Event handlers
+        inputs = [input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source]
+        relight_button.click(fn=process_relight, inputs=inputs, outputs=[result_gallery])
+        example_prompts.click(lambda x: x[0], inputs=example_prompts, outputs=prompt, show_progress=False, queue=False)
+        
+        # Examples
+        gr.Examples(
+            examples=[
+                ["examples/person1.jpg", "examples/bg1.jpg", "beautiful woman, cinematic lighting", "Use Background Image"],
+                ["examples/person2.jpg", None, "handsome man, dramatic lighting", "Left Light"],
+            ],
+            inputs=[input_fg, input_bg, prompt, bg_source],
+            outputs=[result_gallery],
+            fn=process_relight,
+            cache_examples=False,
+        )
+    
+    return demo
+
 if __name__ == "__main__":
-    iface.launch()
+    demo = create_demo()
+    demo.queue(max_size=20)
+    demo.launch(
+        server_name='0.0.0.0', 
+        server_port=7860,
+        show_error=True,
+        share=False
+    )

briarmbg.py

@@ -0,0 +1,176 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
+import numpy as np
+from PIL import Image
+
+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
+
+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)
+        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):
+        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
+
+class BriaRMBG(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, config: dict = {"in_ch": 3, "out_ch": 1}):
+        super(BriaRMBG, self).__init__()
+        self.config = config
+        # Simplified architecture for fallback
+        self.stage1 = RSU7(3, 32, 64)
+        self.stage2 = RSU7(64, 32, 128)
+        self.stage3 = RSU7(128, 64, 256)
+        self.stage4 = RSU7(256, 128, 512)
+        
+        # Decoder
+        self.stage4d = RSU7(1024, 128, 256)
+        self.stage3d = RSU7(512, 64, 128)
+        self.stage2d = RSU7(256, 32, 64)
+        self.stage1d = RSU7(128, 16, 64)
+        
+        self.side1 = nn.Conv2d(64, 1, 3, padding=1)
+        self.side2 = nn.Conv2d(64, 1, 3, padding=1)
+        self.side3 = nn.Conv2d(128, 1, 3, padding=1)
+        self.side4 = nn.Conv2d(256, 1, 3, padding=1)
+        self.side5 = nn.Conv2d(512, 1, 3, padding=1)
+        self.side6 = nn.Conv2d(256, 1, 3, padding=1)
+        
+        self.outconv = nn.Conv2d(6, 1, 1)
+
+    def forward(self, x):
+        hx = x
+        
+        # Encoder
+        hx1 = self.stage1(hx)
+        hx = F.max_pool2d(hx1, 2, stride=2, ceil_mode=True)
+        
+        hx2 = self.stage2(hx)
+        hx = F.max_pool2d(hx2, 2, stride=2, ceil_mode=True)
+        
+        hx3 = self.stage3(hx)
+        hx = F.max_pool2d(hx3, 2, stride=2, ceil_mode=True)
+        
+        hx4 = self.stage4(hx)
+        hx = F.max_pool2d(hx4, 2, stride=2, ceil_mode=True)
+        
+        # Decoder
+        hx4d = self.stage4d(torch.cat((hx, 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 outputs
+        side1 = self.side1(hx1d)
+        side2 = _upsample_like(self.side2(hx2d), side1)
+        side3 = _upsample_like(self.side3(hx3d), side1)
+        side4 = _upsample_like(self.side4(hx4d), side1)
+        side5 = _upsample_like(self.side5(hx), side1)
+        side6 = _upsample_like(self.side6(hx4d), side1)
+        
+        # Fusion
+        out = self.outconv(torch.cat((side1, side2, side3, side4, side5, side6), 1))
+        
+        return torch.sigmoid(out)
+
+def simple_background_removal(image):
+    """
+    Simple fallback background removal using edge detection and thresholding
+    """
+    if isinstance(image, np.ndarray):
+        img = image
+    else:
+        img = np.array(image)
+    
+    # Convert to grayscale
+    gray = np.mean(img, axis=2)
+    
+    # Simple edge detection
+    from scipy import ndimage
+    edges = ndimage.sobel(gray)
+    
+    # Create a simple mask based on intensity
+    mask = np.ones_like(gray)
+    
+    # Simple thresholding - assume foreground is in center and has more edges
+    h, w = gray.shape
+    center_mask = np.zeros_like(gray)
+    center_mask[h//4:3*h//4, w//4:3*w//4] = 1
+    
+    # Combine edge information with center bias
+    mask = (edges > np.percentile(edges, 70)) * center_mask
+    mask = ndimage.binary_fill_holes(mask)
+    mask = ndimage.gaussian_filter(mask.astype(float), sigma=2)
+    
+    return mask 

requirements.txt

@@ -1,12 +1,16 @@
-
-gradio==3.41.2
-torch
-Pillow
-opencv-python
-einops
-transformers==4.36.2
+torch>=2.0.0
+torchvision
 diffusers==0.27.2
+transformers==4.36.2
+accelerate
+opencv-python-headless
 safetensors
+pillow==10.2.0
+einops
 peft
+gradio==3.41.2
 protobuf==3.20
-git+https://github.com/lllyasviel/IC-Light.git
+numpy
+scipy
+huggingface_hub
+xformers