Update main.py

main.py

@@ -1,11 +1,10 @@
 import torch
 from diffusers import StableDiffusionXLPipeline
-from huggingface_hub import login
 
 
 class StableDiffusionEngine:
     """
-    Core Stable Diffusion engine for research and experimentation.
+    Stable Diffusion XL Engine (HF Spaces safe version)
     """
 
     def __init__(
@@ -13,17 +12,6 @@ class StableDiffusionEngine:
         model_id: str = "stabilityai/sdxl-base-1.0",
         hf_token: str | None = None,
     ):
-        """
-        Initialize the Stable Diffusion pipeline.
-
-        Args:
-            model_id (str): Hugging Face model ID
-            hf_token (str): Hugging Face access token
-        """
-
-        if hf_token:
-            login(token=hf_token)
-
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
         self.pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -31,6 +19,7 @@ class StableDiffusionEngine:
             torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
             use_safetensors=True,
             safety_checker=True,
+            token=hf_token,   # ✅ CRITICAL FIX
         )
 
         self.pipe.to(self.device)
@@ -43,20 +32,6 @@ class StableDiffusionEngine:
         width: int = 1024,
         height: int = 1024,
     ):
-        """
-        Generate an image from a text prompt.
-
-        Args:
-            prompt (str): Text prompt
-            num_inference_steps (int): Diffusion steps
-            guidance_scale (float): CFG scale
-            width (int): Image width
-            height (int): Image height
-
-        Returns:
-            PIL.Image.Image
-        """
-
         with torch.no_grad():
             image = self.pipe(
                 prompt=prompt,