Upload folder using huggingface_hub

app.py

@@ -82,7 +82,7 @@ class ConceptAlignerModel:
         self.text_encoder.load_state_dict(adapter_state, strict=True)
         print("    ✓ T5 encoder loaded")
 
-        # Only download VAE (small ~330MB) - not fine-tuned
+        # Only download VAE (small ~168MB)
         print("  Loading VAE from FLUX.1-dev...")
         vae = AutoencoderKL.from_pretrained(
             'black-forest-labs/FLUX.1-dev',
@@ -90,25 +90,23 @@ class ConceptAlignerModel:
             torch_dtype=self.dtype,
             token=HF_TOKEN
         ).to(self.device)
-        print("    ✓ VAE loaded (~330MB download)")
+        print("    ✓ VAE loaded")
 
-        # Create transformer architecture WITHOUT downloading base weights
-        print("  Initializing transformer architecture...")
-
-        # Get config only (no weights download)
-        from diffusers.models.transformers.transformer_flux import FluxTransformerConfig
+        # Create transformer from config only (download config.json but not weights)
+        print("  Downloading transformer config only...")
         config = FluxTransformer2DModel.load_config(
             'black-forest-labs/FLUX.1-dev',
             subfolder="transformer",
             token=HF_TOKEN
         )
 
-        # Initialize empty transformer from config
-        transformer = FluxTransformer2DModel.from_config(config)
-        print("    ✓ Transformer architecture initialized")
+        # Initialize transformer from config (no weights)
+        print("  Initializing transformer architecture from config...")
+        transformer = FluxTransformer2DModel.from_config(config, torch_dtype=self.dtype)
+        print("    ✓ Empty transformer initialized")
 
-        # Add LoRA config (needed for architecture)
-        print("  Adding LoRA adapter config...")
+        # Add LoRA adapter config
+        print("  Adding LoRA adapter layers...")
         transformer_lora_config = LoraConfig(
             r=256, lora_alpha=256, lora_dropout=0.0, init_lora_weights="gaussian",
             target_modules=[
@@ -120,15 +118,25 @@ class ConceptAlignerModel:
         )
         transformer.add_adapter(transformer_lora_config)
         transformer.context_embedder.requires_grad_(True)
+        print("    ✓ LoRA adapters added")
 
-        # Load YOUR FULL fine-tuned transformer weights (no base model needed!)
-        print("  Loading YOUR fine-tuned transformer weights...")
+        # Load YOUR FULL fine-tuned transformer weights
+        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)
+
+        # Load with strict=False in case of minor key mismatches
+        missing_keys, unexpected_keys = transformer.load_state_dict(transformer_state, strict=False)
+
+        if missing_keys:
+            print(f"    ⚠️  Missing keys: {len(missing_keys)}")
+        if unexpected_keys:
+            print(f"    ⚠️  Unexpected keys: {len(unexpected_keys)}")
+
         transformer = transformer.to(self.device).to(self.dtype)
-        print("    ✓ Fine-tuned transformer loaded (~26GB from your checkpoint)")
+        print("    ✓ Fine-tuned transformer loaded")
 
         # Load empty pooled clip
+        print("  Loading empty pooled clip...")
         self.empty_pooled_clip = torch.load(
             os.path.join(self.checkpoint_path, "empty_pooled_clip.pt"),
             map_location=self.device,
@@ -136,8 +144,8 @@ class ConceptAlignerModel:
         ).to(self.dtype)
         print("    ✓ Empty pooled clip loaded")
 
-        # Create scheduler (just config, no weights)
-        print("  Loading scheduler config...")
+        # Create scheduler (just config)
+        print("  Loading scheduler...")
         noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             'black-forest-labs/FLUX.1-dev',
             subfolder="scheduler",
@@ -146,7 +154,7 @@ class ConceptAlignerModel:
         print("    ✓ Scheduler loaded")
 
         # Create pipeline
-        print("  Creating pipeline...")
+        print("  Assembling pipeline...")
         self.pipe = CustomFluxKontextPipeline(
             scheduler=noise_scheduler,
             aligner=self.model,
@@ -156,10 +164,8 @@ class ConceptAlignerModel:
         ).to(self.device)
 
         print("="*60)
-        print("✓ ALL MODELS LOADED SUCCESSFULLY!")
+        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():
@@ -211,7 +217,6 @@ with gr.Blocks(title="ConceptAligner", theme=gr.themes.Soft()) as demo:
     # 🎨 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():