Update app.py

app.py

@@ -19,8 +19,7 @@ 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
+    rmbg_pipeline = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
 except Exception as e:
     print(f"Failed to load RMBG pipeline: {e}")
     USE_RMBG_PIPELINE = False
@@ -47,10 +46,10 @@ 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")
+tokenizer = CLIPTokenizer.from_pretrained(sd15_name, subfolder="tokenizer", trust_remote_code=True)
+text_encoder = CLIPTextModel.from_pretrained(sd15_name, subfolder="text_encoder", trust_remote_code=True)
+vae = AutoencoderKL.from_pretrained(sd15_name, subfolder="vae", trust_remote_code=True)
+unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet", trust_remote_code=True)
 
 # Modify UNet for IC-Light
 with torch.no_grad():