Upload folder using huggingface_hub

app.py

@@ -1,6 +1,5 @@
 """
-ConceptAligner Hugging Face Demo - Minimal downloads
-Only downloads VAE, uses your fine-tuned weights for everything else
+ConceptAligner Hugging Face Demo
 """
 
 import torch
@@ -14,6 +13,14 @@ from pipeline import CustomFluxKontextPipeline
 from diffusers import FluxTransformer2DModel, FlowMatchEulerDiscreteScheduler, AutoencoderKL
 from peft import LoraConfig
 
+# For HF Spaces GPU support
+try:
+    import spaces
+    GPU_AVAILABLE = True
+except ImportError:
+    GPU_AVAILABLE = False
+    print("⚠️  spaces package not available, running without @spaces.GPU decorator")
+
 # Login with token from environment
 HF_TOKEN = os.environ.get("HF_TOKEN")
 if HF_TOKEN:
@@ -73,7 +80,7 @@ class ConceptAlignerModel:
         self.model.load_state_dict(adapter_state, strict=True)
         print("    ✓ ConceptAligner loaded")
 
-        # Load T5 encoder (your fine-tuned version with full weights)
+        # Load T5 encoder
         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"))
@@ -82,7 +89,7 @@ class ConceptAlignerModel:
         self.text_encoder.load_state_dict(adapter_state, strict=True)
         print("    ✓ T5 encoder loaded")
 
-        # Only download VAE (small ~168MB)
+        # Download VAE
         print("  Loading VAE from FLUX.1-dev...")
         vae = AutoencoderKL.from_pretrained(
             'black-forest-labs/FLUX.1-dev',
@@ -92,21 +99,18 @@ class ConceptAlignerModel:
         ).to(self.device)
         print("    ✓ VAE loaded")
 
-        # Create transformer from config only (download config.json but not weights)
-        print("  Downloading transformer config only...")
+        # Create transformer from config
+        print("  Downloading transformer config...")
         config = FluxTransformer2DModel.load_config(
             'black-forest-labs/FLUX.1-dev',
             subfolder="transformer",
             token=HF_TOKEN
         )
 
-        # Initialize transformer from config (no weights)
-        print("  Initializing transformer architecture from config...")
+        print("  Initializing transformer...")
         transformer = FluxTransformer2DModel.from_config(config, torch_dtype=self.dtype)
-        print("    ✓ Empty transformer initialized")
 
-        # Add LoRA adapter config
-        print("  Adding LoRA adapter layers...")
+        print("  Adding LoRA adapters...")
         transformer_lora_config = LoraConfig(
             r=256, lora_alpha=256, lora_dropout=0.0, init_lora_weights="gaussian",
             target_modules=[
@@ -118,43 +122,28 @@ 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
-        print("  Loading your fine-tuned transformer weights...")
+        print("  Loading fine-tuned transformer weights...")
         transformer_state = load_file(os.path.join(self.checkpoint_path, "model.safetensors"))
-
-        # 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.load_state_dict(transformer_state, strict=False)
         transformer = transformer.to(self.device).to(self.dtype)
-        print("    ✓ Fine-tuned transformer loaded")
+        print("    ✓ 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,
             weights_only=True
         ).to(self.dtype)
-        print("    ✓ Empty pooled clip loaded")
 
-        # Create scheduler (just config)
-        print("  Loading scheduler...")
+        # Create scheduler
         noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             'black-forest-labs/FLUX.1-dev',
             subfolder="scheduler",
             token=HF_TOKEN
         )
-        print("    ✓ Scheduler loaded")
 
         # Create pipeline
-        print("  Assembling pipeline...")
         self.pipe = CustomFluxKontextPipeline(
             scheduler=noise_scheduler,
             aligner=self.model,
@@ -163,15 +152,11 @@ class ConceptAlignerModel:
             text_embedder=self.text_encoder,
         ).to(self.device)
 
-        print("="*60)
-        print("✅ ALL MODELS LOADED SUCCESSFULLY!")
-        print("="*60)
+        print("✅ ALL MODELS LOADED!")
 
-        # 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")
+            print(f"📊 GPU Memory: {allocated:.2f}GB allocated")
 
     @torch.no_grad()
     def generate_image(self, prompt, threshold=0.0, topk=0, height=512, width=512,
@@ -196,7 +181,7 @@ class ConceptAlignerModel:
             return prev_image, current_image, prev_prompt
         except Exception as e:
             import traceback
-            print(f"❌ Generation error: {e}")
+            print(f"❌ Error: {e}")
             print(traceback.format_exc())
             return self.previous_image, None, self.previous_prompt or ""
 
@@ -206,18 +191,18 @@ class ConceptAlignerModel:
         return None, None, "No previous generation"
 
 # Initialize model
-print("="*60)
-print("🚀 Initializing ConceptAligner Demo")
-print("="*60)
+print("🚀 Initializing ConceptAligner...")
 model = ConceptAlignerModel()
 
-# Create Gradio interface (without theme in constructor for Gradio 6.0)
+# Wrap generation function with @spaces.GPU if available
+if GPU_AVAILABLE:
+    generate_fn = spaces.GPU(model.generate_image)
+else:
+    generate_fn = model.generate_image
+
+# Create Gradio interface
 with gr.Blocks(title="ConceptAligner") as demo:
-    gr.Markdown("""
-    # 🎨 ConceptAligner Demo
-    
-    Generate images with fine-tuned concept alignment using FLUX!
-    """)
+    gr.Markdown("# 🎨 ConceptAligner Demo\nGenerate images with fine-tuned concept alignment!")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -255,7 +240,7 @@ with gr.Blocks(title="ConceptAligner") as demo:
     gr.Examples(examples=EXAMPLE_PROMPTS, inputs=prompt_input)
 
     generate_btn.click(
-        fn=model.generate_image,
+        fn=generate_fn,
         inputs=[prompt_input, threshold, topk, height, width, guidance_scale, true_cfg_scale, num_steps, seed],
         outputs=[prev_image, current_image, prev_prompt_display]
     )
@@ -263,9 +248,4 @@ with gr.Blocks(title="ConceptAligner") as demo:
     reset_btn.click(fn=model.reset_history, outputs=[prev_image, current_image, prev_prompt_display])
 
 if __name__ == "__main__":
-    # Launch with proper configuration for Hugging Face Spaces
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+    demo.launch()

requirements.txt

@@ -18,4 +18,5 @@ httpx==0.28.1
 requests==2.32.5
 numpy==1.26.4
 pydantic==2.11.9
-python-multipart==0.0.20
+python-multipart==0.0.20
+spaces