Abdullah commited on
Commit
db8e26c
·
1 Parent(s): a0ed50b

- Added cpu and Cuda usage

Browse files
Files changed (1) hide show
  1. app.py +6 -2
app.py CHANGED
@@ -3,6 +3,8 @@ import spaces
3
  from gradio_imageslider import ImageSlider
4
  from image_gen_aux import UpscaleWithModel
5
  from image_gen_aux.utils import load_image
 
 
6
 
7
  # This uses https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/upscalers/README.md
8
  # Also this space has been duplicated from their official huggingface space, https://huggingface.co/spaces/OzzyGT/basic_upscaler
@@ -118,7 +120,9 @@ MODELS = {
118
  def upscale_image(image, model_selection):
119
  original = load_image(image)
120
 
121
- upscaler = UpscaleWithModel.from_pretrained(MODELS[model_selection]).to("cuda")
 
 
122
  image = upscaler(original, tiling=True, tile_width=1024, tile_height=1024)
123
 
124
  return original, image
@@ -165,4 +169,4 @@ with gr.Blocks() as demo:
165
  )
166
 
167
 
168
- demo.launch(share=False)
 
3
  from gradio_imageslider import ImageSlider
4
  from image_gen_aux import UpscaleWithModel
5
  from image_gen_aux.utils import load_image
6
+ import torch
7
+
8
 
9
  # This uses https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/upscalers/README.md
10
  # Also this space has been duplicated from their official huggingface space, https://huggingface.co/spaces/OzzyGT/basic_upscaler
 
120
  def upscale_image(image, model_selection):
121
  original = load_image(image)
122
 
123
+ device = "cuda" if torch.cuda.is_available() else "cpu"
124
+ upscaler = UpscaleWithModel.from_pretrained(MODELS[model_selection]).to(device)
125
+
126
  image = upscaler(original, tiling=True, tile_width=1024, tile_height=1024)
127
 
128
  return original, image
 
169
  )
170
 
171
 
172
+ demo.launch(share=True)