Upload folder using huggingface_hub

app.py

@@ -1,11 +1,11 @@
 """
-ConceptAligner Hugging Face Demo - ZeroGPU Compatible
+ConceptAligner - Same GPU behavior as FLUX demo
+Models loaded at startup, GPU allocated only for inference
 """
 
 # CRITICAL: Import spaces FIRST
 import spaces
 
-# Now import everything else
 import torch
 import gradio as gr
 import os
@@ -17,7 +17,7 @@ from pipeline import CustomFluxKontextPipeline
 from diffusers import FluxTransformer2DModel, FlowMatchEulerDiscreteScheduler, AutoencoderKL
 from peft import LoraConfig
 
-# Login with token from environment
+# Login
 HF_TOKEN = os.environ.get("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
@@ -26,15 +26,16 @@ if HF_TOKEN:
 # Configuration
 MODEL_REPO = "Shaoan/ConceptAligner-Weights"
 CHECKPOINT_DIR = "./checkpoint"
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 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."""]
 ]
 
 def download_checkpoint():
-    """Download checkpoint files from HF model repo"""
+    """Download checkpoint files"""
     print("Downloading checkpoint files...")
-
     files = ["model.safetensors", "model_1.safetensors", "model_2.safetensors", "empty_pooled_clip.pt"]
     os.makedirs(CHECKPOINT_DIR, exist_ok=True)
 
@@ -48,118 +49,101 @@ def download_checkpoint():
                 local_dir=CHECKPOINT_DIR,
                 token=HF_TOKEN
             )
-            print(f"    ✓ {filename} downloaded")
-
-    print("✓ All checkpoint files ready!")
-
-# 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
-    )
+    print("✓ Checkpoint files ready!")
 
-    # 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 at startup
 download_checkpoint()
 
-# ZeroGPU decorator - this moves computation to GPU when called
-@spaces.GPU(duration=60)  # 60 seconds of GPU time per generation
+# Load models at startup (like FLUX does)
+print("Loading models...")
+
+# Load ConceptAligner
+aligner_model = ConceptAligner().to(device).to(dtype)
+adapter_state = load_file(os.path.join(CHECKPOINT_DIR, "model_1.safetensors"))
+aligner_model.load_state_dict(adapter_state, strict=True)
+print("  ✓ ConceptAligner")
+
+# Load T5 encoder
+text_encoder = LoraT5Embedder(device=device).to(dtype)
+adapter_state = load_file(os.path.join(CHECKPOINT_DIR, "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)
+print("  ✓ T5 Encoder")
+
+# Load VAE
+vae = AutoencoderKL.from_pretrained(
+    'black-forest-labs/FLUX.1-dev',
+    subfolder="vae",
+    torch_dtype=dtype,
+    token=HF_TOKEN
+).to(device)
+print("  ✓ VAE")
+
+# Load 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_DIR, "model.safetensors"))
+transformer.load_state_dict(transformer_state, strict=False)
+transformer = transformer.to(device).to(dtype)
+print("  ✓ Transformer")
+
+# Load scheduler
+noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+    'black-forest-labs/FLUX.1-dev',
+    subfolder="scheduler",
+    token=HF_TOKEN
+)
+
+# Create pipeline
+pipe = CustomFluxKontextPipeline(
+    scheduler=noise_scheduler,
+    aligner=aligner_model,
+    transformer=transformer,
+    vae=vae,
+    text_embedder=text_encoder,
+).to(device)
+
+print("✅ Models loaded and ready!")
+torch.cuda.empty_cache()
+
+# History tracking
+previous_image = None
+previous_prompt = None
+
+@spaces.GPU(duration=75)
 @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"""
+def generate_image(prompt, height=512, width=512, guidance_scale=3.5,
+                   true_cf_scale=1.0, num_inference_steps=20, seed=0,
+                   progress=gr.Progress(track_tqdm=True)):
+    """Generate image - models already loaded"""
+    global previous_image, previous_prompt
 
     if not prompt.strip():
-        return None, None, "Please enter a prompt"
+        return previous_image, None, previous_prompt or "No previous generation", seed
 
     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(
+        current_image = pipe(
             prompt=prompt,
             guidance_scale=guidance_scale,
             true_cfg_scale=true_cf_scale,
@@ -170,66 +154,147 @@ def generate_image(prompt, threshold=0.0, topk=0, height=512, width=512,
             generator=generator,
         ).images[0]
 
-        return None, image, prompt
+        # Store for comparison
+        prev_image = previous_image
+        prev_prompt = previous_prompt or "No previous generation"
+
+        previous_image = current_image
+        previous_prompt = prompt
+
+        return prev_image, current_image, prev_prompt, seed
 
     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)}"
+        print(f"❌ Error: {e}")
+        print(traceback.format_exc())
+        return previous_image, None, previous_prompt or "", seed
 
-# Create Gradio interface
-with gr.Blocks(title="ConceptAligner") as demo:
-    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..."
-            )
+def reset_history():
+    """Clear generation history"""
+    global previous_image, previous_prompt
+    previous_image = None
+    previous_prompt = None
+    return None, None, "No previous generation"
 
