Initial clean commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,14 @@
+# Python virtual environment
+venv/
+.venv/
+
+# Python cache files
+__pycache__/
+*.pyc
+
+# Environment variables (IMPORTANT for security)
+.env
+
+# IDE files (optional, but good practice)
+.idea/
+.vscode/

README.md

@@ -0,0 +1,10 @@
+---
+title: Virtual Try On
+emoji: 🐠
+colorFrom: blue
+colorTo: indigo
+sdk: docker
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app/api/image_router.py

@@ -0,0 +1,160 @@
+import logging
+from fastapi import APIRouter, HTTPException, Response, UploadFile, File, Form
+from fastapi.responses import StreamingResponse, JSONResponse
+from app.schemas import image as image_schemas
+from app.services import image_service
+from io import BytesIO
+from typing import Optional, Dict, Any
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(
+    prefix="/api/v1",
+    tags=["Image Generation"]
+)
+
+# --- Helper function for streaming image responses ---
+def create_image_streaming_response(image_data: BytesIO, media_type: str = "image/png"):
+    """
+    Creates a StreamingResponse for an image from BytesIO data.
+    """
+    if not image_data:
+        raise HTTPException(status_code=404, detail="Image not found or could not be generated.")
+    
+    image_data.seek(0) # Ensure the pointer is at the beginning of the BytesIO object
+    return StreamingResponse(image_data, media_type=media_type)
+
+# --- Endpoints ---
+
+@router.post("/enhance-prompt", response_model=image_schemas.EnhancePromptResponse)
+async def enhance_prompt(request: image_schemas.EnhancePromptRequest):
+    logger.info(f"Received request to /enhance-prompt for: {request.raw_prompt[:50]}...")
+    try:
+        enhanced_prompt = image_service.enhance_user_prompt(request.raw_prompt)
+        logger.info("Successfully enhanced prompt.")
+        return image_schemas.EnhancePromptResponse(
+            raw_prompt=request.raw_prompt,
+            enhanced_prompt=enhanced_prompt
+        )
+    except Exception as e:
+        logger.error(f"Error in /enhance-prompt: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.post("/generate-image", response_class=StreamingResponse) # Response class directly for image
+async def generate_image(request: image_schemas.GenerateImageRequest):
+    """
+    Generates an image from either a raw or an enhanced prompt.
+    Returns the image directly.
+    """
+    logger.info("Received request to /generate-image")
+    try:
+        if request.enhanced_prompt:
+            image_prompt = request.enhanced_prompt
+        elif request.raw_prompt:
+            image_prompt = request.raw_prompt
+        else:
+            logger.warning("Bad request to /generate-image: No prompt provided.")
+            raise HTTPException(status_code=400, detail="Either raw_prompt or enhanced_prompt must be provided.")
+        
+        # Service returns text and BytesIO
+        generated_text, image_bytes_io = image_service.generate_image_from_text(image_prompt)
+        
+        if not image_bytes_io:
+            logger.error("Image generation failed or returned no image.")
+            # If text was returned, perhaps send that in a JSON error response
+            if generated_text:
+                 raise HTTPException(status_code=500, detail=f"Image generation failed. Model response: {generated_text}")
+            raise HTTPException(status_code=500, detail="Image generation failed: No image data received.")
+
+        logger.info("Successfully generated image. Streaming response.")
+        return create_image_streaming_response(image_bytes_io)
+    
+    except HTTPException: # Re-raise HTTPExceptions directly
+        raise
+    except Exception as e:
+        logger.error(f"Error in /generate-image: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+# --- update_image (MODIFIED) ---
+@router.post("/update-image", response_class=StreamingResponse) # Response class is already correct
+async def update_image(
+    image: UploadFile = File(..., description="The image to update (PNG, JPG)"),
+    text_instruction: str = Form(..., description="The text instruction for what to change.")
+):
+    """
+    Updates an existing image using a text instruction.
+    Returns the updated image directly.
+    """
+    logger.info("Received request to /update-image")
+
+    # Check if the uploaded file is an image
+    if not image.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Image must be an image type (e.g., image/png, image/jpeg).")
+
+    try:
+        # Read the image bytes from the uploaded file
+        image_bytes = await image.read()
+
+        # Service returns text and BytesIO
+        updated_text, updated_image_bytes_io = image_service.update_image_with_text(
+            text_instruction=text_instruction,
+            image_bytes=image_bytes  # Pass the raw bytes to the service
+        )
+        
+        if not updated_image_bytes_io:
+            logger.error("Image update failed or returned no image.")
+            if updated_text:
+                raise HTTPException(status_code=500, detail=f"Image update failed. Model response: {updated_text}")
+            raise HTTPException(status_code=500, detail="Image update failed: No image data received.")
+            
+        logger.info("Successfully updated image. Streaming response.")
+        return create_image_streaming_response(updated_image_bytes_io)
+    
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error in /update-image: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+    
+    
+# --- New Virtual Try-On Endpoint with File Uploads ---
+@router.post("/virtual-try-on")
+async def virtual_try_on(
+    dress_image: UploadFile = File(..., description="The dress image for try-on (PNG, JPG)"),
+    person_image: UploadFile = File(..., description="The person image for try-on (PNG, JPG)")
+) -> Response: # Return type is Response because it can be JSON or Image
+    """
+    Performs a virtual try-on using an uploaded dress image and a person image.
+    Returns the try-on image directly or a JSON response with a summary if no image.
+    """
+    logger.info("Received request to /virtual-try-on with image uploads.")
+
+    if not dress_image.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Dress image must be an image type (e.g., image/png, image/jpeg).")
+    if not person_image.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Person image must be an image type (e.g., image/png, image/jpeg).")
+
+    try:
+        dress_image_bytes = await dress_image.read()
+        person_image_bytes = await person_image.read()
+
+        # Service returns JSON summary (dict) and BytesIO object
+        summary, try_on_image_bytes_io = image_service.virtual_try_on(
+            dress_image_bytes=dress_image_bytes,
+            person_image_bytes=person_image_bytes
+        )
+        
+        if try_on_image_bytes_io:
+            logger.info("Virtual try-on successful. Streaming image response.")
+            # If an image is returned, stream it directly
+            return create_image_streaming_response(try_on_image_bytes_io)
+        else:
+            logger.warning("Virtual try-on returned no image. Sending JSON summary.")
+            # If no image, return the JSON summary as a JSON response
+            return JSONResponse(content=summary)
+            
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error in /virtual-try-on: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))

