Update app.py

app.py

@@ -18,18 +18,98 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 # Initialize text encoder client ONCE at module level to avoid thread exhaustion
-text_encoder_client = client = Client("Gemini899/mistral-text-encoder")
+text_encoder_client = Client("Gemini899/mistral-text-encoder")
+
+# ============================================================================
+# HARDCODED PROMPTS - Exact match only
+# ============================================================================
+
+HARDCODED_PROMPTS = {
+    "relief": "Clay bas-relief sculpture. PRESERVE exact facial features and proportions. Uniform matte gray material, NO black areas, NO dark shadows, NO outlines. Soft smooth depth only. Light gray to white tones. Like carved marble or clay relief.",
+    
+    "details": "Enhance this depth map by adding surface micro-details (skin pores, fabric texture, hair strands, stone grain, wrinkles, fingernails, knuckles) using ONLY tonal variations within ±10% of local gray values. Keep EXACT same outline, silhouette, and overall tonal range. NO shadows, NO reflections, NO new light sources, NO dark areas under nose/eyes/lips. Output must overlay perfectly on original as bump map detail layer."
+}
+
+# Pre-load embeddings at startup
+_cached_embeddings = {}
+
+def load_cached_embeddings():
+    """Load pre-generated embeddings at startup."""
+    global _cached_embeddings
+    
+    embedding_files = {
+        "relief": "relief.pt",
+        "details": "details.pt"
+    }
+    
+    for key, filename in embedding_files.items():
+        # Try multiple possible paths
+        possible_paths = [
+            filename,  # Current directory
+            f"/home/user/app/{filename}",  # HF Spaces app directory
+            os.path.join(os.path.dirname(__file__), filename),  # Same dir as script
+        ]
+        
+        for path in possible_paths:
+            if os.path.exists(path):
+                try:
+                    _cached_embeddings[key] = torch.load(path, map_location='cpu')
+                    print(f"✓ Loaded cached embedding: {key} from {path}")
+                    break
+                except Exception as e:
+                    print(f"✗ Error loading {path}: {e}")
+        else:
+            print(f"⚠ Warning: {filename} not found - will use API for '{key}' prompt")
+
+def normalize_prompt(prompt: str) -> str:
+    """Normalize prompt by stripping whitespace for comparison."""
+    return prompt.strip()
+
+def get_cached_embedding(prompt: str) -> torch.Tensor | None:
+    """
+    Check if prompt EXACTLY matches a hardcoded prompt.
+    Returns cached embedding if exact match, None otherwise.
+    """
+    normalized_input = normalize_prompt(prompt)
+    
+    for key, hardcoded_prompt in HARDCODED_PROMPTS.items():
+        if normalized_input == normalize_prompt(hardcoded_prompt):
+            if key in _cached_embeddings:
+                print(f"⚡ Exact match found: using cached '{key}' embedding (no API call)")
+                return _cached_embeddings[key]
+            else:
+                print(f"⚠ Exact match for '{key}' but no cached embedding - using API")
+                return None
+    
+    return None
 
 def remote_text_encoder(prompts):
+    """
+    Encode text prompts to embeddings.
+    Uses cached embeddings for exact hardcoded prompt matches.
+    Falls back to Mistral API for all other prompts.
+    """
+    # Check for exact match with hardcoded prompts
+    cached = get_cached_embedding(prompts)
+    if cached is not None:
+        return cached
+    
+    # Not an exact match - use API
+    print(f"🌐 Calling Mistral API for prompt encoding...")
     result = text_encoder_client.predict(
         prompt=prompts,
         api_name="/encode_text"
     )
-
     prompt_embeds = torch.load(result[0])
     return prompt_embeds
 
