Update app.py

app.py

@@ -25,20 +25,22 @@ async def lifespan(app: FastAPI):
     # Load the model on startup
     hf_token = os.getenv("HF_TOKEN")
     if not hf_token:
-        raise RuntimeError("HF_TOKEN environment variable not set!")
+        raise RuntimeError("HF_TOKEN environment variable not set! Please add it in the Space settings.")
 
-    model_id = "rxmha125/sdxl-base-1.0-private" # Your private model ID
+    # --- *** THIS IS THE ONLY LINE THAT CHANGES *** ---
+    # We now point directly to the original, public model repository.
+    model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+    # --- ****************************************** ---
+    
     print(f"Loading model: {model_id}")
     
-    # --- *** THIS IS THE CORRECTED PART *** ---
-    # We removed variant="fp16" and use_safetensors=True
-    # to load the available .bin files instead of the missing .safetensors.
     pipe = AutoPipelineForText2Image.from_pretrained(
         model_id,
-        torch_dtype=torch.float16, # Keep for memory optimization
+        torch_dtype=torch.float16,
+        variant="fp16", # Use the optimized fp16 variant
+        use_safetensors=True,
         token=hf_token
     ).to("cuda")
-    # --- *********************************** ---
 
     # Optimization for speed and memory
     pipe.enable_model_cpu_offload()
@@ -66,14 +68,12 @@ def generate_image(request: ImageRequest):
     
     print(f"Generating image for prompt: '{request.prompt}'")
     try:
-        # Generate the image
         image = pipe(
             prompt=request.prompt,
             negative_prompt=request.negative_prompt,
             num_inference_steps=request.steps
         ).images[0]
 
-        # Convert image to Base64
         buffer = io.BytesIO()
         image.save(buffer, format="PNG")
         img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")