app/core/clients.py

@@ -0,0 +1,57 @@
+import os
+import logging
+from langchain_groq import ChatGroq
+from langchain_core.prompts import PromptTemplate
+from google import genai
+from .config import settings
+
+logger = logging.getLogger(__name__)
+
+# --- LLM (LangChain) Client ---
+
+def get_llm():
+    """Initializes and returns the ChatGroq LLM instance."""
+    logger.info("Initializing ChatGroq LLM...")
+    return ChatGroq(
+        model="meta-llama/llama-4-scout-17b-16e-instruct",
+        temperature=1,
+        max_tokens=1024,
+        api_key=settings.LANGCHAIN_API_KEY
+    )
+
+def get_prompt_template():
+    """Returns the prompt enhancement template."""
+    return PromptTemplate.from_template('''
+You are a Prompt Enhancement AI Assistant. Your task is to take the
+user's raw image prompt and convert it into a detailed,
+professional prompt optimized for generating high-quality AI image.
+Enhance the prompt by including relevant details such as:
+- Camera specifications (e.g., lens type, aperture, focal length)
+- Lighting setup (e.g., natural light, studio lighting, soft shadows)
+- Camera angle (e.g., top-down, macro, isometric, side view)
+- Background style (e.g., plain white, minimalistic, outdoor, studio backdrop)
+- Scene composition (e.g., centered product, depth of field, reflections)
+Focus only on *Poster photography* — do not include humans or models.
+Raw Prompt:
+{Raw_Prompt}
+Enhanced Prompt:
+''')
+
+# --- Google GenAI Client (FIXED) ---
+def get_genai_client():
+    """Initializes and returns the Google GenAI client (older SDK version)."""
+    try:
+        logger.info("Initializing Google GenAI Client...")
+        # Use the genai.Client from your original code
+        client = genai.Client(api_key=settings.GOOGLE_API_KEY)
+        logger.info("Google GenAI Client initialized successfully.")
+        return client
+    except Exception as e:
+        logger.error(f"Failed to initialize GenAI client: {e}")
+        raise
+
+# --- Initialize Clients ---
+llm = get_llm()
+prompt_template = get_prompt_template()
+# Export the client object
+genai_client = get_genai_client()

