Update app.py

app.py

@@ -5,30 +5,41 @@ import spaces
 import random
 from diffusers import StableDiffusion3Pipeline
 from diffusers.loaders import SD3LoraLoaderMixin
+from safetensors.torch import load_file, save_file
 
 # Device selection
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Load the Hugging Face token securely
+# Load Hugging Face token securely
 token = os.getenv("HF_TOKEN")
 
 # Model ID for SD 3.5 Large
 model_repo_id = "stabilityai/stable-diffusion-3.5-large"
 
-# Load Stable Diffusion pipeline once at the start
+# Convert .pt to .safetensors if needed
+lora_pt_path = "lora_trained_model.pt"
+lora_safetensors_path = "lora_trained_model.safetensors"
+
+if os.path.exists(lora_pt_path) and not os.path.exists(lora_safetensors_path):
+    print("🔄 Converting LoRA .pt to .safetensors...")
+    lora_weights = torch.load(lora_pt_path, map_location="cpu")
+    save_file(lora_weights, lora_safetensors_path)
+    print(f"✅ LoRA saved as {lora_safetensors_path}")
+
+# Load Stable Diffusion pipeline
 pipeline = StableDiffusion3Pipeline.from_pretrained(
     model_repo_id,
     torch_dtype=torch_dtype,
     use_safetensors=True,  # Use safetensors format if supported
 ).to(device)
 
-# Load the LoRA trained weights once at the start
-lora_path = "lora_trained_model.safetensors"  # Use the correct file name
-if os.path.exists(lora_path):
+# Load and fuse LoRA trained weights
+if os.path.exists(lora_safetensors_path):
     try:
-        SD3LoraLoaderMixin.load_lora_into_model(pipeline, lora_path)  # Correct method
-        print("✅ LoRA weights loaded successfully!")
+        pipeline.load_lora_weights(".", weight_name="lora_trained_model.safetensors")  # Corrected loading method
+        pipeline.fuse_lora()  # Merges LoRA into the base model
+        print("✅ LoRA weights loaded and fused successfully!")
     except Exception as e:
         print(f"❌ Error loading LoRA: {e}")
 else:
@@ -37,7 +48,7 @@ else:
 # Verify if LoRA is applied
 for name, param in pipeline.text_encoder.named_parameters():
     if "lora" in name.lower():
-        print(f"LoRA applied to: {name}, requires_grad={param.requires_grad}")
+        print(f"✅ LoRA applied to: {name}, requires_grad={param.requires_grad}")
 
 # Ensure GPU allocation in Hugging Face Spaces
 @spaces.GPU(duration=65)