Update app.py

app.py

@@ -32,7 +32,7 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # --- Configuration ---
-MODEL_PATH = os.environ.get("OV_MODEL_PATH", "/app/models/base")
+MODEL_PATH = os.environ.get("OV_MODEL_PATH", "/app/models/LCM-dreamshaper-v7-openvino")
 CONTROLNET_ID = "lllyasviel/sd-controlnet-canny"
 OUTPUT_DIR = Path("/tmp/outputs")
 LORA_CACHE_DIR = Path("/app/models/loras")
@@ -162,28 +162,6 @@ AVAILABLE_LORAS = [
 ]
 
 # --- Helper Functions ---
-async def download_lora(lora_id: str) -> Path:
-    """Download a LoRA from Hugging Face if not already cached."""
-    lora_path = LORA_CACHE_DIR / f"{lora_id.replace('/', '_')}.safetensors"
-    if lora_path.exists():
-        return lora_path
-
-    logger.info(f"Downloading LoRA: {lora_id}")
-    from huggingface_hub import hf_hub_download
-    try:
-        downloaded_path = hf_hub_download(
-            repo_id=lora_id,
-            filename="pytorch_lora_weights.safetensors",
-            cache_dir=str(LORA_CACHE_DIR)
-        )
-        # Create a symlink to our expected path
-        if not lora_path.exists():
-            os.symlink(downloaded_path, lora_path)
-        return lora_path
-    except Exception as e:
-        logger.error(f"Failed to download LoRA {lora_id}: {e}")
-        raise HTTPException(status_code=400, detail=f"LoRA {lora_id} not found or invalid")
-
 def download_lora_sync(lora_id: str) -> Path:
     """Download a LoRA from Hugging Face if not already cached (synchronous)."""
     lora_path = LORA_CACHE_DIR / f"{lora_id.replace('/', '_')}.safetensors"
@@ -206,6 +184,30 @@ def download_lora_sync(lora_id: str) -> Path:
         logger.error(f"Failed to download LoRA {lora_id}: {e}")
         raise HTTPException(status_code=400, detail=f"LoRA {lora_id} not found or invalid")
 
+def apply_loras_to_pipe(pipe, lora_ids: str, lora_scales: Optional[str] = None) -> list:
+    """Apply LoRAs to a PyTorch pipeline and return list of applied LoRAs."""
+    lora_list = []
+    if not lora_ids:
+        return lora_list
+
+    lora_ids_list = [lid.strip() for lid in lora_ids.split(",")]
+    if lora_scales:
+        scales_list = [float(s.strip()) for s in lora_scales.split(",")]
+    else:
+        scales_list = [1.0] * len(lora_ids_list)
+
+    if len(lora_ids_list) != len(scales_list):
+        raise HTTPException(status_code=400, detail="Number of LoRA IDs must match number of scales")
+
+    # Download and load each LoRA synchronously
+    for lora_id, scale in zip(lora_ids_list, scales_list):
+        lora_path = download_lora_sync(lora_id)
+        pipe.load_lora_weights(str(lora_path))
+        pipe.fuse_lora(lora_scale=scale)
+        lora_list.append({"id": lora_id, "scale": scale})
+
+    return lora_list
+
 def load_ov_pipeline():
     """Load the OpenVINO-optimized LCM pipeline from local path."""
     global ov_pipeline, tokenizer
@@ -305,30 +307,6 @@ def cleanup_temp_files(*paths: Path):
         except Exception as e:
             logger.warning(f"Failed to delete {path}: {e}")
 
-def apply_loras_to_pipe(pipe, lora_ids: str, lora_scales: Optional[str] = None) -> list:
-    """Apply LoRAs to a PyTorch pipeline and return list of applied LoRAs."""
-    lora_list = []
-    if not lora_ids:
-        return lora_list
-
-    lora_ids_list = [lid.strip() for lid in lora_ids.split(",")]
-    if lora_scales:
-        scales_list = [float(s.strip()) for s in lora_scales.split(",")]
-    else:
-        scales_list = [1.0] * len(lora_ids_list)
-
-    if len(lora_ids_list) != len(scales_list):
-        raise HTTPException(status_code=400, detail="Number of LoRA IDs must match number of scales")
-
-    # Download and load each LoRA synchronously
-    for lora_id, scale in zip(lora_ids_list, scales_list):
-        lora_path = download_lora_sync(lora_id)
-        pipe.load_lora_weights(str(lora_path))
-        pipe.fuse_lora(lora_scale=scale)
-        lora_list.append({"id": lora_id, "scale": scale})
-
-    return lora_list
-
 # --- API Endpoints ---
 @app.get("/", include_in_schema=False)
 async def root():
@@ -639,10 +617,10 @@ async def get_image(filename: str):
     summary="Health check endpoint"
 )
 async def health_check():
-    base_model_exists = (Path("/app/models/base") / "openvino_model.xml").exists()
+    model_index_exists = (Path(MODEL_PATH) / "model_index.json").exists()
     return {
-        "status": "healthy" if base_model_exists else "degraded",
-        "base_model_loaded": base_model_exists,
+        "status": "healthy" if model_index_exists else "degraded",
+        "base_model_loaded": model_index_exists,
         "available_loras": len(AVAILABLE_LORAS)
     }