Muhammad Anas Akhtar
commited on
Update app.py
Browse files
app.py
CHANGED
@@ -1,26 +1,31 @@
|
|
1 |
import gradio as gr
|
2 |
import torch
|
3 |
-
from diffusers import StableDiffusionPipeline
|
4 |
-
|
|
|
|
|
|
|
5 |
|
6 |
-
# Retrieve the Hugging Face token stored in Hugging Face Spaces secrets
|
7 |
-
HUGGINGFACE_TOKEN = os.getenv("keyss")
|
8 |
|
9 |
-
if not HUGGINGFACE_TOKEN:
|
10 |
-
raise ValueError("Hugging Face token not found! Make sure it's added in Hugging Face Secrets.")
|
11 |
|
12 |
def image_generation(prompt):
|
13 |
# Check if GPU is available
|
14 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
15 |
|
16 |
# Load the Stable Diffusion 3 pipeline
|
17 |
-
pipeline =
|
18 |
-
"stabilityai/stable-diffusion-3-medium-diffusers",
|
19 |
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
|
20 |
use_auth_token=HUGGINGFACE_TOKEN, # Use the Hugging Face token for authentication
|
21 |
text_encoder_3=None,
|
22 |
-
tokenizer_3=None
|
23 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
24 |
|
25 |
# Enable efficient model execution
|
26 |
pipeline.enable_model_cpu_offload()
|
|
|
1 |
import gradio as gr
|
2 |
import torch
|
3 |
+
#from diffusers import StableDiffusionPipeline
|
4 |
+
|
5 |
+
|
6 |
+
from diffusers import DiffusionPipeline
|
7 |
+
|
8 |
|
|
|
|
|
9 |
|
|
|
|
|
10 |
|
11 |
def image_generation(prompt):
|
12 |
# Check if GPU is available
|
13 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
14 |
|
15 |
# Load the Stable Diffusion 3 pipeline
|
16 |
+
pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1",
|
|
|
17 |
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
|
18 |
use_auth_token=HUGGINGFACE_TOKEN, # Use the Hugging Face token for authentication
|
19 |
text_encoder_3=None,
|
20 |
+
tokenizer_3=None)
|
21 |
+
|
22 |
+
#pipeline = StableDiffusionPipeline.from_pretrained(
|
23 |
+
# "stabilityai/stable-diffusion-3-medium-diffusers",
|
24 |
+
# torch_dtype=torch.float16 if device == "cuda" else torch.float32,
|
25 |
+
# use_auth_token=HUGGINGFACE_TOKEN, # Use the Hugging Face token for authentication
|
26 |
+
# text_encoder_3=None,
|
27 |
+
# tokenizer_3=None
|
28 |
+
# )
|
29 |
|
30 |
# Enable efficient model execution
|
31 |
pipeline.enable_model_cpu_offload()
|