Update app.py

app.py

@@ -1,31 +1,68 @@
-import os
-from huggingface_hub import HfApi, HfFolder, upload_folder
 from diffusers import StableDiffusionXLPipeline
-import torch
+from huggingface_hub import hf_hub_download, HfApi
+import torch, os
 
-base_model = "stabilityai/stable-diffusion-xl-base-1.0"
-lora_model_id = "CoolKrishh/Comic-SDXL-LoRA"
-output_dir = "comic-sdxl-full"
+# ------------------------------
+# SETTINGS
+# ------------------------------
+LORA_REPO = "CoolKrishh/Comic-SDXL-LoRA"  # LoRA repo on HuggingFace
+LORA_FILENAME = "Comic-SDXL.safetensors"  # exact filename in repo
+OUTPUT_REPO = "CoolKrishh/mythic-sdxl"    # final full model repo
+SAVE_DIR = "mythic-sdxl-pipeline"
+# ------------------------------
 
-print("⬇️ Loading base model...")
-pipe = StableDiffusionXLPipeline.from_pretrained(base_model, torch_dtype=torch.float16)
-print("🔗 Loading LoRA weights...")
-pipe.load_lora_weights(lora_model_id)
+print("🔑 Fetching HF token from secrets...")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
 
-print("💾 Saving merged pipeline...")
-pipe.save_pretrained(output_dir)
+if HF_TOKEN is None:
+    raise ValueError("❌ HF_TOKEN not found! Add in HF Space → Secrets → New variable → HF_TOKEN")
 
-# ✅ Use token from environment
-token = os.getenv("HF_TOKEN")
+# ------------------------------
+# DOWNLOAD BASE MODEL
+# ------------------------------
+print("🚀 Downloading SDXL base pipeline...")
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16
+)
+pipe.to("cpu")
 
-if not token:
-    raise ValueError("❌ No HF_TOKEN found. Add it in your Space secrets.")
+# ------------------------------
+# DOWNLOAD LORA WEIGHTS
+# ------------------------------
+print("⬇️ Downloading LoRA weights from:", LORA_REPO)
+lora_path = hf_hub_download(
+    repo_id=LORA_REPO,
+    filename=LORA_FILENAME
+)
 
-api = HfApi()
-repo_id = "CoolKrishh/Comic-SDXL-Full"
+# ------------------------------
+# MERGE LoRA INTO PIPELINE
+# ------------------------------
+print("🔄 Applying LoRA weights...")
+pipe.load_lora_weights(lora_path)
+pipe.fuse_lora()  # <---- applies LoRA permanently into UNet/TextEncoder
 
-print("☁️ Uploading to Hugging Face Hub...")
-api.create_repo(repo_id=repo_id, repo_type="model", private=False, exist_ok=True, token=token)
-upload_folder(folder_path=output_dir, repo_id=repo_id, repo_type="model", token=token)
+print("✅ LoRA successfully merged into the base model!")
 
-print(f"🚀 Uploaded successfully to https://huggingface.co/{repo_id}")
+# ------------------------------
+# SAVE FULL PIPELINE
+# ------------------------------
+print("💾 Saving full pipeline to:", SAVE_DIR)
+pipe.save_pretrained(SAVE_DIR)
+
+# ------------------------------
+# UPLOAD MODEL TO HF
+# ------------------------------
+from huggingface_hub import upload_folder
+
+print(f"⬆️ Uploading model to: {OUTPUT_REPO}")
+
+upload_folder(
+    repo_id=OUTPUT_REPO,
+    folder_path=SAVE_DIR,
+    repo_type="model",
+    token=HF_TOKEN
+)
+
+print("✅ DONE — Full model uploaded to Hugging Face successfully!")