fix: bug fixes and improvements

app.py

@@ -0,0 +1,29 @@
+import os
+import sys
+import logging
+
+# Configure logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+logger = logging.getLogger("faceforge_app")
+
+# Add the project root to the Python path
+sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
+
+try:
+    logger.info("Starting FaceForge app")
+    # Import the demo from the UI module
+    from faceforge_ui.app import demo
+    
+    # Launch the app
+    if __name__ == "__main__":
+        logger.info("Launching Gradio interface")
+        demo.launch(server_name="0.0.0.0")
+except Exception as e:
+    logger.critical(f"Failed to start app: {e}")
+    import traceback
+    logger.debug(traceback.format_exc())
+    raise 

faceforge_api/main.py

@@ -1,13 +1,168 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import List, Optional
+import numpy as np
+import base64
+import logging
+import sys
+import traceback
+import io
+from PIL import Image
+
+from faceforge_core.latent_explorer import LatentSpaceExplorer
+from faceforge_core.attribute_directions import LatentDirectionFinder
+from faceforge_core.custom_loss import attribute_preserving_loss
+
+# Configure logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+logger = logging.getLogger("faceforge_api")
+
+# --- Models for API ---
+
+class PointIn(BaseModel):
+    text: str
+    encoding: Optional[List[float]] = None
+    xy_pos: Optional[List[float]] = None
+
+class GenerateRequest(BaseModel):
+    prompts: List[str]
+    positions: Optional[List[List[float]]] = None
+    mode: str = "distance"
+    player_pos: Optional[List[float]] = None
+
+class ManipulateRequest(BaseModel):
+    encoding: List[float]
+    direction: List[float]
+    alpha: float
+
+class AttributeDirectionRequest(BaseModel):
+    latents: List[List[float]]
+    labels: Optional[List[int]] = None
+    n_components: Optional[int] = 10
+
+# --- FastAPI app ---
 
 app = FastAPI()
 
+# Add CORS middleware to allow requests from any origin
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Global explorer instance
+explorer = LatentSpaceExplorer()
+
+# Error handling middleware
+@app.middleware("http")
+async def error_handling_middleware(request: Request, call_next):
+    try:
+        return await call_next(request)
+    except Exception as e:
+        logger.error(f"Unhandled exception: {str(e)}")
+        logger.debug(traceback.format_exc())
+        return JSONResponse(
+            status_code=500,
+            content={"detail": "Internal server error", "error": str(e)},
+        )
+
 @app.get("/")
 def read_root():
+    logger.debug("Root endpoint called")
     return {"message": "FaceForge API is running"}
 
 @app.post("/generate")
-def generate_image():
-    # Placeholder for image generation logic
-    return JSONResponse(content={"status": "success", "image": None}) 
+def generate_image(req: GenerateRequest):
+    try:
+        logger.debug(f"Generate image request: {req}")
+        
+        # Clear existing points
+        explorer.points = []
+        
+        # Add points for each prompt
+        for i, prompt in enumerate(req.prompts):
+            logger.debug(f"Processing prompt {i}: {prompt}")
+            
+            # Generate a mock encoding (in production, this would use a real model)
+            encoding = np.random.randn(512)  # Stub: replace with real encoding
+            
+            # Get position if provided, otherwise None
+            xy_pos = req.positions[i] if req.positions and i < len(req.positions) else None
+            logger.debug(f"Position for prompt {i}: {xy_pos}")
+            
+            # Add point to explorer
+            explorer.add_point(prompt, encoding, xy_pos)
+        
+        # Get player position
+        if req.player_pos is None:
+            player_pos = [0.0, 0.0]
+        else:
+            player_pos = req.player_pos
+        logger.debug(f"Player position: {player_pos}")
+        
+        # Sample encoding
+        logger.debug(f"Sampling with mode: {req.mode}")
+        sampled = explorer.sample_encoding(tuple(player_pos), mode=req.mode)
+        
+        # Generate mock image (in production, this would use the sampled encoding)
+        img = (np.random.rand(256, 256, 3) * 255).astype(np.uint8)
+        
+        # Convert to base64
+        logger.debug("Converting image to base64")
+        pil_img = Image.fromarray(img)
+        buffer = io.BytesIO()
+        pil_img.save(buffer, format="PNG")
+        img_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        
+        logger.debug("Image generated successfully")
+        return {"status": "success", "image": img_b64}
+        
+    except Exception as e:
+        logger.error(f"Error in generate_image: {str(e)}")
+        logger.debug(traceback.format_exc())
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/manipulate")
+def manipulate(req: ManipulateRequest):
+    try:
+        logger.debug(f"Manipulate request: {req}")
+        encoding = np.array(req.encoding)
+        direction = np.array(req.direction)
+        manipulated = encoding + req.alpha * direction
+        logger.debug("Manipulation successful")
+        return {"manipulated_encoding": manipulated.tolist()}
+    except Exception as e:
+        logger.error(f"Error in manipulate: {str(e)}")
+        logger.debug(traceback.format_exc())
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/attribute_direction")
+def attribute_direction(req: AttributeDirectionRequest):
+    try:
+        logger.debug(f"Attribute direction request: {req}")
+        latents = np.array(req.latents)
+        finder = LatentDirectionFinder(latents)
+        
+        if req.labels is not None:
+            logger.debug("Using classifier-based direction finding")
+            direction = finder.classifier_direction(req.labels)
+            logger.debug("Direction found successfully")
+            return {"direction": direction.tolist()}
+        else:
+            logger.debug(f"Using PCA with {req.n_components} components")
+            components, explained = finder.pca_direction(n_components=req.n_components)
+            logger.debug("PCA completed successfully")
+            return {"components": components.tolist(), "explained_variance": explained.tolist()}
+    except Exception as e:
+        logger.error(f"Error in attribute_direction: {str(e)}")
+        logger.debug(traceback.format_exc())
+        raise HTTPException(status_code=500, detail=str(e)) 

