Upload folder using huggingface_hub

app.py

@@ -1,5 +1,6 @@
 """
-ConceptAligner Hugging Face Demo - Optimized for storage
+ConceptAligner Hugging Face Demo - Minimal downloads
+Only downloads VAE, uses your fine-tuned weights for everything else
 """
 
 import torch
@@ -23,11 +24,6 @@ if HF_TOKEN:
 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."""]
 ]
@@ -47,11 +43,11 @@ def download_checkpoint():
                 repo_id=MODEL_REPO,
                 filename=filename,
                 local_dir=CHECKPOINT_DIR,
-                local_dir_use_symlinks=False,
                 token=HF_TOKEN
             )
+            print(f"    ✓ {filename} downloaded")
 
-    print("✓ Checkpoint files ready!")
+    print("✓ All checkpoint files ready!")
 
 class ConceptAlignerModel:
     def __init__(self):
@@ -71,14 +67,14 @@ class ConceptAlignerModel:
         print(f"Loading models on {self.device}...")
 
         # Load ConceptAligner
-        print("  Loading ConceptAligner adapter...")
+        print("  Loading ConceptAligner...")
         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...")
+        # Load T5 encoder (your fine-tuned version with full weights)
+        print("  Loading fine-tuned T5 encoder...")
         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:
@@ -86,30 +82,33 @@ class ConceptAlignerModel:
         self.text_encoder.load_state_dict(adapter_state, strict=True)
         print("    ✓ T5 encoder loaded")
 
-        # Load VAE (will use shared cache)
+        # Only download VAE (small ~330MB) - not fine-tuned
         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,
-            cache_dir="/data/.huggingface/hub",
-            low_cpu_mem_usage=True
+            token=HF_TOKEN
         ).to(self.device)
-        print("    ✓ VAE loaded")
+        print("    ✓ VAE loaded (~330MB download)")
+
+        # Create transformer architecture WITHOUT downloading base weights
+        print("  Initializing transformer architecture...")
 
-        # Load transformer (will use shared cache)
-        print("  Loading transformer from FLUX.1-dev...")
-        transformer = FluxTransformer2DModel.from_pretrained(
+        # Get config only (no weights download)
+        from diffusers.models.transformers.transformer_flux import FluxTransformerConfig
+        config = FluxTransformer2DModel.load_config(
             'black-forest-labs/FLUX.1-dev',
             subfolder="transformer",
-            torch_dtype=self.dtype,
-            token=HF_TOKEN,
-            cache_dir="/data/.huggingface/hub",
-            low_cpu_mem_usage=True
+            token=HF_TOKEN
         )
 
-        print("  Adding LoRA adapters to transformer...")
+        # Initialize empty transformer from config
+        transformer = FluxTransformer2DModel.from_config(config)
+        print("    ✓ Transformer architecture initialized")
+
+        # Add LoRA config (needed for architecture)
+        print("  Adding LoRA adapter config...")
         transformer_lora_config = LoraConfig(
             r=256, lora_alpha=256, lora_dropout=0.0, init_lora_weights="gaussian",
             target_modules=[
@@ -122,27 +121,32 @@ class ConceptAlignerModel:
         transformer.add_adapter(transformer_lora_config)
         transformer.context_embedder.requires_grad_(True)
 
-        print("  Loading fine-tuned transformer weights...")
+        # Load YOUR FULL fine-tuned transformer weights (no base model needed!)
+        print("  Loading YOUR 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")
+        transformer = transformer.to(self.device).to(self.dtype)
+        print("    ✓ Fine-tuned transformer loaded (~26GB from your checkpoint)")
 
         # Load empty pooled clip
         self.empty_pooled_clip = torch.load(
             os.path.join(self.checkpoint_path, "empty_pooled_clip.pt"),
-            map_location=self.device
+            map_location=self.device,
+            weights_only=True
         ).to(self.dtype)
+        print("    ✓ Empty pooled clip loaded")
 
-        # Create pipeline
-        print("  Creating pipeline...")
+        # Create scheduler (just config, no weights)
+        print("  Loading scheduler config...")
         noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             'black-forest-labs/FLUX.1-dev',
             subfolder="scheduler",
-            token=HF_TOKEN,
-            cache_dir="/data/.huggingface/hub"
+            token=HF_TOKEN
         )
+        print("    ✓ Scheduler loaded")
 
+        # Create pipeline
+        print("  Creating pipeline...")
         self.pipe = CustomFluxKontextPipeline(
             scheduler=noise_scheduler,
             aligner=self.model,
@@ -151,7 +155,11 @@ class ConceptAlignerModel:
             text_embedder=self.text_encoder,
         ).to(self.device)
 
-        print("✓ All models loaded successfully!")
+        print("="*60)
+        print("✓ ALL MODELS LOADED SUCCESSFULLY!")
+        print("="*60)
+        print(f"Total downloads: ~330MB VAE + ~26GB your checkpoints")
+        print(f"Saved: ~24GB by not downloading base FLUX transformer!")
 
         # Print memory usage
         if torch.cuda.is_available():
@@ -193,13 +201,18 @@ class ConceptAlignerModel:
 
 # Initialize model
 print("="*60)
-print("Initializing ConceptAligner...")
+print("🚀 Initializing ConceptAligner Demo")
 print("="*60)
 model = ConceptAlignerModel()
 
 # Create Gradio interface
 with gr.Blocks(title="ConceptAligner", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎨 ConceptAligner Demo\nGenerate images with fine-tuned concept alignment!")
+    gr.Markdown("""
+    # 🎨 ConceptAligner Demo
+    
+    Generate images with fine-tuned concept alignment using FLUX!
+    This demo uses fully fine-tuned weights - no base model downloads needed.
+    """)
 
     with gr.Row():
         with gr.Column(scale=1):