app/core/config.py

@@ -0,0 +1,11 @@
+from pydantic_settings import BaseSettings
+
+class Settings(BaseSettings):
+    LANGCHAIN_API_KEY: str
+    GOOGLE_API_KEY: str
+
+    class Config:
+        env_file = ".env"
+
+# Create a single, validated settings instance
+settings = Settings()

app/main.py

@@ -0,0 +1,35 @@
+import uvicorn
+import logging
+import sys
+from fastapi import FastAPI
+from app.api.image_router import router as image_router
+
+# --- Logging Configuration ---
+# We configure logging to output to stdout (standard output)
+# This is best practice for services, as the environment (like Docker or Uvicorn)
+# can then capture and manage the logs.
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+# ---
+
+# Get a logger for this file
+logger = logging.getLogger(__name__)
+
+# Initialize the FastAPI app
+app = FastAPI(title="Image Generation & Update API")
+
+# Include your API routers
+app.include_router(image_router)
+
+# 1. Root endpoint
+@app.get("/")
+async def root():
+    logger.info("Root endpoint '/' was hit")
+    return {"message": "Welcome to the Image Generation API! Visit /docs for API documentation."}
+
+# --- Run the app ---
+if __name__ == "__main__":
+    uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True)

app/schemas/image.py

@@ -0,0 +1,43 @@
+from pydantic import BaseModel
+from typing import Optional, Dict, Any
+from fastapi import UploadFile, File # Import UploadFile and File
+
+# --- Request Models ---
+
+class EnhancePromptRequest(BaseModel):
+    raw_prompt: str
+
+class GenerateImageRequest(BaseModel):
+    raw_prompt: Optional[str] = None
+    enhanced_prompt: Optional[str] = None
+
+# UpdateImageRequest still needs base64 as it receives existing image for modification
+class UpdateImageRequest(BaseModel):
+    text_instruction: str
+    image_base64: str
+
+# VirtualTryOnRequest now expects UploadFile objects
+# Note: For file uploads, it's often simpler to define them directly in the API endpoint function
+# using `File()` or `Form()` for metadata, rather than in a Pydantic model for the request body.
+# However, for consistency, we'll keep a schema here and handle parsing in the router.
+# But for multipart/form-data, Pydantic models for request *body* are slightly different.
+# The endpoint itself will directly take the UploadFile.
+
+
+# --- Response Models ---
+
+class EnhancePromptResponse(BaseModel):
+    raw_prompt: str
+    enhanced_prompt: str
+
+# These response models no longer contain image_base64
+# The image itself will be the response body.
+class GenerateImageOnlyResponse(BaseModel):
+    text: Optional[str] = None # Text can still be returned if the model generates it
+
+class UpdateImageOnlyResponse(BaseModel):
+    text: Optional[str] = None
+
+class VirtualTryOnResponse(BaseModel):
+    summary: Dict[str, Any]
+    # try_on_image will be the response body, so no base64 here

app/services/image_service.py

