nenene commited on
Commit
a34f623
·
1 Parent(s): 554790e

change to cpu

Browse files
Files changed (1) hide show
  1. app.py +6 -9
app.py CHANGED
@@ -2,14 +2,12 @@ import gradio as gr
2
  from diffusers import DiffusionPipeline
3
  import torch
4
  from PIL import Image, ImageOps
5
- import requests
6
- from io import BytesIO
7
  from transparent_background import Remover
8
 
9
  # Initialize the Diffusion Pipeline
10
  model_id = "yahoo-inc/photo-background-generation"
11
  pipeline = DiffusionPipeline.from_pretrained(model_id, custom_pipeline=model_id)
12
- pipeline = pipeline.to('cpu')
13
 
14
  def resize_with_padding(img, expected_size):
15
  img.thumbnail((expected_size[0], expected_size[1]))
@@ -21,7 +19,6 @@ def resize_with_padding(img, expected_size):
21
  return ImageOps.expand(img, padding)
22
 
23
  def process_image(input_image, prompt):
24
- # Resize and process the input image
25
  img = resize_with_padding(input_image, (512, 512))
26
 
27
  # Load background detection model
@@ -33,10 +30,10 @@ def process_image(input_image, prompt):
33
  seed = 13
34
  mask = ImageOps.invert(fg_mask)
35
  img = resize_with_padding(img, (512, 512))
36
- generator = torch.Generator(device='cuda').manual_seed(seed)
37
  cond_scale = 1.0
38
 
39
- with torch.autocast("cuda"):
40
  controlnet_image = pipeline(
41
  prompt=prompt,
42
  image=img,
@@ -55,10 +52,10 @@ def process_image(input_image, prompt):
55
  iface = gr.Interface(
56
  fn=process_image,
57
  inputs=[
58
- gr.inputs.Image(type="pil", label="Upload Image"),
59
- gr.inputs.Textbox(label="Enter Prompt")
60
  ],
61
- outputs=gr.outputs.Image(label="Generated Image"),
62
  title="Image Processing with Diffusion Pipeline",
63
  description="Upload an image and enter a prompt to generate a new image using the diffusion model."
64
  )
 
2
  from diffusers import DiffusionPipeline
3
  import torch
4
  from PIL import Image, ImageOps
 
 
5
  from transparent_background import Remover
6
 
7
  # Initialize the Diffusion Pipeline
8
  model_id = "yahoo-inc/photo-background-generation"
9
  pipeline = DiffusionPipeline.from_pretrained(model_id, custom_pipeline=model_id)
10
+ pipeline = pipeline.to('cpu') # Use CPU instead of CUDA
11
 
12
  def resize_with_padding(img, expected_size):
13
  img.thumbnail((expected_size[0], expected_size[1]))
 
19
  return ImageOps.expand(img, padding)
20
 
21
  def process_image(input_image, prompt):
 
22
  img = resize_with_padding(input_image, (512, 512))
23
 
24
  # Load background detection model
 
30
  seed = 13
31
  mask = ImageOps.invert(fg_mask)
32
  img = resize_with_padding(img, (512, 512))
33
+ generator = torch.Generator(device='cpu').manual_seed(seed) # Use CPU generator
34
  cond_scale = 1.0
35
 
36
+ with torch.no_grad(): # Disable gradient calculations for inference
37
  controlnet_image = pipeline(
38
  prompt=prompt,
39
  image=img,
 
52
  iface = gr.Interface(
53
  fn=process_image,
54
  inputs=[
55
+ gr.Image(type="pil", label="Upload Image"),
56
+ gr.Textbox(label="Enter Prompt")
57
  ],
58
+ outputs=gr.Image(label="Generated Image"),
59
  title="Image Processing with Diffusion Pipeline",
60
  description="Upload an image and enter a prompt to generate a new image using the diffusion model."
61
  )