diff --git a/Caddyfile b/Caddyfile index 379004e..10c40c7 100644 --- a/Caddyfile +++ b/Caddyfile @@ -1,15 +1,122 @@ -# Caddy 2 Configuration for AI Photo Edit -# Simple version - frontend nginx handles API routing internally +# Caddy 2 Configuration for EditmaskwithAI +# ========================================== +# +# SETUP INSTRUCTIONS: +# 1. Replace 'your-subdomain.yourdomain.com' with your actual domain +# 2. Make sure DNS CNAME record points to your server +# 3. Ensure ports 80 and 443 are open (Caddy handles SSL automatically) +# 4. The frontend runs on port 3080 by default (docker-compose) +# +# Common Issues: +# - "Connection refused": Check if the frontend container is running +# - "Bad gateway": Check if localhost:3080 is accessible +# - "SSL error": Make sure ports 80/443 are open for Let's Encrypt -# Option 1: With domain name -ai-photo-edit.yourdomain.com { - reverse_proxy localhost:3080 +# ============================================ +# OPTION 1: Domain with automatic HTTPS (recommended) +# ============================================ +# Replace with your actual domain +your-subdomain.yourdomain.com { + # Reverse proxy to frontend (nginx serves both frontend and proxies API) + reverse_proxy localhost:3080 { + # Health checks + health_uri /health + health_interval 30s + health_timeout 10s + + # Headers for proper proxying + header_up Host {upstream_hostport} + header_up X-Real-IP {remote_host} + header_up X-Forwarded-For {remote_host} + header_up X-Forwarded-Proto {scheme} + } + + # Enable compression + encode gzip zstd + + # Logging (optional - uncomment for debugging) + # log { + # output file /var/log/caddy/access.log + # format json + # } } -# Option 2: IP:Port (comment out option 1 if using this) +# ============================================ +# OPTION 2: IP address or localhost (no HTTPS) +# ============================================ +# Uncomment this block and comment out Option 1 if you don't have a domain +# or want to test locally + # :8080 { -# reverse_proxy localhost:3080 +# reverse_proxy localhost:3080 { +# header_up Host {upstream_hostport} +# header_up X-Real-IP {remote_host} +# header_up X-Forwarded-For {remote_host} +# } +# encode gzip zstd # } -# Note: You don't need to specify /projects, /edits, etc. -# The frontend's nginx is already configured to proxy those to the backend internally. +# ============================================ +# OPTION 3: Multiple subdomains +# ============================================ +# If you want both www and non-www versions + +# yourdomain.com, www.yourdomain.com { +# reverse_proxy localhost:3080 { +# header_up Host {upstream_hostport} +# header_up X-Real-IP {remote_host} +# header_up X-Forwarded-For {remote_host} +# header_up X-Forwarded-Proto {scheme} +# } +# encode gzip zstd +# } + +# ============================================ +# OPTION 4: Behind another reverse proxy (Cloudflare, etc.) +# ============================================ +# Use this if Caddy is behind Cloudflare or another proxy + +# your-subdomain.yourdomain.com { +# # Trust proxy headers from upstream +# servers { +# trusted_proxies static 173.245.48.0/20 103.21.244.0/22 103.22.200.0/22 103.31.4.0/22 141.101.64.0/18 108.162.192.0/18 190.93.240.0/20 188.114.96.0/20 197.234.240.0/22 198.41.128.0/17 162.158.0.0/15 104.16.0.0/13 104.24.0.0/14 172.64.0.0/13 131.0.72.0/22 +# } +# +# reverse_proxy localhost:3080 { +# header_up Host {upstream_hostport} +# header_up X-Real-IP {http.request.header.CF-Connecting-IP} +# header_up X-Forwarded-For {http.request.header.CF-Connecting-IP} +# header_up X-Forwarded-Proto {scheme} +# } +# encode gzip zstd +# } + +# ============================================ +# TROUBLESHOOTING +# ============================================ +# +# 1. Check Caddy logs: +# docker logs caddy +# OR: journalctl -u caddy -f +# +# 2. Test backend connectivity: +# curl -I http://localhost:3080 +# +# 3. Check DNS resolution: +# dig your-subdomain.yourdomain.com +# nslookup your-subdomain.yourdomain.com +# +# 4. Verify ports are open: +# sudo netstat -tlnp | grep -E ':(80|443|3080)' +# +# 5. Check firewall: +# sudo ufw status +# sudo iptables -L -n +# +# 6. For Let's Encrypt issues: +# - Ensure ports 80 and 443 are accessible from internet +# - Check if domain resolves to your server's IP +# - Try: caddy validate --config /path/to/Caddyfile +# +# 7. Force reload Caddy config: +# caddy reload --config /path/to/Caddyfile diff --git a/backend/app/routers/tools.py b/backend/app/routers/tools.py index 062e2bf..c6eed00 100644 --- a/backend/app/routers/tools.py +++ b/backend/app/routers/tools.py @@ -128,53 +128,171 @@ async def inpaint_base64(request: InpaintRequest): raise HTTPException(status_code=500, detail=str(e)) +class RemoveBackgroundRequestV2(BaseModel): + image: str # Base64 encoded image + model: Optional[str] = "auto" # "auto", "u2net", "rembg", "birefnet" + + @router.post("/remove-background-base64") async def remove_background_base64(request: RemoveBackgroundRequest): """ - Remove background from a base64 encoded image using rembg with BiRefNet. - BiRefNet is state-of-the-art for background removal (better than u2net). + Remove background from a base64 encoded image. + Tries multiple methods: U2Net (direct), rembg with BiRefNet, rembg default. Returns base64 encoded PNG with transparent background. Used by miniPaint frontend. """ - try: - from rembg import remove, new_session - except ImportError: - raise HTTPException( - status_code=500, - detail="rembg not installed. Run: pip install rembg" - ) - try: # Decode base64 image image_bytes = base64.b64decode(request.image) + img = Image.open(BytesIO(image_bytes)).convert('RGB') - # Use BiRefNet model for best quality (state-of-the-art) - # Falls back to default model if BiRefNet not available + result_bytes = None + method_used = None + + # Try U2Net first (direct implementation, no rembg dependency issues) try: - session = new_session("birefnet-general") - result_bytes = remove(image_bytes, session=session) - except Exception: - # Fallback to default model - result_bytes = remove(image_bytes) + result_bytes = await _remove_background_u2net(img) + method_used = "u2net" + except Exception as e: + print(f"U2Net failed: {e}") + + # Fall back to rembg if U2Net failed + if result_bytes is None: + try: + from rembg import remove, new_session + try: + session = new_session("birefnet-general") + result_bytes = remove(image_bytes, session=session) + method_used = "birefnet" + except Exception: + result_bytes = remove(image_bytes) + method_used = "rembg-default" + except ImportError: + pass + except Exception as e: + print(f"rembg failed: {e}") + + if result_bytes is None: + raise HTTPException( + status_code=500, + detail="No background removal method available. Install u2net or rembg." + ) # Convert result to base64 result_b64 = base64.b64encode(result_bytes).decode('utf-8') # Get dimensions - img = Image.open(BytesIO(result_bytes)) + result_img = Image.open(BytesIO(result_bytes)) return { "result": result_b64, - "width": img.width, - "height": img.height + "width": result_img.width, + "height": result_img.height, + "method": method_used } + except HTTPException: + raise except Exception as e: import traceback traceback.print_exc() raise HTTPException(status_code=500, detail=str(e)) +# Global U2Net model cache +_u2net_model = None + + +async def _remove_background_u2net(img: Image.Image) -> bytes: + """ + Remove background using U2Net model directly. + This avoids rembg dependency issues while providing good quality. + """ + global _u2net_model + + import torch + from pathlib import Path + + # Check for U2Net model + models_dir = Path('/app/data/models') + u2net_path = models_dir / 'u2net.pth' + + # Also check alternative names + if not u2net_path.exists(): + for alt_name in ['u2net.onnx', 'u2netp.pth', 'u2net_human_seg.pth']: + alt_path = models_dir / alt_name + if alt_path.exists(): + u2net_path = alt_path + break + + if not u2net_path.exists(): + raise FileNotFoundError( + f"U2Net model not found at {u2net_path}. " + "Download from: https://github.com/xuebinqin/U-2-Net" + ) + + # Load model if not cached + if _u2net_model is None: + print(f"Loading U2Net model from {u2net_path}") + + if str(u2net_path).endswith('.onnx'): + # Use ONNX runtime + import onnxruntime as ort + _u2net_model = ort.InferenceSession(str(u2net_path)) + else: + # Use PyTorch + from app.services.u2net_model import U2NET + _u2net_model = U2NET(3, 1) + _u2net_model.load_state_dict(torch.load(str(u2net_path), map_location='cpu')) + _u2net_model.eval() + + print("U2Net model loaded") + + # Preprocess image + img_np = np.array(img) + original_size = img.size + + # Resize to model input size + input_size = 320 + img_resized = img.resize((input_size, input_size), Image.Resampling.BILINEAR) + img_np = np.array(img_resized).astype(np.float32) + + # Normalize + img_np = img_np / 255.0 + img_np = (img_np - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] + img_np = img_np.transpose(2, 0, 1) # HWC to CHW + img_np = np.expand_dims(img_np, 0) # Add batch dimension + + # Run inference + if hasattr(_u2net_model, 'run'): + # ONNX runtime + input_name = _u2net_model.get_inputs()[0].name + outputs = _u2net_model.run(None, {input_name: img_np}) + mask = outputs[0][0, 0] + else: + # PyTorch + with torch.no_grad(): + input_tensor = torch.from_numpy(img_np).float() + d1, d2, d3, d4, d5, d6, d7 = _u2net_model(input_tensor) + mask = d1[0, 0].numpy() + + # Post-process mask + mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8) + mask = (mask * 255).astype(np.uint8) + + # Resize mask back to original size + mask_img = Image.fromarray(mask).resize(original_size, Image.Resampling.BILINEAR) + + # Apply mask to original image + result = img.convert('RGBA') + result.putalpha(mask_img) + + # Save to bytes + buffer = BytesIO() + result.save(buffer, format='PNG') + return buffer.getvalue() + + @router.post("/remove-background") async def remove_background( project_id: Optional[int] = Form(None), diff --git a/backend/app/services/u2net_model.py b/backend/app/services/u2net_model.py new file mode 100644 index 0000000..fa269b1 --- /dev/null +++ b/backend/app/services/u2net_model.py @@ -0,0 +1,500 @@ +""" +U2Net Model Definition for Background Removal +Based on: https://github.com/xuebinqin/U-2-Net + +This is a simplified implementation that works with the standard U2Net weights. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class REBNCONV(nn.Module): + def __init__(self, in_ch=3, out_ch=3, dirate=1): + super(REBNCONV, self).__init__() + + self.conv_s1 = nn.Conv2d(in_ch, out_ch, 3, padding=1*dirate, dilation=1*dirate) + 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', align_corners=False) + return src + + +class RSU7(nn.Module): + def __init__(self, in_ch=3, mid_ch=12, out_ch=3): + super(RSU7, 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.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 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 + + +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 + + +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 + + +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 U2NET(nn.Module): + def __init__(self, in_ch=3, out_ch=1): + super(U2NET, self).__init__() + + self.stage1 = RSU7(in_ch, 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 + + # stage 1 + hx1 = self.stage1(hx) + 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) + + d2 = self.side2(hx2d) + d2 = _upsample_like(d2, d1) + + d3 = self.side3(hx3d) + d3 = _upsample_like(d3, d1) + + d4 = self.side4(hx4d) + d4 = _upsample_like(d4, d1) + + d5 = self.side5(hx5d) + d5 = _upsample_like(d5, d1) + + d6 = self.side6(hx6) + d6 = _upsample_like(d6, d1) + + d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1)) + + return torch.sigmoid(d0), torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid(d4), torch.sigmoid(d5), torch.sigmoid(d6) + + +class U2NETP(nn.Module): + """Smaller/faster U2Net variant (u2netp)""" + + def __init__(self, in_ch=3, out_ch=1): + super(U2NETP, self).__init__() + + self.stage1 = RSU7(in_ch, 16, 64) + self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage2 = RSU6(64, 16, 64) + self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage3 = RSU5(64, 16, 64) + self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage4 = RSU4(64, 16, 64) + self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage5 = RSU4F(64, 16, 64) + self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage6 = RSU4F(64, 16, 64) + + # decoder + self.stage5d = RSU4F(128, 16, 64) + self.stage4d = RSU4(128, 16, 64) + self.stage3d = RSU5(128, 16, 64) + self.stage2d = RSU6(128, 16, 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(64, out_ch, 3, padding=1) + self.side4 = nn.Conv2d(64, out_ch, 3, padding=1) + self.side5 = nn.Conv2d(64, out_ch, 3, padding=1) + self.side6 = nn.Conv2d(64, out_ch, 3, padding=1) + + self.outconv = nn.Conv2d(6*out_ch, out_ch, 1) + + def forward(self, x): + hx = x + + hx1 = self.stage1(hx) + hx = self.pool12(hx1) + + hx2 = self.stage2(hx) + hx = self.pool23(hx2) + + hx3 = self.stage3(hx) + hx = self.pool34(hx3) + + hx4 = self.stage4(hx) + hx = self.pool45(hx4) + + hx5 = self.stage5(hx) + hx = self.pool56(hx5) + + hx6 = self.stage6(hx) + hx6up = _upsample_like(hx6, hx5) + + 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)) + + d1 = self.side1(hx1d) + + d2 = self.side2(hx2d) + d2 = _upsample_like(d2, d1) + + d3 = self.side3(hx3d) + d3 = _upsample_like(d3, d1) + + d4 = self.side4(hx4d) + d4 = _upsample_like(d4, d1) + + d5 = self.side5(hx5d) + d5 = _upsample_like(d5, d1) + + d6 = self.side6(hx6) + d6 = _upsample_like(d6, d1) + + d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1)) + + return torch.sigmoid(d0), torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid(d4), torch.sigmoid(d5), torch.sigmoid(d6) diff --git a/frontend/src/css/popup.css b/frontend/src/css/popup.css index 2ac6e82..5a72e74 100644 --- a/frontend/src/css/popup.css +++ b/frontend/src/css/popup.css @@ -245,6 +245,143 @@ margin-right: auto; } +/* Shape/Library Tabs */ +#popups .popup .shape-tabs { + display: flex; + gap: 0; + margin-bottom: 1rem; + border-bottom: 2px solid var(--border-color); +} + +#popups .popup .shape-tab { + padding: 0.8rem 1.5rem; + background: transparent; + border: none; + border-bottom: 2px solid transparent; + margin-bottom: -2px; + cursor: pointer; + color: var(--text-color-muted); + font-size: 1rem; + transition: color 0.2s, border-color 0.2s; +} + +#popups .popup .shape-tab:hover { + color: var(--text-color); +} + +#popups .popup .shape-tab.active { + color: var(--link-color); + border-bottom-color: var(--link-color); +} + +#popups .popup .library-loading { + text-align: center; + padding: 2rem; + color: var(--text-color-muted); +} + +/* My Library Browser Styles */ +#popups .popup .library-browser { + max-height: calc(60vh - 100px); + overflow-y: auto; +} + +#popups .popup .library-categories { + display: flex; + flex-direction: column; + gap: 1.5rem; +} + +#popups .popup .library-category h3 { + color: var(--text-color); + font-size: 1.4rem; + margin-bottom: 0.8rem; + padding-bottom: 0.4rem; + border-bottom: 1px solid var(--border-color); +} + +#popups .popup .library-items { + display: flex; + flex-wrap: wrap; + gap: 1rem; +} + +#popups .popup .library-item { + display: flex; + flex-direction: column; + align-items: center; + width: 120px; + padding: 0.8rem; + background: var(--input-background-color); + border: 1px solid var(--border-color); + border-radius: 4px; + cursor: pointer; + transition: background 0.2s, border-color 0.2s; +} + +#popups .popup .library-item:hover { + background: var(--input-background-color-hover); + border-color: var(--link-color); +} + +#popups .popup .library-item img { + width: 100px; + height: 80px; + object-fit: contain; + background: repeating-conic-gradient(#666 0% 25%, #888 0% 50%) 50% / 10px 10px; + border-radius: 2px; + margin-bottom: 0.5rem; +} + +#popups .popup .library-item-name { + font-size: 0.85rem; + text-align: center; + color: var(--text-color); + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; + width: 100%; + margin-bottom: 0.5rem; +} + +#popups .popup .library-item-actions { + display: flex; + gap: 0.5rem; +} + +#popups .popup .library-item-actions button { + font-size: 0.75rem; + padding: 0.3rem 0.6rem; + background: var(--button-background-color); + border: 1px solid var(--border-color); + border-radius: 3px; + cursor: pointer; + color: var(--text-color); +} + +#popups .popup .library-item-actions .insert-btn { + background: #2a6d2a; +} + +#popups .popup .library-item-actions .insert-btn:hover { + background: #3a8d3a; +} + +#popups .popup .library-item-actions .delete-btn { + background: #6d2a2a; +} + +#popups .popup .library-item-actions .delete-btn:hover { + background: #8d3a3a; +} + +/* Library empty state */ +#popups .popup .library-empty { + text-align: center; + padding: 2rem; + color: var(--text-color-muted); +} + @media screen and (max-width:500px){ #popups .popup { max-height: calc(80vh - 20px); /* mobile phones has bottom menu */ diff --git a/frontend/src/js/modules/file/my_library.js b/frontend/src/js/modules/file/my_library.js index 4d60869..40d9574 100644 --- a/frontend/src/js/modules/file/my_library.js +++ b/frontend/src/js/modules/file/my_library.js @@ -334,7 +334,7 @@ class File_my_library_class { title: 'My Library (' + assets.length + ' assets)', params: [], html: html, - className: 'library-dialog', + className: 'wide', on_load: function(el) { // Add click handlers el.querySelectorAll('.insert-btn').forEach(function(btn) { diff --git a/frontend/src/js/tools/ai_inpaint.js b/frontend/src/js/tools/ai_inpaint.js index 1bff763..54e2f56 100644 --- a/frontend/src/js/tools/ai_inpaint.js +++ b/frontend/src/js/tools/ai_inpaint.js @@ -49,14 +49,20 @@ class Ai_inpaint_class extends Base_tools_class { } var settings = { - title: 'AI Inpaint', + title: 'AI Edit Selection', params: [ + { + name: "mode", + title: "Edit Mode:", + value: "inpaint", + values: ["inpaint", "transform"] + }, { name: "prompt", - title: "Describe what you want:", + title: "AI Inpaint - Describe replacement:", type: "textarea", value: "", - placeholder: "e.g., 'a red rose', 'remove the object', 'blue sky with clouds'" + placeholder: "AI will REPLACE the selection with what you describe.\nExamples: 'a red rose', 'empty background', 'blue sky'" }, { name: "negative_prompt", @@ -66,21 +72,164 @@ class Ai_inpaint_class extends Base_tools_class { }, { name: "strength", - title: "Edit Strength:", + title: "AI Edit Strength:", type: "range", value: 80, range: [1, 100], step: 1 + }, + { + name: "scale", + title: "Transform - Scale %:", + type: "range", + value: 100, + range: [10, 200], + step: 5 } ], + on_load: function(el) { + // Add info text + var infoDiv = document.createElement('div'); + infoDiv.className = 'ai-inpaint-info'; + infoDiv.innerHTML = '

