Ashoka74 commited on
Commit
54e8f2c
·
verified ·
1 Parent(s): 398d68f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -20
app.py CHANGED
@@ -34,25 +34,12 @@ pipe = prepare_pipeline(
34
  dtype=dtype,
35
  )
36
 
37
- # remove bg
38
- birefnet = AutoModelForImageSegmentation.from_pretrained(
39
- "ZhengPeng7/BiRefNet", trust_remote_code=True
40
- )
41
- birefnet.to(device)
42
- transform_image = transforms.Compose(
43
- [
44
- transforms.Resize((1024, 1024)),
45
- transforms.ToTensor(),
46
- transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
47
- ]
48
- )
49
-
50
 
51
  @spaces.GPU()
52
  def infer(
53
  prompt,
54
  image,
55
- do_rembg=True,
56
  seed=42,
57
  randomize_seed=False,
58
  guidance_scale=3.0,
@@ -61,10 +48,6 @@ def infer(
61
  negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
62
  progress=gr.Progress(track_tqdm=True),
63
  ):
64
- if do_rembg:
65
- remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, device)
66
- else:
67
- remove_bg_fn = None
68
  if randomize_seed:
69
  seed = random.randint(0, MAX_SEED)
70
  images, preprocessed_image = run_pipeline(
@@ -77,12 +60,12 @@ def infer(
77
  num_inference_steps=num_inference_steps,
78
  guidance_scale=guidance_scale,
79
  seed=seed,
80
- remove_bg_fn=remove_bg_fn,
81
  reference_conditioning_scale=reference_conditioning_scale,
82
  negative_prompt=negative_prompt,
83
  device=device,
84
  )
85
- return images, preprocessed_image, seed
86
 
87
 
88
  # examples = [
 
34
  dtype=dtype,
35
  )
36
 
 
 
 
 
 
 
 
 
 
 
 
 
 
37
 
38
  @spaces.GPU()
39
  def infer(
40
  prompt,
41
  image,
42
+ do_rembg=False,
43
  seed=42,
44
  randomize_seed=False,
45
  guidance_scale=3.0,
 
48
  negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
49
  progress=gr.Progress(track_tqdm=True),
50
  ):
 
 
 
 
51
  if randomize_seed:
52
  seed = random.randint(0, MAX_SEED)
53
  images, preprocessed_image = run_pipeline(
 
60
  num_inference_steps=num_inference_steps,
61
  guidance_scale=guidance_scale,
62
  seed=seed,
63
+ remove_bg_fn=None,
64
  reference_conditioning_scale=reference_conditioning_scale,
65
  negative_prompt=negative_prompt,
66
  device=device,
67
  )
68
+ return images
69
 
70
 
71
  # examples = [