Update app_2.py

app_2.py

@@ -54,7 +54,13 @@ import torch
 from PIL import Image
 
 
-
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
 
 # Load
 
@@ -351,7 +357,7 @@ def encode_prompt_pair(positive_prompt, negative_prompt):
 def infer(
     prompt,
     image,
-    do_rembg=False,
+    do_rembg=True,
     seed=42,
     randomize_seed=False,
     guidance_scale=3.0,
@@ -360,8 +366,7 @@ def infer(
     negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
     progress=gr.Progress(track_tqdm=True),
 ):
-    # if do_rembg:
-    #     remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, device)
+    remove_bg_fn = lambda x: remove_bg(x, rmbg, transform_image, device)
     # else:
     #     remove_bg_fn = None
     if randomize_seed:
@@ -377,7 +382,7 @@ def infer(
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
         seed=seed,
-        remove_bg_fn=None,
+        remove_bg_fn=remove_bg_fn,
         reference_conditioning_scale=reference_conditioning_scale,
         negative_prompt=negative_prompt,
         device=device,