Update app.py

app.py

@@ -7,6 +7,7 @@ from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
 import os
 import logging
+from safetensors.torch import load_file  # Import safetensors loader
 
 # Set up detailed logging
 logging.basicConfig(level=logging.INFO)
@@ -74,7 +75,7 @@ ip_adapter_weights = download_file(
     ip_adapter_path
 )
 
-# Load the pipeline with verbose logging
+# Load the pipeline
 logger.info("Loading Stable Diffusion XL base model...")
 try:
     pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -83,18 +84,18 @@ try:
         safety_checker=None,
         local_files_only=False,
         cache_dir=cache_dir,
-        variant="fp16",  # Use FP16 weights to reduce memory usage
-        use_safetensors=True  # Prefer safetensors format if available
+        variant="fp16",
+        use_safetensors=True
     )
     logger.info("SDXL base model loaded successfully.")
 except Exception as e:
     logger.error(f"Failed to load SDXL base model: {e}")
     raise
 
-# Load Kolors unet weights
+# Load Kolors unet weights using safetensors
 logger.info(f"Loading Kolors unet weights from {kolors_weights}...")
 try:
-    state_dict = torch.load(kolors_weights, map_location=device, weights_only=False)
+    state_dict = load_file(kolors_weights, device=device)  # Use safetensors loader
     pipe.unet.load_state_dict(state_dict)
     logger.info("Kolors unet weights loaded successfully.")
 except Exception as e:
@@ -131,9 +132,9 @@ def generate_image(uploaded_image, prompt):
         image = pipe(
             prompt=prompt,
             image_embeds=face_emb,
-            num_inference_steps=10,  # Reduced steps for faster execution
+            num_inference_steps=10,
             guidance_scale=7.5,
-            height=256,  # Smaller resolution to fit memory
+            height=256,
             width=256
         ).images[0]
         logger.info("Image generated successfully.")