faceforge_ui/app.py

@@ -1,19 +1,154 @@
 import gradio as gr
 import requests
+import numpy as np
+from PIL import Image
+import io
+import base64
+import logging
+import sys
+import traceback
+import os
 
-def generate_image():
-    response = requests.post("http://localhost:8000/generate")
-    if response.ok:
-        return None  # Placeholder: should return image from backend
-    return None
-
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=[],
-    outputs="image",
-    title="FaceForge Latent Space Explorer",
-    description="Interactively explore and edit faces in latent space."
+# Configure logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
 )
+logger = logging.getLogger("faceforge_ui")
+
+# API configuration
+API_URL = os.environ.get("API_URL", "http://localhost:8000")
+logger.info(f"Using API URL: {API_URL}")
+
+def generate_image(prompts, mode, player_x, player_y):
+    """
+    Generate an image based on prompts and player position.
+    
+    Args:
+        prompts: Comma-separated list of prompts
+        mode: Sampling mode ('distance' or 'circle')
+        player_x: X-coordinate of player position
+        player_y: Y-coordinate of player position
+    
+    Returns:
+        PIL.Image or None: Generated image or None if generation failed
+    """
+    try:
+        logger.debug(f"Generating image with prompts: {prompts}, mode: {mode}, position: ({player_x}, {player_y})")
+        
+        # Parse prompts
+        prompt_list = [p.strip() for p in prompts.split(",") if p.strip()]
+        if not prompt_list:
+            logger.warning("No valid prompts provided")
+            return None
+        
+        logger.debug(f"Parsed prompts: {prompt_list}")
+        
+        # Prepare request
+        req = {
+            "prompts": prompt_list,
+            "mode": mode,
+            "player_pos": [float(player_x), float(player_y)]
+        }
+        
+        logger.debug(f"Sending request to API: {req}")
+        
+        # Make API call
+        try:
+            resp = requests.post(f"{API_URL}/generate", json=req, timeout=30)
+            logger.debug(f"API response status: {resp.status_code}")
+            
+            if resp.ok:
+                data = resp.json()
+                logger.debug("Successfully received API response")
+                
+                if "image" in data:
+                    img_b64 = data["image"]
+                    img_bytes = base64.b64decode(img_b64)
+                    
+                    try:
+                        img = Image.frombytes("RGB", (256, 256), img_bytes)
+                        logger.debug("Successfully decoded image")
+                        return img
+                    except Exception as e:
+                        logger.error(f"Error decoding image: {e}")
+                        logger.debug(traceback.format_exc())
+                        return None
+                else:
+                    logger.warning("No image in API response")
+                    return None
+            else:
+                logger.error(f"API error: {resp.status_code} - {resp.text}")
+                return None
+                
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Request failed: {e}")
+            logger.debug(traceback.format_exc())
+            return None
+            
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        logger.debug(traceback.format_exc())
+        return None
+
+# Create Gradio interface
+logger.info("Initializing Gradio interface")
+with gr.Blocks() as demo:
+    gr.Markdown("# FaceForge Latent Space Explorer")
+    
+    with gr.Row():
+        with gr.Column():
+            prompts = gr.Textbox(
+                label="Prompts (comma-separated)", 
+                value="A photo of a cat, A photo of a dog",
+                info="Enter prompts separated by commas"
+            )
+            mode = gr.Radio(
+                choices=["distance", "circle"], 
+                value="distance", 
+                label="Sampling Mode",
+                info="Choose how to sample the latent space"
+            )
+            player_x = gr.Slider(-1.0, 1.0, value=0.0, label="Player X")
+            player_y = gr.Slider(-1.0, 1.0, value=0.0, label="Player Y")
+            btn = gr.Button("Generate")
+        
+        with gr.Column():
+            img = gr.Image(label="Generated Image")
+            status = gr.Textbox(label="Status", interactive=False)
+    
+    def on_generate_click(prompts, mode, player_x, player_y):
+        try:
+            logger.info("Generate button clicked")
+            result = generate_image(prompts, mode, player_x, player_y)
+            if result is not None:
+                return [result, "Image generated successfully"]
+            else:
+                return [None, "Failed to generate image. Check logs for details."]
+        except Exception as e:
+            logger.error(f"Error in generate button handler: {e}")
+            logger.debug(traceback.format_exc())
+            return [None, f"Error: {str(e)}"]
+    
+    btn.click(
+        fn=on_generate_click,
+        inputs=[prompts, mode, player_x, player_y],
+        outputs=[img, status]
+    )
+    
+    demo.load(lambda: "Ready to generate images", outputs=status)
 
 if __name__ == "__main__":
-    iface.launch() 
+    logger.info("Starting Gradio app")
+    try:
+        # Check if we're running in Hugging Face Spaces
+        if "SPACE_ID" in os.environ:
+            logger.info("Running in Hugging Face Space")
+            demo.launch(server_name="0.0.0.0", share=False)
+        else:
+            logger.info("Running locally")
+            demo.launch(server_name="0.0.0.0", share=False)
+    except Exception as e:
+        logger.critical(f"Failed to launch Gradio app: {e}")
+        logger.debug(traceback.format_exc())