Avijit Ghosh commited on
Commit
4e76f82
1 Parent(s): 5bf723e

add more time to zerogpu

Browse files
Files changed (1) hide show
  1. app.py +3 -6
app.py CHANGED
@@ -12,11 +12,11 @@ import stone
12
  import os
13
  import spaces
14
 
15
- api_key = os.getenv("AccessTokenSD3")
16
 
17
 
18
  from huggingface_hub import login
19
- login(token = api_key)
20
 
21
 
22
  # Define model initialization functions
@@ -52,12 +52,9 @@ def load_model(model_name):
52
  variant="fp16"
53
  ).to("cuda")
54
  elif model_name == "stabilityai/stable-diffusion-3-medium-diffusers":
55
- if api_key is None:
56
- raise ValueError("Hugging Face token is required to access this model")
57
  pipeline = StableDiffusion3Pipeline.from_pretrained(
58
  model_name,
59
  torch_dtype=torch.float16,
60
- use_auth_token=api_key
61
  ).to("cuda")
62
  else:
63
  raise ValueError("Unknown model name")
@@ -123,7 +120,7 @@ def skintoneplot(hex_codes):
123
  ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=sorted_hex_codes[i]))
124
  return fig
125
 
126
- @spaces.GPU
127
  def generate_images_plots(prompt, model_name):
128
  global pipeline_text2image
129
  pipeline_text2image = load_model(model_name)
 
12
  import os
13
  import spaces
14
 
15
+ access_token = os.getenv("AccessTokenSD3")
16
 
17
 
18
  from huggingface_hub import login
19
+ login(token = access_token)
20
 
21
 
22
  # Define model initialization functions
 
52
  variant="fp16"
53
  ).to("cuda")
54
  elif model_name == "stabilityai/stable-diffusion-3-medium-diffusers":
 
 
55
  pipeline = StableDiffusion3Pipeline.from_pretrained(
56
  model_name,
57
  torch_dtype=torch.float16,
 
58
  ).to("cuda")
59
  else:
60
  raise ValueError("Unknown model name")
 
120
  ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=sorted_hex_codes[i]))
121
  return fig
122
 
123
+ @spaces.GPU(duration=400)
124
  def generate_images_plots(prompt, model_name):
125
  global pipeline_text2image
126
  pipeline_text2image = load_model(model_name)