shanetx commited on
Commit
407b04d
1 Parent(s): a0dace5

update cuda

Browse files
Files changed (2) hide show
  1. app.py +6 -2
  2. requirements.txt +3 -1
app.py CHANGED
@@ -2,13 +2,17 @@ import os
2
  import gradio as gr
3
  import torch
4
 
 
 
 
 
5
  from diffusers import StableDiffusionImg2ImgPipeline
6
 
7
  os.environ['GRADIO_THEME'] = 'default'
8
  # load the pipeline
9
- device = "cpu"
10
  model_id_or_path = "runwayml/stable-diffusion-v1-5"
11
- img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float32)
12
  img2img_pipe = img2img_pipe.to(device)
13
 
14
 
 
2
  import gradio as gr
3
  import torch
4
 
5
+ print(f"Is CUDA available: {torch.cuda.is_available()}")
6
+ # True
7
+ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
8
+
9
  from diffusers import StableDiffusionImg2ImgPipeline
10
 
11
  os.environ['GRADIO_THEME'] = 'default'
12
  # load the pipeline
13
+ device = "cuda"
14
  model_id_or_path = "runwayml/stable-diffusion-v1-5"
15
+ img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
16
  img2img_pipe = img2img_pipe.to(device)
17
 
18
 
requirements.txt CHANGED
@@ -1,4 +1,6 @@
 
1
  torch
2
  diffusers
3
  transformers
4
- accelerate
 
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu113
2
  torch
3
  diffusers
4
  transformers
5
+ accelerate
6
+