Update app.py

app.py

@@ -13,13 +13,16 @@ from PIL import Image
 import numpy as np
 import cv2
 from transformers import pipeline as transformers_pipeline
+from huggingface_hub import hf_hub_download
 import os
 
-# Device configuration
+# Configuration
+MODEL_REPO = "primerz/pixagram"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
 print(f"Using device: {device}")
+print(f"Loading models from: {MODEL_REPO}")
 
 class RetroArtConverter:
     def __init__(self):
@@ -42,16 +45,22 @@ class RetroArtConverter:
             torch_dtype=self.dtype
         ).to(self.device)
         
-        # Load custom VAE
-        print("Loading custom VAE (pixelate)...")
-        vae_path = "./models/vae/pixelate.safetensors"
-        if os.path.exists(vae_path):
+        # Load custom VAE from HuggingFace Hub
+        print("Loading custom VAE (pixelate) from HuggingFace Hub...")
+        try:
+            vae_path = hf_hub_download(
+                repo_id=MODEL_REPO,
+                filename="pixelate.safetensors",
+                repo_type="model"
+            )
             self.vae = AutoencoderKL.from_single_file(
                 vae_path,
                 torch_dtype=self.dtype
             ).to(self.device)
-        else:
-            print("Warning: Custom VAE not found, using default SDXL VAE")
+            print("✓ Custom VAE loaded successfully")
+        except Exception as e:
+            print(f"Warning: Could not load custom VAE: {e}")
+            print("Using default SDXL VAE")
             self.vae = AutoencoderKL.from_pretrained(
                 "madebyollin/sdxl-vae-fp16-fix",
                 torch_dtype=self.dtype
@@ -64,11 +73,14 @@ class RetroArtConverter:
             model="Intel/dpt-hybrid-midas"
         )
         
-        # Load SDXL base model with custom checkpoint
-        print("Loading SDXL model (horizon)...")
-        model_path = "./models/checkpoints/horizon.safetensors"
-        
-        if os.path.exists(model_path):
+        # Load SDXL checkpoint from HuggingFace Hub
+        print("Loading SDXL checkpoint (horizon) from HuggingFace Hub...")
+        try:
+            model_path = hf_hub_download(
+                repo_id=MODEL_REPO,
+                filename="horizon.safetensors",
+                repo_type="model"
+            )
             self.pipe = StableDiffusionXLControlNetPipeline.from_single_file(
                 model_path,
                 controlnet=self.controlnet_depth,
@@ -76,8 +88,10 @@ class RetroArtConverter:
                 torch_dtype=self.dtype,
                 use_safetensors=True
             ).to(self.device)
-        else:
-            print("Warning: Custom checkpoint not found, using default SDXL")
+            print("✓ Custom checkpoint loaded successfully")
+        except Exception as e:
+            print(f"Warning: Could not load custom checkpoint: {e}")
+            print("Using default SDXL")
             self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
                 "stabilityai/stable-diffusion-xl-base-1.0",
                 controlnet=self.controlnet_depth,
@@ -86,14 +100,19 @@ class RetroArtConverter:
                 use_safetensors=True
             ).to(self.device)
         
-        # Load custom LORA
-        print("Loading LORA (retroart)...")
-        lora_path = "./models/lora/retroart.safetensors"
-        if os.path.exists(lora_path):
+        # Load LORA from HuggingFace Hub
+        print("Loading LORA (retroart) from HuggingFace Hub...")
+        try:
+            lora_path = hf_hub_download(
+                repo_id=MODEL_REPO,
+                filename="retroart.safetensors",
+                repo_type="model"
+            )
             self.pipe.load_lora_weights(lora_path)
-            print("LORA loaded successfully")
-        else:
-            print("Warning: Custom LORA not found")
+            print("✓ LORA loaded successfully")
+        except Exception as e:
+            print(f"Warning: Could not load LORA: {e}")
+            print("Running without LORA")
         
         # Optimize pipeline
         self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
@@ -331,4 +350,4 @@ if __name__ == "__main__":
         server_port=7860,
         share=False,
         show_api=True  # Enable API
-    )
+    )