Upload folder using huggingface_hub

app.py

@@ -1,6 +1,5 @@
 """
-ConceptAligner Hugging Face Demo
-Downloads weights from model repo at startup
+ConceptAligner Hugging Face Demo - Optimized for storage
 """
 
 import torch
@@ -19,13 +18,16 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
     print("✓ Logged in to Hugging Face")
-else:
-    print("⚠️  Warning: No HF_TOKEN found in environment")
 
 # Configuration
 MODEL_REPO = "Shaoan/ConceptAligner-Weights"
 CHECKPOINT_DIR = "./checkpoint"
 
+# Use HF cache directory to avoid duplication
+os.environ["HF_HOME"] = "/data/.huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/data/.huggingface/hub"
+os.environ["HF_HUB_CACHE"] = "/data/.huggingface/hub"
+
 EXAMPLE_PROMPTS = [
     ["""In the image, a single white duck walks proudly across a cobblestone street. It wears a red ribbon around its neck, and the morning sun glints off puddles from a recent rain. In the background, a few people watch and smile, giving the scene a playful charm. The duck's confident stride and upright posture make it appear oddly dignified."""]
 ]
@@ -49,7 +51,7 @@ def download_checkpoint():
                 token=HF_TOKEN
             )
 
-    print("✓ All files ready!")
+    print("✓ Checkpoint files ready!")
 
 class ConceptAlignerModel:
     def __init__(self):
@@ -69,33 +71,45 @@ class ConceptAlignerModel:
         print(f"Loading models on {self.device}...")
 
         # Load ConceptAligner
+        print("  Loading ConceptAligner adapter...")
         self.model = ConceptAligner().to(self.device).to(self.dtype)
         adapter_state = load_file(os.path.join(self.checkpoint_path, "model_1.safetensors"))
         self.model.load_state_dict(adapter_state, strict=True)
+        print("    ✓ ConceptAligner loaded")
 
         # Load T5 encoder
+        print("  Loading T5 encoder adapter...")
         self.text_encoder = LoraT5Embedder(device=self.device).to(self.dtype)
         adapter_state = load_file(os.path.join(self.checkpoint_path, "model_2.safetensors"))
         if "t5_encoder.shared.weight" in adapter_state:
             adapter_state["t5_encoder.encoder.embed_tokens.weight"] = adapter_state["t5_encoder.shared.weight"]
         self.text_encoder.load_state_dict(adapter_state, strict=True)
+        print("    ✓ T5 encoder loaded")
 
-        # Load VAE with token
+        # Load VAE (will use shared cache)
+        print("  Loading VAE from FLUX.1-dev...")
         vae = AutoencoderKL.from_pretrained(
             'black-forest-labs/FLUX.1-dev',
             subfolder="vae",
             torch_dtype=self.dtype,
-            token=HF_TOKEN
+            token=HF_TOKEN,
+            cache_dir="/data/.huggingface/hub",
+            low_cpu_mem_usage=True
         ).to(self.device)
+        print("    ✓ VAE loaded")
 
-        # Load transformer with token
+        # Load transformer (will use shared cache)
+        print("  Loading transformer from FLUX.1-dev...")
         transformer = FluxTransformer2DModel.from_pretrained(
             'black-forest-labs/FLUX.1-dev',
             subfolder="transformer",
             torch_dtype=self.dtype,
-            token=HF_TOKEN
+            token=HF_TOKEN,
+            cache_dir="/data/.huggingface/hub",
+            low_cpu_mem_usage=True
         )
 
+        print("  Adding LoRA adapters to transformer...")
         transformer_lora_config = LoraConfig(
             r=256, lora_alpha=256, lora_dropout=0.0, init_lora_weights="gaussian",
             target_modules=[
@@ -108,9 +122,11 @@ class ConceptAlignerModel:
         transformer.add_adapter(transformer_lora_config)
         transformer.context_embedder.requires_grad_(True)
 
+        print("  Loading fine-tuned transformer weights...")
         transformer_state = load_file(os.path.join(self.checkpoint_path, "model.safetensors"))
         transformer.load_state_dict(transformer_state, strict=True)
         transformer = transformer.to(self.device)
+        print("    ✓ Transformer loaded")
 
         # Load empty pooled clip
         self.empty_pooled_clip = torch.load(
@@ -119,8 +135,12 @@ class ConceptAlignerModel:
         ).to(self.dtype)
 
         # Create pipeline
+        print("  Creating pipeline...")
         noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            'black-forest-labs/FLUX.1-dev', subfolder="scheduler", token=HF_TOKEN
+            'black-forest-labs/FLUX.1-dev',
+            subfolder="scheduler",
+            token=HF_TOKEN,
+            cache_dir="/data/.huggingface/hub"
         )
 
         self.pipe = CustomFluxKontextPipeline(
@@ -131,7 +151,13 @@ class ConceptAlignerModel:
             text_embedder=self.text_encoder,
         ).to(self.device)
 
-        print("✓ Model loaded!")
+        print("✓ All models loaded successfully!")
+
+        # Print memory usage
+        if torch.cuda.is_available():
+            allocated = torch.cuda.memory_allocated(0) / 1024**3
+            reserved = torch.cuda.memory_reserved(0) / 1024**3
+            print(f"📊 GPU Memory: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved")
 
     @torch.no_grad()
     def generate_image(self, prompt, threshold=0.0, topk=0, height=512, width=512,
@@ -155,7 +181,9 @@ class ConceptAlignerModel:
 
             return prev_image, current_image, prev_prompt
         except Exception as e:
-            print(f"Error: {e}")
+            import traceback
+            print(f"❌ Generation error: {e}")
+            print(traceback.format_exc())
             return self.previous_image, None, self.previous_prompt or ""
 
     def reset_history(self):
@@ -164,7 +192,9 @@ class ConceptAlignerModel:
         return None, None, "No previous generation"
 
 # Initialize model
+print("="*60)
 print("Initializing ConceptAligner...")
+print("="*60)
 model = ConceptAlignerModel()
 
 # Create Gradio interface