Commit
813561e
1 Parent(s): e049b9f

Only use GPU when needed (#2)

Browse files

- Only use GPU when needed (060c3c709b24c0f8d6512f3a4a56522e88ba158c)


Co-authored-by: Fabrice TIERCELIN <[email protected]>

Files changed (1) hide show
  1. app.py +6 -4
app.py CHANGED
@@ -19,22 +19,24 @@ aura_sr = AuraSR.from_pretrained("fal-ai/AuraSR")
19
  # Restore original torch.load
20
  torch.load = original_load
21
 
22
- @spaces.GPU
23
  def process_image(input_image):
24
  if input_image is None:
25
- return None
26
 
27
  # Convert to PIL Image for resizing
28
  pil_image = Image.fromarray(input_image)
29
 
30
  # Upscale the image using AuraSR
31
- with torch.no_grad():
32
- upscaled_image = aura_sr.upscale_4x(pil_image)
33
 
34
  # Convert result to numpy array if it's not already
35
  result_array = np.array(upscaled_image)
36
 
37
  return [input_image, result_array]
 
 
 
 
38
 
39
  title = """<h1 align="center">AuraSR - An open reproduction of the GigaGAN Upscaler from fal.ai</h1>
40
  <p><center>
 
19
  # Restore original torch.load
20
  torch.load = original_load
21
 
 
22
  def process_image(input_image):
23
  if input_image is None:
24
+ raise gr.Error("Please provide an image to upscale.")
25
 
26
  # Convert to PIL Image for resizing
27
  pil_image = Image.fromarray(input_image)
28
 
29
  # Upscale the image using AuraSR
30
+ upscaled_image = process_image_on_gpu(pil_image)
 
31
 
32
  # Convert result to numpy array if it's not already
33
  result_array = np.array(upscaled_image)
34
 
35
  return [input_image, result_array]
36
+
37
+ @spaces.GPU
38
+ def process_image_on_gpu(pil_image):
39
+ return aura_sr.upscale_4x(pil_image)
40
 
41
  title = """<h1 align="center">AuraSR - An open reproduction of the GigaGAN Upscaler from fal.ai</h1>
42
  <p><center>