Update app.py

app.py

@@ -1,18 +1,31 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusion3Pipeline
+from diffusers import StableDiffusionPipeline
+from huggingface_hub import hf_api
 
+# Retrieve the Hugging Face token stored in Hugging Face Spaces secrets
+HUGGINGFACE_TOKEN = hf_api.get_secret("keyss")
+
+if not HUGGINGFACE_TOKEN:
+    raise ValueError("Hugging Face token not found! Make sure it's added in the Hugging Face Secrets.")
 
 def image_generation(prompt):
+    # Check if GPU is available
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipeline = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers",
-                                                        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-                                                        use_auth_token=True,
-                                                        text_encoder_3 =None,
-                                                        tokenizer_3 =None)
+    
+    # Load the Stable Diffusion 3 pipeline
+    pipeline = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-3-medium-diffusers",
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        use_auth_token=HUGGINGFACE_TOKEN,  # Use the Hugging Face token for authentication
+        text_encoder_3=None,
+        tokenizer_3=None
+    )
+    
+    # Enable efficient model execution
     pipeline.enable_model_cpu_offload()
-    # pipeline.to(device)
-
+    
+    # Generate an image based on the prompt
     image = pipeline(
         prompt=prompt,
         negative_prompt="blurred, ugly, watermark, low resolution, blurry",
@@ -21,17 +34,17 @@ def image_generation(prompt):
         width=1024,
         guidance_scale=9.0
     ).images[0]
-
+    
     return image
 
-# image_generation("A magician cat doing spell")
-
-interface= gr.Interface(
+# Define the Gradio interface
+interface = gr.Interface(
     fn=image_generation,
-    inputs = gr.Textbox(lines=2, placeholder="Enter your Prompt..."),
-    outputs =gr.Image(type="pil"),
-    title =" Image creation using Stable Diffusion 3 Model",
-    description="This application will be used to generate awesome images using SD3 model"
+    inputs=gr.Textbox(lines=2, placeholder="Enter your Prompt..."),
+    outputs=gr.Image(type="pil"),
+    title="Image Creation using Stable Diffusion 3 Model",
+    description="This application generates awesome images using the Stable Diffusion 3 model."
 )
 
-interface.launch()
+# Launch the Gradio app
+interface.launch()