Spaces:
Running
on
Zero
Running
on
Zero
Avijit Ghosh
commited on
Commit
•
4e76f82
1
Parent(s):
5bf723e
add more time to zerogpu
Browse files
app.py
CHANGED
@@ -12,11 +12,11 @@ import stone
|
|
12 |
import os
|
13 |
import spaces
|
14 |
|
15 |
-
|
16 |
|
17 |
|
18 |
from huggingface_hub import login
|
19 |
-
login(token =
|
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)
|