@@ -0,0 +1,181 @@
+import base64
+import logging
+import json
+from io import BytesIO
+from PIL import Image
+from google.genai import types
+from typing import Optional, Dict, Any
+
+# Import our initialized clients
+from app.core.clients import llm, prompt_template, genai_client
+
+logger = logging.getLogger(__name__)
+
+# --- New Constant for Virtual Try-On Prompt ---
+VIRTUAL_TRY_ON_PROMPT = """
+You are an expert image-editing agent. Perform a high-fidelity virtual try-on using two inputs:
+
+- Input 1 (dress reference): the exact garment image to be transferred.
+- Input 2 (person): the target person who must wear the garment.
+
+Positive instructions (Do):
+1. Produce exactly one photorealistic output image showing the person wearing the dress from Input 1.
+2. Preserve the person's face, facial features, expression, hair, gender, skin tone, and body proportions — do NOT change identity.
+3. Reproduce the dress design, color, pattern, fabric texture, and major details (buttons, seams, prints) faithfully — do NOT change these.
+4. Match lighting, shadows, perspective, and scale so the garment appears naturally worn by the person.
+5. If minor geometric adjustments are required to account for pose, make them subtle and keep dress details intact.
+6. Also return a short text JSON summary with the form: {"success": true|false, "notes": "explain any limitations or changes"}.
+
+Negative instructions (Do NOT):
+1. Do NOT modify the person's face, gender, or identifying features.
+2. Do NOT change the dress color, pattern, or main texture.
+3. Do NOT add extra clothing items, logos, watermarks, offensive symbols, or unrelated props.
+4. Do NOT produce empty responses, placeholders, or images containing text overlays.
+5. Do NOT blur, heavily distort, or crop important parts of the person or dress.
+
+Output requirements:
+- Return a single photorealistic image (same orientation as the person image) and a short text JSON summary.
+- No collages, no multiple thumbnail images, no visible watermarks.
+- If the try-on cannot be done reliably, set "success": false and explain why in "notes".
+
+When uncertain, err on the side of preserving the person's identity and the dress fidelity.
+"""
+
+def enhance_user_prompt(raw_prompt: str) -> str:
+    logger.info(f"Enhancing prompt: {raw_prompt[:50]}...")
+    formatted_prompt = prompt_template.invoke({"Raw_Prompt": raw_prompt})
+    
+    try:
+        response = llm.invoke(formatted_prompt)
+        logger.info("Prompt enhancement successful.")
+        return response.content
+    except Exception as e:
+        logger.error(f"Error during prompt enhancement: {e}", exc_info=True)
+        raise
+
+# Returns BytesIO object, not base64 string
+def generate_image_from_text(image_prompt: str) -> tuple[Optional[str], Optional[BytesIO]]:
+    """Generates an image using the GenAI client, returning text and BytesIO."""
+    logger.info(f"Generating image with prompt: {image_prompt[:50]}...")
+    
+    generation_config = types.GenerateContentConfig(
+        response_modalities=['Text', 'Image']
+    )
+    
+    try:
+        response = genai_client.models.generate_content(
+            model="gemini-2.0-flash-exp-image-generation",
+            contents=image_prompt,
+            config=generation_config
+        )
+        
+        logger.info("Image generation call successful.")
+        
+        generated_text = None
+        generated_image_bytes = None
+        
+        for part in response.candidates[0].content.parts:
+            if part.text is not None:
+                generated_text = part.text
+            elif part.inline_data is not None:
+                logger.info("Processing generated image data...")
+                # Image is directly from model, convert to BytesIO
+                generated_image_bytes = BytesIO(part.inline_data.data)
+                
+        return generated_text, generated_image_bytes
+    except Exception as e:
+        logger.error(f"Error during image generation: {e}", exc_info=True)
+        raise
+
+# Returns BytesIO object, not base64 string
+def update_image_with_text(text_instruction: str, image_bytes: bytes) -> tuple[Optional[str], Optional[BytesIO]]:
+    """Updates an existing image based on a text instruction, returning text and BytesIO."""
+    logger.info("Opening image from bytes for update...")
+    
+    try:
+        # MODIFICATION: Open image directly from bytes, not base64
+        image = Image.open(BytesIO(image_bytes))
+    except Exception as e:
+        logger.error(f"Failed to open image from bytes: {e}", exc_info=True)
+        raise ValueError("Invalid image data. The uploaded file is corrupt or not an image.")
+
+    logger.info(f"Updating image with instruction: {text_instruction[:50]}...")
+    
+    generation_config = types.GenerateContentConfig(
+        response_modalities=['Text', 'Image']
+    )
+    
+    try:
+        response = genai_client.models.generate_content(
+            model="gemini-2.0-flash-exp-image-generation",
+            contents=[text_instruction, image],
+            config=generation_config
+        )
+        
+        logger.info("Image update call successful.")
+        
+        updated_text = None
+        updated_image_bytes = None
+        
+        for part in response.candidates[0].content.parts:
+            if part.text is not None:
+                updated_text = part.text
+            elif part.inline_data is not None:
+                logger.info("Processing updated image data...")
+                updated_image_bytes = BytesIO(part.inline_data.data)
+                
+        return updated_text, updated_image_bytes
+    except Exception as e:
+        logger.error(f"Error during image update: {e}", exc_info=True)
+        raise
+# --- New Virtual Try-On Service ---
+# Accepts raw bytes for dress and person images
+def virtual_try_on(dress_image_bytes: bytes, person_image_bytes: bytes) -> tuple[Dict[str, Any], Optional[BytesIO]]:
+    """Performs virtual try-on using raw image bytes, returning summary and BytesIO."""
+    
+    logger.info("Opening images for virtual try-on.")
+    try:
+        dress_image = Image.open(BytesIO(dress_image_bytes))
+        person_image = Image.open(BytesIO(person_image_bytes))
+    except Exception as e:
+        logger.error(f"Failed to open image from bytes: {e}", exc_info=True)
+        raise ValueError("Invalid image data. One or both uploaded files are corrupt or not images.")
+
+    try:
+        logger.info("Calling GenAI model for virtual try-on...")
+        config = types.GenerateContentConfig(
+            response_modalities=["Text", "Image"]
+        )
+        
+        response = genai_client.models.generate_content(
+            model="gemini-2.0-flash-exp",
+            contents=[VIRTUAL_TRY_ON_PROMPT, dress_image, person_image],
+            config=config
+        )
+        
+        logger.info("Virtual try-on call successful.")
+        
+        result_summary = {"success": False, "notes": "No text response from model."}
+        result_image_bytes = None
+        
+        for part in response.candidates[0].content.parts:
+            if part.text is not None:
+                try:
+                    result_summary = json.loads(part.text)
+                    logger.info("Parsed JSON summary from model.")
+                except json.JSONDecodeError:
+                    logger.warning(f"Model returned non-JSON text: {part.text}")
+                    result_summary = {
+                        "success": False, 
+                        "notes": "Model returned non-JSON text.", 
+                        "raw_text": part.text
+                    }
+            elif part.inline_data is not None:
+                logger.info("Processing generated try-on image data...")
+                result_image_bytes = BytesIO(part.inline_data.data)
+                
+        return result_summary, result_image_bytes
+        
+    except Exception as e:
+        logger.error(f"Error during virtual try-on model call: {e}", exc_info=True)
+        raise

