Upload folder using huggingface_hub

app.py

@@ -1,7 +1,11 @@
 """
-ConceptAligner Hugging Face Demo
+ConceptAligner Hugging Face Demo - ZeroGPU Compatible
 """
 
+# CRITICAL: Import spaces FIRST
+import spaces
+
+# Now import everything else
 import torch
 import gradio as gr
 import os
@@ -13,14 +17,6 @@ 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:
@@ -56,161 +52,151 @@ def download_checkpoint():
 
     print("✓ All checkpoint files ready!")
 
-class ConceptAlignerModel:
-    def __init__(self):
-        download_checkpoint()
-
-        self.checkpoint_path = CHECKPOINT_DIR
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        self.dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
-        self.previous_image = None
-        self.previous_prompt = None
-
-        self.setup_models()
-
-    def setup_models(self):
-        """Load all models"""
-        print(f"Loading models on {self.device}...")
-
-        # Load ConceptAligner
-        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 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:
-            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")
-
-        # Download VAE
-        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
-        ).to(self.device)
-        print("    ✓ VAE loaded")
-
-        # Create transformer from config
-        print("  Downloading transformer config...")
-        config = FluxTransformer2DModel.load_config(
-            'black-forest-labs/FLUX.1-dev',
-            subfolder="transformer",
-            token=HF_TOKEN
-        )
-
-        print("  Initializing transformer...")
-        transformer = FluxTransformer2DModel.from_config(config, torch_dtype=self.dtype)
-
-        print("  Adding LoRA adapters...")
-        transformer_lora_config = LoraConfig(
-            r=256, lora_alpha=256, lora_dropout=0.0, init_lora_weights="gaussian",
-            target_modules=[
-                "attn.to_k", "attn.to_q", "attn.to_v", "attn.to_out.0",
-                "attn.add_k_proj", "attn.add_q_proj", "attn.add_v_proj", "attn.to_add_out",
-                "ff.net.0.proj", "ff.net.2", "ff_context.net.0.proj", "ff_context.net.2",
-                "proj_mlp", "proj_out", "norm.linear", "norm1.linear"
-            ],
-        )
-        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=False)
-        transformer = transformer.to(self.device).to(self.dtype)
-        print("    ✓ Transformer loaded")
-
-        # Load 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)
-
-        # Create scheduler
-        noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            'black-forest-labs/FLUX.1-dev',
-            subfolder="scheduler",
-            token=HF_TOKEN
-        )
-
-        # Create pipeline
-        self.pipe = CustomFluxKontextPipeline(
-            scheduler=noise_scheduler,
-            aligner=self.model,
-            transformer=transformer,
-            vae=vae,
-            text_embedder=self.text_encoder,
-        ).to(self.device)
-
-        print("✅ ALL MODELS LOADED!")
-
-        if torch.cuda.is_available():
-            allocated = torch.cuda.memory_allocated(0) / 1024**3
-            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,
-                       guidance_scale=3.5, true_cf_scale=1.0, num_inference_steps=20, seed=1995):
-        if not prompt.strip():
-            return self.previous_image, None, self.previous_prompt or ""
-
-        try:
-            generator = torch.Generator(device=self.device).manual_seed(int(seed))
-            current_image = self.pipe(
-                prompt=prompt, guidance_scale=guidance_scale, true_cfg_scale=true_cf_scale,
-                max_sequence_length=512, num_inference_steps=num_inference_steps,
-                height=height, width=width, generator=generator,
-            ).images[0]
-
-            prev_image = self.previous_image
-            prev_prompt = self.previous_prompt or "No previous generation"
-
-            self.previous_image = current_image
-            self.previous_prompt = prompt
-
-            return prev_image, current_image, prev_prompt
-        except Exception as e:
-            import traceback
-            print(f"❌ Error: {e}")
-            print(traceback.format_exc())
-            return self.previous_image, None, self.previous_prompt or ""
-
-    def reset_history(self):
-        self.previous_image = None
-        self.previous_prompt = None
-        return None, None, "No previous generation"
-
-# Initialize model
-print("🚀 Initializing ConceptAligner...")
-model = ConceptAlignerModel()
-
-# Wrap generation function with @spaces.GPU if available
-if GPU_AVAILABLE:
-    generate_fn = spaces.GPU(model.generate_image)
-else:
-    generate_fn = model.generate_image
+# Global model variable
+model_pipeline = None
+
+def load_models():
+    """Load all models - called once at startup"""
+    global model_pipeline
+
+    if model_pipeline is not None:
+        return model_pipeline
+
+    print("🚀 Loading models...")
+
+    checkpoint_path = CHECKPOINT_DIR
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
+    # Load ConceptAligner
+    print("  Loading ConceptAligner...")
+    aligner_model = ConceptAligner().to(device).to(dtype)
+    adapter_state = load_file(os.path.join(checkpoint_path, "model_1.safetensors"))
+    aligner_model.load_state_dict(adapter_state, strict=True)
+
+    # Load T5 encoder
+    print("  Loading T5 encoder...")
+    text_encoder = LoraT5Embedder(device=device).to(dtype)
+    adapter_state = load_file(os.path.join(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"]
+    text_encoder.load_state_dict(adapter_state, strict=True)
+
+    # Load VAE
+    print("  Loading VAE...")
+    vae = AutoencoderKL.from_pretrained(
+        'black-forest-labs/FLUX.1-dev',
+        subfolder="vae",
+        torch_dtype=dtype,
+        token=HF_TOKEN
+    ).to(device)
+
+    # Load transformer
+    print("  Loading transformer...")
+    config = FluxTransformer2DModel.load_config(
+        'black-forest-labs/FLUX.1-dev',
+        subfolder="transformer",
+        token=HF_TOKEN
+    )
+
+    transformer = FluxTransformer2DModel.from_config(config, torch_dtype=dtype)
+
+    transformer_lora_config = LoraConfig(
+        r=256, lora_alpha=256, lora_dropout=0.0, init_lora_weights="gaussian",
+        target_modules=[
+            "attn.to_k", "attn.to_q", "attn.to_v", "attn.to_out.0",
+            "attn.add_k_proj", "attn.add_q_proj", "attn.add_v_proj", "attn.to_add_out",
+            "ff.net.0.proj", "ff.net.2", "ff_context.net.0.proj", "ff_context.net.2",
+            "proj_mlp", "proj_out", "norm.linear", "norm1.linear"
+        ],
+    )
+    transformer.add_adapter(transformer_lora_config)
+    transformer.context_embedder.requires_grad_(True)
+
+    transformer_state = load_file(os.path.join(checkpoint_path, "model.safetensors"))
+    transformer.load_state_dict(transformer_state, strict=False)
+    transformer = transformer.to(device).to(dtype)
+
+    # Load scheduler
+    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        'black-forest-labs/FLUX.1-dev',
+        subfolder="scheduler",
+        token=HF_TOKEN
+    )
+
+    # Create pipeline
+    pipeline = CustomFluxKontextPipeline(
+        scheduler=noise_scheduler,
+        aligner=aligner_model,
+        transformer=transformer,
+        vae=vae,
+        text_embedder=text_encoder,
+    ).to(device)
+
+    model_pipeline = pipeline
+    print("✅ Models loaded!")
+    return pipeline
+
+# Download checkpoint at startup
+download_checkpoint()
+
+# ZeroGPU decorator - this moves computation to GPU when called
+@spaces.GPU(duration=60)  # 60 seconds of GPU time per generation
+@torch.no_grad()
+def generate_image(prompt, threshold=0.0, topk=0, height=512, width=512,
+                   guidance_scale=3.5, true_cf_scale=1.0, num_inference_steps=20, seed=1995):
+    """Generate image using the model"""
+
+    if not prompt.strip():
+        return None, None, "Please enter a prompt"
+
+    try:
+        # Load models (will use cached version after first call)
+        pipe = load_models()
+
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        generator = torch.Generator(device=device).manual_seed(int(seed))
+
+        print(f"Generating image: {prompt[:50]}...")
+
+        image = pipe(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            true_cfg_scale=true_cf_scale,
+            max_sequence_length=512,
+            num_inference_steps=num_inference_steps,
+            height=height,
+            width=width,
+            generator=generator,
+        ).images[0]
+
+        return None, image, prompt
+
+    except Exception as e:
+        import traceback
+        error_msg = f"❌ Error: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
+        return None, None, f"Error: {str(e)}"
 
 # Create Gradio interface
 with gr.Blocks(title="ConceptAligner") 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!
+    
+    ⚡ Running on ZeroGPU - GPU allocated on-demand for each generation
+    """)
 
     with gr.Row():
         with gr.Column(scale=1):
-            prompt_input = gr.Textbox(label="Prompt", lines=6, placeholder="Describe your image...")
+            prompt_input = gr.Textbox(
+                label="Prompt",
+                lines=6,
+                placeholder="Describe your image..."
+            )
 
-            with gr.Row():
-                generate_btn = gr.Button("✨ Generate", variant="primary", size="lg", scale=3)
-                reset_btn = gr.Button("🔄 Reset", variant="secondary", size="lg", scale=1)
+            generate_btn = gr.Button("✨ Generate", variant="primary", size="lg")
 
             with gr.Accordion("⚙️ Settings", open=True):
                 guidance_scale = gr.Slider(1.0, 10.0, value=3.5, step=0.5, label="Guidance Scale")
@@ -226,26 +212,25 @@ with gr.Blocks(title="ConceptAligner") as demo:
                     width = gr.Slider(256, 1024, value=512, step=64, label="Width")
 
         with gr.Column(scale=2):
-            gr.Markdown("### 📊 Comparison View")
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown("**Previous**")
-                    prev_image = gr.Image(label="Previous", type="pil", height=450)
-                    prev_prompt_display = gr.Textbox(label="Previous Prompt", lines=3, interactive=False)
-                with gr.Column():
-                    gr.Markdown("**Current**")
-                    current_image = gr.Image(label="Current", type="pil", height=450)
-
-    gr.Markdown("### 📝 Example")
+            gr.Markdown("### 🖼️ Generated Image")
+            output_image = gr.Image(label="Output", type="pil", height=512)
+            status_text = gr.Textbox(label="Status", interactive=False, visible=False)
+
+    gr.Markdown("### 📝 Example Prompt")
     gr.Examples(examples=EXAMPLE_PROMPTS, inputs=prompt_input)
 
+    # Hidden components for compatibility
+    prev_image_hidden = gr.Image(visible=False)
+    prev_prompt_hidden = gr.Textbox(visible=False)
+
     generate_btn.click(
-        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]
+        fn=generate_image,
+        inputs=[
+            prompt_input, threshold, topk,
+            height, width, guidance_scale, true_cfg_scale, num_steps, seed
+        ],
+        outputs=[prev_image_hidden, output_image, prev_prompt_hidden]
     )
 
-    reset_btn.click(fn=model.reset_history, outputs=[prev_image, current_image, prev_prompt_display])
-
 if __name__ == "__main__":
     demo.launch()