Lookingsoft commited on
Commit
d2191a8
Β·
verified Β·
1 Parent(s): 73dfd2c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -9
app.py CHANGED
@@ -2,24 +2,20 @@ from diffusers import StableDiffusionPipeline
2
  import gradio as gr
3
  import torch
4
  import os
5
-
6
  from huggingface_hub import login
7
-
8
- # 3. Login to Hugging Face (replace YOUR_HF_TOKEN with your actual token)
9
-
10
 
11
  login(token=os.getenv("HF_TOKEN"))
12
 
13
- # 4. Load private model using your token
14
- model_id ="stabilityai/stable-diffusion-2-1-base"
15
-
16
  device = "cuda" if torch.cuda.is_available() else "cpu"
17
 
18
  pipe = StableDiffusionPipeline.from_pretrained(
19
  model_id,
20
- torch_dtype=torch.float16 # smaller memory
21
  ).to(device)
22
 
 
23
  def generate_image(prompt):
24
  image = pipe(prompt).images[0]
25
  return image
@@ -31,4 +27,4 @@ iface = gr.Interface(
31
  title="Private Stable Diffusion Demo"
32
  )
33
 
34
- iface.launch(share=True)
 
2
  import gradio as gr
3
  import torch
4
  import os
 
5
  from huggingface_hub import login
6
+ from spaces import GPU # πŸ‘ˆ only for HF Spaces
 
 
7
 
8
  login(token=os.getenv("HF_TOKEN"))
9
 
10
+ model_id = "stabilityai/stable-diffusion-2-1-base"
 
 
11
  device = "cuda" if torch.cuda.is_available() else "cpu"
12
 
13
  pipe = StableDiffusionPipeline.from_pretrained(
14
  model_id,
15
+ torch_dtype=torch.float16
16
  ).to(device)
17
 
18
+ @GPU # πŸ‘ˆ this tells HF Space to start with GPU
19
  def generate_image(prompt):
20
  image = pipe(prompt).images[0]
21
  return image
 
27
  title="Private Stable Diffusion Demo"
28
  )
29
 
30
+ iface.launch()