requirements.txt

@@ -0,0 +1,61 @@
+annotated-doc==0.0.3
+annotated-types==0.7.0
+anyio==4.11.0
+cachetools==6.2.1
+certifi==2025.10.5
+charset-normalizer==3.4.4
+click==8.3.0
+distro==1.9.0
+fastapi==0.120.0
+google-ai-generativelanguage==0.6.15
+google-api-core==2.27.0
+google-api-python-client==2.185.0
+google-auth==2.41.1
+google-auth-httplib2==0.2.0
+google-genai==1.46.0
+google-generativeai==0.8.5
+googleapis-common-protos==1.71.0
+groq==0.33.0
+grpcio==1.76.0
+grpcio-status==1.71.2
+h11==0.16.0
+httpcore==1.0.9
+httplib2==0.31.0
+httptools==0.7.1
+httpx==0.28.1
+idna==3.11
+jsonpatch==1.33
+jsonpointer==3.0.0
+langchain-core==1.0.1
+langchain-groq==1.0.0
+langsmith==0.4.38
+orjson==3.11.4
+packaging==25.0
+pillow==12.0.0
+proto-plus==1.26.1
+protobuf==5.29.5
+pyasn1==0.6.1
+pyasn1_modules==0.4.2
+pydantic==2.12.3
+pydantic-settings==2.11.0
+pydantic_core==2.41.4
+pyparsing==3.2.5
+python-dotenv==1.2.0
+python-multipart==0.0.20
+PyYAML==6.0.3
+requests==2.32.5
+requests-toolbelt==1.0.0
+rsa==4.9.1
+sniffio==1.3.1
+starlette==0.48.0
+tenacity==9.1.2
+tqdm==4.67.1
+typing-inspection==0.4.2
+typing_extensions==4.15.0
+uritemplate==4.2.0
+urllib3==2.5.0
+uvicorn==0.38.0
+uvloop==0.22.1
+watchfiles==1.1.1
+websockets==15.0.1
+zstandard==0.25.0