-# Load model
+# Load cached embeddings at startup
+load_cached_embeddings()
+
+# ============================================================================
+# Model Loading
+# ============================================================================
+
 repo_id = "black-forest-labs/FLUX.2-dev"
 
 dit = Flux2Transformer2DModel.from_pretrained(
@@ -46,8 +126,9 @@ pipe = Flux2Pipeline.from_pretrained(
 )
 pipe.to(device)
 
-# AOTI blocks temporarily disabled - HuggingFace needs to recompile for new ZeroGPU environment
-# spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/FLUX.2", variant="fa3")
+# ============================================================================
+# Image Generation Functions
+# ============================================================================
 
 def update_dimensions_from_image(image_list):
     """Update width/height sliders based on uploaded image aspect ratio."""
@@ -97,7 +178,7 @@ def generate_image(prompt_embeds, image_list, width, height, num_inference_steps
     
     if progress:
         progress(0, desc="Starting generation...")
-        
+    
     image = pipe(**pipe_kwargs).images[0]
     return image
 
@@ -111,26 +192,30 @@ def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024,
         image_list = []
         for item in input_images:
             image_list.append(item[0])
-
-    # Text Encoding
+    
+    # Text Encoding (checks for cached embeddings first)
     progress(0.1, desc="Encoding prompt...")
     prompt_embeds = remote_text_encoder(prompt)
     
     # Image Generation
     progress(0.3, desc="Waiting for GPU...")
     image = generate_image(
-        prompt_embeds, 
-        image_list, 
-        width, 
-        height, 
-        num_inference_steps, 
-        guidance_scale, 
-        seed, 
+        prompt_embeds,
+        image_list,
+        width,
+        height,
+        num_inference_steps,
+        guidance_scale,
+        seed,
         progress
     )
     
     return image, seed
 
+# ============================================================================
+# Gradio UI
+# ============================================================================
+
 examples = [
     ["Create a vase on a table in living room, the color of the vase is a gradient of color, starting with #02eb3c color and finishing with #edfa3c. The flowers inside the vase have the color #ff0088"],
     ["Photorealistic infographic showing the complete Berlin TV Tower (Fernsehturm) from ground base to antenna tip, full vertical view with entire structure visible including concrete shaft, metallic sphere, and antenna spire. Slight upward perspective angle looking up toward the iconic sphere, perfectly centered on clean white background. Left side labels with thin horizontal connector lines: the text '368m' in extra large bold dark grey numerals (#2D3748) positioned at exactly the antenna tip with 'TOTAL HEIGHT' in small caps below. The text '207m' in extra large bold with 'TELECAFÉ' in small caps below, with connector line touching the sphere precisely at the window level. Right side label with horizontal connector line touching the sphere's equator: the text '32m' in extra large bold dark grey numerals with 'SPHERE DIAMETER' in small caps below. Bottom section arranged in three balanced columns: Left - Large text '986' in extra bold dark grey with 'STEPS' in caps below. Center - 'BERLIN TV TOWER' in bold caps with 'FERNSEHTURM' in lighter weight below. Right - 'INAUGURATED' in bold caps with 'OCTOBER 3, 1969' below. All typography in modern sans-serif font (such as Inter or Helvetica), color #2D3748, clean minimal technical diagram style. Horizontal connector lines are thin, precise, and clearly visible, touching the tower structure at exact corresponding measurement points. Professional architectural elevation drawing aesthetic with dynamic low angle perspective creating sense of height and grandeur, poster-ready infographic design with perfect visual hierarchy."],
@@ -157,7 +242,7 @@ with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""# FLUX.2 [dev]
 FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and combining images based on text instructions model [[model](https://huggingface.co/black-forest-labs/FLUX.2-dev)], [[blog](https://bfl.ai/blog/flux-2)]
-        """)
+""")
         with gr.Row():
             with gr.Column():
                 with gr.Row():
@@ -171,7 +256,7 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
                     )
                     
                     run_button = gr.Button("Run", scale=1)
-                    
+                
                 with gr.Accordion("Input image(s) (optional)", open=True):
                     input_images = gr.Gallery(
                         label="Input Image(s)",
@@ -236,7 +321,7 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
             cache_examples=True,
             cache_mode="lazy"
         )
-
+        
         gr.Examples(
             examples=examples_images,
             fn=infer,
@@ -245,13 +330,13 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
             cache_examples=True,
             cache_mode="lazy"
         )
-
+    
     input_images.upload(
         fn=update_dimensions_from_image,
         inputs=[input_images],
         outputs=[width, height]
     )
-
+    
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,