Update app.py

app.py

@@ -1,61 +1,72 @@
-import os, torch
+import os, torch, shutil
 from diffusers import StableDiffusionXLPipeline
 from huggingface_hub import HfApi, hf_hub_download, create_repo
 
 # ====================== CONFIG ======================
 BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
 LORA_REPO  = "CoolKrishh/Comic-SDXL-LoRA"
-LORA_FILE  = "Comic-SDXL.safetensors"
+LORA_FILE  = "comic-sdxl.safetensors"
 DEST_REPO  = "CoolKrishh/mythic-sdxl"
 OUTPUT_CKPT = "mythic-sdxl.safetensors"
+OUTPUT_DIR  = "/app/out"
 # ===================================================
 
-print("🔑 Authenticating...")
 TOKEN = os.getenv("HF_TOKEN")
 if not TOKEN:
     raise ValueError("❌ HF_TOKEN not set!")
 
 api = HfApi(token=TOKEN)
-
-print("📦 Creating repo if not exists...")
 create_repo(repo_id=DEST_REPO, exist_ok=True, token=TOKEN)
 
-print("⬇️ Downloading LoRA from Hugging Face...")
+print("⬇ Downloading LoRA...")
 lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=TOKEN)
 
-print("🚀 Loading base SDXL...")
+print("🚀 Loading SDXL base...")
 pipe = StableDiffusionXLPipeline.from_pretrained(
     BASE_MODEL, torch_dtype=torch.float16
 )
 
-print("🧬 Merging LoRA into model...")
+print("🧬 Merging LoRA...")
 pipe.load_lora_weights(lora_path, adapter_name="default")
-pipe.fuse_lora(lora_scale=1.0)
+pipe.fuse_lora()
+
+# Ensure output folder exists
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+print("💾 Saving merged pipeline...")
+pipe.save_pretrained(OUTPUT_DIR, safe_serialization=True)
+
+# ---- FIND REAL SAVED SAFETENSORS ----
+print("🔍 Scanning saved files...")
+found = None
+for f in os.listdir(OUTPUT_DIR):
+    if f.endswith(".safetensors"):
+        found = f
+        break
 
-print("💾 Saving merged model...")
-os.makedirs("./out", exist_ok=True)
-pipe.save_pretrained("./out", safe_serialization=True)
+if not found:
+    raise FileNotFoundError("❌ No .safetensors file was created by diffusers!")
 
-# Rename to final checkpoint name
-for file in os.listdir("./out"):
-    if file.endswith(".safetensors"):
-        os.rename(f"./out/{file}", f"./out/{OUTPUT_CKPT}")
+src = os.path.join(OUTPUT_DIR, found)
+dst = os.path.join(OUTPUT_DIR, OUTPUT_CKPT)
 
-file_path = os.path.abspath(f"./out/{OUTPUT_CKPT}")
+print(f"📌 Renaming:\n  {src}\n  → {dst}")
+shutil.move(src, dst)
 
-# Verify file exists
-if not os.path.exists(file_path):
-    raise FileNotFoundError(f"❌ File NOT found! Path: {file_path}")
+# ---- VERIFY FINAL FILE EXISTS ----
+if not os.path.exists(dst):
+    raise FileNotFoundError(f"❌ Still missing after rename: {dst}")
 
-print(f"✅ Model saved at: {file_path}")
+print(f"✅ Final model file ready at:\n{dst}")
 
-print("☁️ Uploading to Hugging Face...")
+# ---- UPLOAD ----
+print("☁ Uploading to Hugging Face...")
 api.upload_file(
-    path_or_fileobj=file_path,
+    path_or_fileobj=dst,
     path_in_repo=OUTPUT_CKPT,
     repo_id=DEST_REPO,
     token=TOKEN
 )
 
-print("\n🎉 DONE — Your model is now Inference API ready!")
-print(f"👉 https://huggingface.co/{DEST_REPO}")
+print("\n🎉 SUCCESS! Your model is now uploaded & inference-ready.")
+print(f"🔗 https://huggingface.co/{DEST_REPO}")