' + + 'Inpaint Mode: AI replaces the selected area with generated content.
' + + 'Transform Mode: Scale, shrink, or enlarge the selection without AI.
' + + 'Tip: To shrink something by 35%, use Transform mode with Scale at 65%.

'; + + var dialogContent = el.querySelector('.dialog_content'); + if (dialogContent && dialogContent.firstChild) { + dialogContent.insertBefore(infoDiv, dialogContent.firstChild); + } + }, on_finish: async function (params) { - await _this.executeInpaint(params); + if (params.mode === 'transform') { + await _this.executeTransform(params); + } else { + await _this.executeInpaint(params); + } }, }; this.POP.show(settings); } + /** + * Execute transform operation (scale without AI) + */ + async executeTransform(params) { + if (this.isProcessing) { + alertify.warning('Already processing... please wait'); + return; + } + + // Check if we have an image layer + if (config.layer.type != 'image') { + alertify.error('Please select an image layer'); + return; + } + + var maskCanvas = window.smartSelectMask?.canvas; + if (!maskCanvas) { + alertify.error('No selection mask found'); + return; + } + + this.isProcessing = true; + alertify.message('Transforming selection...'); + + try { + var layer = config.layer; + var scale = params.scale / 100; + + // Get mask bounds + var maskCtx = maskCanvas.getContext('2d'); + var imageData = maskCtx.getImageData(0, 0, maskCanvas.width, maskCanvas.height); + var minX = maskCanvas.width, minY = maskCanvas.height; + var maxX = 0, maxY = 0; + + for (var y = 0; y < maskCanvas.height; y++) { + for (var x = 0; x < maskCanvas.width; x++) { + var i = (y * maskCanvas.width + x) * 4; + if (imageData.data[i] > 128) { + minX = Math.min(minX, x); + minY = Math.min(minY, y); + maxX = Math.max(maxX, x); + maxY = Math.max(maxY, y); + } + } + } + + if (maxX <= minX || maxY <= minY) { + throw new Error('Selection is too small'); + } + + var selWidth = maxX - minX + 1; + var selHeight = maxY - minY + 1; + var centerX = minX + selWidth / 2; + var centerY = minY + selHeight / 2; + + // Extract selected pixels + var extractCanvas = document.createElement('canvas'); + extractCanvas.width = layer.width_original; + extractCanvas.height = layer.height_original; + var extractCtx = extractCanvas.getContext('2d'); + extractCtx.drawImage(layer.link, 0, 0); + extractCtx.globalCompositeOperation = 'destination-in'; + extractCtx.drawImage(maskCanvas, 0, 0); + + // Create result canvas + var resultCanvas = document.createElement('canvas'); + resultCanvas.width = layer.width_original; + resultCanvas.height = layer.height_original; + var resultCtx = resultCanvas.getContext('2d'); + + // Draw original image + resultCtx.drawImage(layer.link, 0, 0); + + // Remove original selection (create hole) + resultCtx.globalCompositeOperation = 'destination-out'; + resultCtx.drawImage(maskCanvas, 0, 0); + + // Calculate scaled dimensions + var newWidth = selWidth * scale; + var newHeight = selHeight * scale; + var newX = centerX - newWidth / 2; + var newY = centerY - newHeight / 2; + + // Draw scaled selection back + resultCtx.globalCompositeOperation = 'source-over'; + + // Create temp canvas for just the selection + var selCanvas = document.createElement('canvas'); + selCanvas.width = selWidth; + selCanvas.height = selHeight; + var selCtx = selCanvas.getContext('2d'); + selCtx.drawImage(extractCanvas, minX, minY, selWidth, selHeight, 0, 0, selWidth, selHeight); + + // Draw scaled + resultCtx.drawImage(selCanvas, 0, 0, selWidth, selHeight, newX, newY, newWidth, newHeight); + + // Apply result + app.State.do_action( + new app.Actions.Bundle_action('transform_selection', 'Transform Selection', [ + new app.Actions.Update_layer_image_action(resultCanvas) + ]) + ); + + // Clear selection + window.smartSelectMask = null; + config.need_render = true; + + alertify.success('Transform complete! Selection scaled to ' + params.scale + '%'); + + } catch (error) { + console.error('Transform error:', error); + alertify.error('Transform failed: ' + error.message); + } finally { + this.isProcessing = false; + } + } + /** * Execute the inpainting operation */ diff --git a/frontend/src/js/tools/shape.js b/frontend/src/js/tools/shape.js index 2269df3..e48f792 100644 --- a/frontend/src/js/tools/shape.js +++ b/frontend/src/js/tools/shape.js @@ -4,6 +4,7 @@ import Base_tools_class from './../core/base-tools.js'; import Base_layers_class from './../core/base-layers.js'; import Dialog_class from './../libs/popup.js'; import GUI_tools_class from './../core/gui/gui-tools.js'; +import File_my_library_class from './../modules/file/my_library.js'; var instance = null; @@ -21,11 +22,13 @@ class Shape_class extends Base_tools_class { this.Base_layers = new Base_layers_class(); this.GUI_tools = new GUI_tools_class(); this.POP = new Dialog_class(); + this.My_library = new File_my_library_class(); this.ctx = ctx; this.name = 'shape'; this.layer = {}; this.preview_width = 150; this.preview_height = 120; + this.activeTab = 'shapes'; // 'shapes' or 'library' this.set_events(); } @@ -53,35 +56,84 @@ class Shape_class extends Base_tools_class { async show_shapes(){ var _this = this; - var html = ''; + // Build tabs HTML + var tabsHtml = '
'; + tabsHtml += ''; + tabsHtml += ''; + tabsHtml += '
'; + + // Build shapes HTML + var shapesHtml = '
'; var data = this.get_shapes(); for (var i in data) { - html += '
'; - html += ' '; + shapesHtml += ' '; - html += '
' + data[i].title + '
'; - html += '
'; + shapesHtml += '
' + data[i].title + '
'; + shapesHtml += '
'; } for (var i = 0; i < 4; i++) { - html += '
'; + shapesHtml += '
'; } + shapesHtml += ''; + + // Build library HTML placeholder + var libraryHtml = ''; var settings = { - title: 'Shapes', + title: 'Shapes & Library', className: 'wide', on_load: function (params, popup) { - var node = document.createElement("div"); - node.classList.add('flex-container'); - node.innerHTML = html; - popup.el.querySelector('.dialog_content').appendChild(node); - //events + // Add tabs + var tabsNode = document.createElement("div"); + tabsNode.innerHTML = tabsHtml; + popup.el.querySelector('.dialog_content').insertBefore(tabsNode, popup.el.querySelector('.dialog_content').firstChild); + + // Add shapes container + var shapesNode = document.createElement("div"); + shapesNode.classList.add('flex-container'); + shapesNode.innerHTML = shapesHtml; + popup.el.querySelector('.dialog_content').appendChild(shapesNode); + + // Add library container + var libraryNode = document.createElement("div"); + libraryNode.innerHTML = libraryHtml; + popup.el.querySelector('.dialog_content').appendChild(libraryNode); + + // Tab click events + var tabs = popup.el.querySelectorAll('.shape-tab'); + tabs.forEach(function(tab) { + tab.addEventListener('click', function() { + var targetTab = this.dataset.tab; + + // Update active tab + tabs.forEach(t => t.classList.remove('active')); + this.classList.add('active'); + + // Show/hide content + var shapesContent = popup.el.querySelector('.shapes-content'); + var libraryContent = popup.el.querySelector('.library-content'); + + if (targetTab === 'shapes') { + shapesContent.style.display = ''; + libraryContent.style.display = 'none'; + } else { + shapesContent.style.display = 'none'; + libraryContent.style.display = ''; + _this.loadLibraryContent(libraryContent); + } + }); + }); + + // Shape click events var targets = popup.el.querySelectorAll('.item canvas'); for (var i = 0; i < targets.length; i++) { targets[i].addEventListener('click', function (event) { - //we have click _this.GUI_tools.activate_tool(this.dataset.key); _this.POP.hide(); }); @@ -106,6 +158,89 @@ class Shape_class extends Base_tools_class { } } + /** + * Load library content into the library tab + */ + loadLibraryContent(container) { + var _this = this; + + this.My_library.getAllAssets(function(assets) { + var html = ''; + + if (assets.length === 0) { + html = '
'; + html += '