-            generate_btn = gr.Button("✨ Generate", variant="primary", size="lg")
+# Create Gradio interface
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1200px;
+}
+"""
 
-            with gr.Accordion("⚙️ Settings", open=True):
-                guidance_scale = gr.Slider(1.0, 10.0, value=3.5, step=0.5, label="Guidance Scale")
-                num_steps = gr.Slider(10, 50, value=20, step=1, label="Steps")
-                seed = gr.Number(value=0, label="Seed", precision=0)
+with gr.Blocks(css=css, title="ConceptAligner") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("""
+        # 🎨 ConceptAligner Image Generator
+        
+        Create stunning AI-generated images from text descriptions.
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                prompt_input = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=3,
+                    placeholder="Describe your image...",
+                    container=False,
+                )
 
-            with gr.Accordion("🔬 Advanced", open=False):
-                true_cfg_scale = gr.Slider(1.0, 10.0, value=1.0, step=0.5, label="True CFG")
-                threshold = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Threshold")
-                topk = gr.Slider(0, 300, value=0, step=1, label="Top-K")
                 with gr.Row():
-                    height = gr.Slider(256, 1024, value=512, step=64, label="Height")
-                    width = gr.Slider(256, 1024, value=512, step=64, label="Width")
-
-        with gr.Column(scale=2):
-            gr.Markdown("### 🖼️ Generated Image")
-            output_image = gr.Image(label="Output", type="pil", height=512)
-            status_text = gr.Textbox(label="Status", interactive=False, visible=False)
+                    generate_btn = gr.Button("✨ Generate", variant="primary", scale=3)
+                    reset_btn = gr.Button("🔄 Clear", variant="secondary", scale=1)
+
+                with gr.Accordion("⚙️ Settings", open=False):
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=2147483647,
+                        step=1,
+                        value=0,
+                    )
+
+                    guidance_scale = gr.Slider(
+                        label="Creativity Level",
+                        minimum=1.0,
+                        maximum=10.0,
+                        step=0.5,
+                        value=3.5,
+                    )
+
+                    num_inference_steps = gr.Slider(
+                        label="Quality (steps)",
+                        minimum=10,
+                        maximum=50,
+                        step=1,
+                        value=20,
+                    )
+
+                    with gr.Row():
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=1024,
+                            step=64,
+                            value=512,
+                        )
+
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=1024,
+                            step=64,
+                            value=512,
+                        )
+
+                    true_cfg_scale = gr.Slider(
+                        label="True CFG Scale",
+                        minimum=1.0,
+                        maximum=10.0,
+                        step=0.5,
+                        value=1.0,
+                        visible=False
+                    )
+
+            with gr.Column(scale=2):
+                gr.Markdown("### 📊 Your Generations")
 
-    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(
+                with gr.Row():
+                    with gr.Column():
+                        gr.Markdown("**Previous**")
+                        prev_image = gr.Image(label="Previous", show_label=False, type="pil", height=400)
+                        prev_prompt_display = gr.Textbox(
+                            label="Previous Prompt",
+                            lines=2,
+                            interactive=False,
+                            show_label=False
+                        )
+
+                    with gr.Column():
+                        gr.Markdown("**Latest**")
+                        current_image = gr.Image(label="Current", show_label=False, type="pil", height=400)
+
+        gr.Markdown("### 📝 Try This Example")
+        gr.Examples(
+            examples=EXAMPLE_PROMPTS,
+            inputs=prompt_input,
+            outputs=[prev_image, current_image, prev_prompt_display, seed],
+            fn=generate_image,
+            cache_examples=False
+        )
+
+    # Event handlers
+    gr.on(
+        triggers=[generate_btn.click, prompt_input.submit],
         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]
+        inputs=[prompt_input, height, width, guidance_scale, true_cfg_scale, num_inference_steps, seed],
+        outputs=[prev_image, current_image, prev_prompt_display, seed]
+    )
+
+    reset_btn.click(
+        fn=reset_history,
+        outputs=[prev_image, current_image, prev_prompt_display]
     )
 
 if __name__ == "__main__":