Update app.py

app.py

@@ -1,27 +1,30 @@
-import os
+import os 
 from huggingface_hub import login
 import torch
 import gradio as gr
 from diffusers import StableDiffusion3Pipeline
 
-
 # Get Hugging Face token from environment variables
 hf_token = os.getenv("HF_API_TOKEN")
 
 if hf_token:
     login(token=hf_token)
-    print("lgin sucessfull")
+    print("Login successful")
 else:
     raise ValueError("Hugging Face token is missing. Please set it in the environment variables.")
-    def image_generator(prompt):
+
+def image_generator(prompt):
     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,
-                                                         text_encoder_3=None,
-                                                         tokenizer_3 = None)
-    #pipeline.enable_model_cpu_offload()
+    pipeline = StableDiffusion3Pipeline.from_pretrained(
+        "stabilityai/stable-diffusion-3-medium-diffusers", 
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        text_encoder_3=None,
+        tokenizer_3=None
+    )
+    # Move the pipeline to the appropriate device
     pipeline.to(device)
 
+    # Generate the image
     image = pipeline(
         prompt=prompt,
         negative_prompt="blurred, ugly, watermark, low, resolution, blurry",
@@ -33,11 +36,15 @@ else:
 
     return image
 
+# Create a Gradio interface
 interface = gr.Interface(
     fn=image_generator, 
-    inputs=gr.Textbox(lines=2, placeholder = "Enter your prompt..."), 
-    outputs=gr.Image(type = "pil"),
-    title = "Image Generator App",
-    description = "This is a simple image generator app using HuggingFace's Stable Diffusion 3 model.")
+    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt..."), 
+    outputs=gr.Image(type="pil"),
+    title="Image Generator App",
+    description="This is a simple image generator app using HuggingFace's Stable Diffusion 3 model."
+)
+
+# Launch the interface
 interface.launch()
 print(interface)