Your library is empty.

'; + html += '

Use File > My Library > Save to Library to add assets.

'; + html += '
'; + } else { + // Group by category + var categories = {}; + assets.forEach(function(asset) { + var cat = asset.category || 'General'; + if (!categories[cat]) categories[cat] = []; + categories[cat].push(asset); + }); + + html = '
'; + html += '
'; + + for (var cat in categories) { + html += '
'; + html += '

' + cat + ' (' + categories[cat].length + ')

'; + html += '
'; + + categories[cat].forEach(function(asset) { + html += '
'; + html += '' + asset.name + ''; + html += '
' + asset.name + '
'; + html += '
'; + html += ''; + html += ''; + html += '
'; + html += '
'; + }); + + html += '
'; + } + + html += '
'; + } + + container.innerHTML = html; + + // Add event handlers for library items + container.querySelectorAll('.insert-btn').forEach(function(btn) { + btn.addEventListener('click', function(e) { + e.stopPropagation(); + var id = parseInt(this.dataset.id); + _this.My_library.insertAsset(id); + _this.POP.hide(); + }); + }); + + container.querySelectorAll('.delete-btn').forEach(function(btn) { + btn.addEventListener('click', function(e) { + e.stopPropagation(); + var id = parseInt(this.dataset.id); + if (confirm('Delete this asset?')) { + _this.My_library.deleteAsset(id, function() { + _this.loadLibraryContent(container); + }); + } + }); + }); + + // Double-click to insert + container.querySelectorAll('.library-item').forEach(function(item) { + item.addEventListener('dblclick', function() { + var id = parseInt(this.dataset.id); + _this.My_library.insertAsset(id); + _this.POP.hide(); + }); + }); + }); + } + render(ctx, layer) { } diff --git a/scripts/download_u2net_model.py b/scripts/download_u2net_model.py new file mode 100644 index 0000000..3f44b75 --- /dev/null +++ b/scripts/download_u2net_model.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +""" +Download U2Net model for background removal. + +U2Net is a deep learning model for salient object detection, +commonly used for background removal tasks. + +Usage: + python download_u2net_model.py [model_type] + +Model types: + u2net - Full U2Net model (~176MB, best quality) + u2netp - Lightweight U2Net (~4MB, faster, good quality) + u2net_human_seg - Optimized for human segmentation (~176MB) + +Default: u2netp (good balance of quality and speed) +""" + +import os +import sys +import urllib.request +from pathlib import Path + +# Model URLs (from official U2Net repository releases) +MODEL_URLS = { + 'u2net': { + 'url': 'https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx', + 'filename': 'u2net.onnx', + 'size_mb': 176 + }, + 'u2netp': { + 'url': 'https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx', + 'filename': 'u2netp.onnx', + 'size_mb': 4 + }, + 'u2net_human_seg': { + 'url': 'https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx', + 'filename': 'u2net_human_seg.onnx', + 'size_mb': 176 + } +} + + +def download_with_progress(url: str, dest_path: Path, expected_size_mb: int): + """Download file with progress indicator.""" + + print(f"Downloading from: {url}") + print(f"Expected size: ~{expected_size_mb}MB") + + def progress_hook(count, block_size, total_size): + if total_size > 0: + percent = min(100, count * block_size * 100 // total_size) + downloaded_mb = count * block_size / (1024 * 1024) + total_mb = total_size / (1024 * 1024) + sys.stdout.write(f"\rProgress: {percent}% ({downloaded_mb:.1f}/{total_mb:.1f} MB)") + sys.stdout.flush() + + try: + urllib.request.urlretrieve(url, str(dest_path), progress_hook) + print("\nDownload complete!") + return True + except Exception as e: + print(f"\nDownload failed: {e}") + return False + + +def main(): + # Determine model type + model_type = 'u2netp' # Default to lightweight model + if len(sys.argv) > 1: + model_type = sys.argv[1].lower() + + if model_type not in MODEL_URLS: + print(f"Unknown model type: {model_type}") + print(f"Available models: {', '.join(MODEL_URLS.keys())}") + sys.exit(1) + + model_info = MODEL_URLS[model_type] + + # Determine models directory + # Check if running in Docker container + if os.path.exists('/app/data/models'): + models_dir = Path('/app/data/models') + else: + # Local development + script_dir = Path(__file__).parent + models_dir = script_dir.parent / 'data' / 'models' + + models_dir.mkdir(parents=True, exist_ok=True) + + dest_path = models_dir / model_info['filename'] + + # Check if already downloaded + if dest_path.exists(): + print(f"Model already exists at: {dest_path}") + print("Delete the file to re-download.") + return + + print(f"Downloading U2Net model: {model_type}") + print(f"Destination: {dest_path}") + print("") + + success = download_with_progress( + model_info['url'], + dest_path, + model_info['size_mb'] + ) + + if success: + # Create symlink for easier access + symlink_path = models_dir / 'u2net.onnx' + if not symlink_path.exists() or symlink_path.is_symlink(): + if symlink_path.is_symlink(): + symlink_path.unlink() + try: + symlink_path.symlink_to(dest_path.name) + print(f"Created symlink: {symlink_path} -> {dest_path.name}") + except OSError: + # Symlinks may not work on all systems + pass + + print(f"\nU2Net model ({model_type}) downloaded successfully!") + print(f"Location: {dest_path}") + print("\nYou can now use background removal in the application.") + else: + print("\nFailed to download model. Please try again or download manually from:") + print(f" {model_info['url']}") + print(f" Save to: {dest_path}") + sys.exit(1) + + +if __name__ == '__main__': + main()