Update model.py

model.py

@@ -118,6 +118,45 @@ class ModelHandler:
             print(f"  [WARNING] Failed to enable xFormers: {e}")
         # --- END NEW ---
 
+        # ==============================================================================
+        # 3.5 LOAD PIVOTAL TUNING EMBEDDINGS (TEXTUAL INVERSION)
+        # ==============================================================================
+        print("Loading Textual Inversion Embeddings (Pivotal Tuning)...")
+        
+        # Define the embedding name (assumed to be in the same repo as Config.REPO_ID)
+        embedding_filename = "retroart.safetensors"
+        embedding_path = os.path.join("./models", embedding_filename)
+        
+        # 1. Download if missing
+        if not os.path.exists(embedding_path):
+            print(f"Downloading embedding '{embedding_filename}' to {embedding_path}...")
+            try:
+                hf_hub_download(
+                    repo_id=Config.REPO_ID,  # Or 'primerz/pixagram' if separate
+                    filename=embedding_filename,
+                    local_dir="./models",
+                    local_dir_use_symlinks=False
+                )
+            except Exception as e:
+                print(f"  [WARNING] Could not download embedding: {e}")
+
+        # 2. Load into the pipeline
+        # SDXL pipelines automatically handle loading into both tokenizers/text_encoders
+        if os.path.exists(embedding_path):
+            try:
+                self.pipeline.load_textual_inversion(
+                    embedding_path, 
+                    token="retroart", # Trigger word: <retroart> or retroart
+                    local_files_only=True
+                )
+                print(f"  [OK] Loaded embedding '{embedding_filename}' associated with token 'retroart'.")
+            except Exception as e:
+                print(f"  [ERROR] Failed to load textual inversion: {e}")
+        else:
+            print(f"  [SKIP] Embedding file not found locally.")
+        # ==============================================================================
+
+
         # 4. Set Scheduler
         # --- MODIFIED: Disable clipping to prevent NaN artifacts ---
         print("Configuring LCMScheduler...")