p4vv37 commited on
Commit
59dee13
·
1 Parent(s): a8a4e0f

updated requirements

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -7,10 +7,15 @@ diffusers.Zero123Pipeline = Zero123Pipeline
7
 
8
  def generate_view(source_img, elevation, azimuth, camera_distance, num_inference_steps):
9
 
 
 
 
 
 
10
  # Prepare pipeline
11
  pipeline = DiffusionPipeline.from_pretrained("ashawkey/stable-zero123-diffusers",
12
  torch_dtype=torch.float16, trust_remote_code=True)
13
- pipeline.to('cuda:0')
14
 
15
  # Prepare input data:
16
  image = source_img.resize((256, 256)).convert("RGB")
@@ -18,9 +23,9 @@ def generate_view(source_img, elevation, azimuth, camera_distance, num_inference
18
 
19
  # Generate and save images:
20
  images = pipeline([image],
21
- torch.tensor([elevation], dtype=torch.float16).to('cuda:0'),
22
- torch.tensor([azimuth], dtype=torch.float16).to('cuda:0'),
23
- torch.tensor([camera_distance], dtype=torch.float16).to('cuda:0'),
24
  num_inference_steps=int(num_inference_steps)).images
25
 
26
  return images[0]
 
7
 
8
  def generate_view(source_img, elevation, azimuth, camera_distance, num_inference_steps):
9
 
10
+ if torch.cuda.is_available():
11
+ device = 'cuda:0'
12
+ else:
13
+ device = 'cpu'
14
+
15
  # Prepare pipeline
16
  pipeline = DiffusionPipeline.from_pretrained("ashawkey/stable-zero123-diffusers",
17
  torch_dtype=torch.float16, trust_remote_code=True)
18
+ pipeline.to(device)
19
 
20
  # Prepare input data:
21
  image = source_img.resize((256, 256)).convert("RGB")
 
23
 
24
  # Generate and save images:
25
  images = pipeline([image],
26
+ torch.tensor([elevation], dtype=torch.float16).to(device),
27
+ torch.tensor([azimuth], dtype=torch.float16).to(device),
28
+ torch.tensor([camera_distance], dtype=torch.float16).to(device),
29
  num_inference_steps=int(num_inference_steps)).images
30
 
31
  return images[0]