gokaygokay commited on
Commit
9148f31
1 Parent(s): 1ed4ae0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -2
app.py CHANGED
@@ -6,9 +6,13 @@ import numpy as np
6
 
7
  from aura_sr import AuraSR
8
  import spaces
 
9
 
10
- # Load the AuraSR model
11
- aura_sr = AuraSR.from_pretrained("fal-ai/AuraSR").to("cuda")
 
 
 
12
 
13
  @spaces.GPU
14
  def process_image(input_image):
 
6
 
7
  from aura_sr import AuraSR
8
  import spaces
9
+ import torch
10
 
11
+
12
+ aura_sr = AuraSR.from_pretrained("fal-ai/AuraSR", device_map="cpu")
13
+
14
+ if torch.cuda.is_available():
15
+ aura_sr.to("cuda")
16
 
17
  @spaces.GPU
18
  def process_image(input_image):