Update app/services/image_service.py

app/services/image_service.py

@@ -2,16 +2,17 @@ import base64
 import logging
 import json
 from io import BytesIO
+from typing import Optional, Dict, Any, Tuple, List
+
 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
+# Import initialized clients (ensure flux_client is defined in app/core/clients.py)
+from app.core.clients import llm, prompt_template, genai_client, flux_client, genai_client1
 
 logger = logging.getLogger(__name__)
 
-# --- New Constant for Virtual Try-On Prompt ---
+# --- Virtual Try-On Instruction Template ---
 VIRTUAL_TRY_ON_PROMPT = """
 You are an expert image-editing agent. Perform a high-fidelity virtual try-on using two inputs:
 
@@ -37,14 +38,17 @@ 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.
 """
 
+
+# ===============================================================
+# 🔹 PROMPT ENHANCEMENT
+# ===============================================================
 def enhance_user_prompt(raw_prompt: str) -> str:
+    """Enhance user prompt using LLM"""
     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.")
@@ -53,129 +57,210 @@ def enhance_user_prompt(raw_prompt: str) -> str:
         logger.error(f"Error during prompt enhancement: {e}", exc_info=True)
         raise
 
-# Returns BytesIO object, not base64 string
+
+# ===============================================================
+# 🔹 IMAGE GENERATION (TEXT → IMAGE) WITH FALLBACK
+# ===============================================================
 def generate_image_from_text(image_prompt: str) -> tuple[Optional[str], Optional[BytesIO]]:
-    """Generates an image using the GenAI client, returning text and BytesIO."""
+    """Generate an image from a text prompt with Gemini → Flux fallback"""
     logger.info(f"Generating image with prompt: {image_prompt[:50]}...")
-    
-    generation_config = types.GenerateContentConfig(
-        response_modalities=['Text', 'Image']
-    )
-    
+
+    generation_config = types.GenerateContentConfig(response_modalities=['Text', 'Image'])
+
     try:
+        # Try Gemini first
         response = genai_client.models.generate_content(
-            model="gemini-2.0-flash-exp-image-generation",
+            model="gemini-2.0-flash-preview-image-generation",
             contents=image_prompt,
             config=generation_config
         )
-        
-        logger.info("Image generation call successful.")
-        
-        generated_text = None
-        generated_image_bytes = None
-        
+
+        logger.info("Gemini image generation successful.")
+
+        generated_text, generated_image_bytes = None, None
         for part in response.candidates[0].content.parts:
-            if part.text is not None:
+            if part.text:
                 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
+            elif getattr(part, "inline_data", None):
                 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
+        logger.warning(f"Gemini image generation failed: {e}. Falling back to Flux...")
+
+        try:
+            image = flux_client.text_to_image(
+                image_prompt,
+                model="black-forest-labs/FLUX.1-dev"
+            )
+            buf = BytesIO()
+            image.save(buf, format="PNG")
+            buf.seek(0)
+            return None, buf
+        except Exception as flux_error:
+            logger.error(f"Flux fallback failed: {flux_error}", exc_info=True)
+            raise
+
 
-# Returns BytesIO object, not base64 string
+# ===============================================================
+# 🔹 IMAGE UPDATE (EDIT EXISTING IMAGE WITH TEXT)
+# ===============================================================
 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...")
-    
+    """Update an existing image using Gemini, fallback to Flux if fails."""
+    logger.info("Opening image 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']
-    )
-    
+    except Exception:
+        raise ValueError("Invalid image data. Upload a valid image file.")
+
+    generation_config = types.GenerateContentConfig(response_modalities=['Text', 'Image'])
+
     try:
         response = genai_client.models.generate_content(
-            model="gemini-2.0-flash-exp-image-generation",
+            model="gemini-2.0-flash-preview-image-generation",
             contents=[text_instruction, image],
             config=generation_config
         )
-        
-        logger.info("Image update call successful.")
-        
-        updated_text = None
-        updated_image_bytes = None
-        
+
+        logger.info("Gemini image update successful.")
+
+        updated_text, updated_image_bytes = None, None
         for part in response.candidates[0].content.parts:
-            if part.text is not None:
+            if part.text:
                 updated_text = part.text
-            elif part.inline_data is not None:
-                logger.info("Processing updated image data...")
+            elif getattr(part, "inline_data", None):
                 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
+        logger.warning(f"Gemini update failed: {e}. Falling back to Flux edit...")
+
+        try:
+            image = flux_client.text_to_image(
+                f"Edit image based on instruction: {text_instruction}",
+                model="black-forest-labs/FLUX.1-dev"
+            )
+            buf = BytesIO()
+            image.save(buf, format="PNG")
+            buf.seek(0)
+            return None, buf
+        except Exception as flux_error:
+            logger.error(f"Flux fallback failed: {flux_error}", exc_info=True)
+            raise
+
+
+# ===============================================================
+# 🔹 VIRTUAL TRY-ON (DRESS + PERSON)
+# ===============================================================
 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.")
+    """Perform virtual try-on with Gemini, fallback to Flux if fails."""
+    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.")
+    except Exception:
+        raise ValueError("Invalid image data provided.")
 
     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",
+        config = types.GenerateContentConfig(response_modalities=["Text", "Image"])
+        response = genai_client1.models.generate_content(
+            model="gemini-2.0-flash-preview-image-generation",
             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."}
+
+        logger.info("Gemini virtual try-on successful.")
+
+        result_summary = {"success": False, "notes": "No text response."}
         result_image_bytes = None
-        
+
         for part in response.candidates[0].content.parts:
-            if part.text is not None:
+            if part.text:
                 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_summary = {"success": False, "notes": part.text}
+            elif getattr(part, "inline_data", None):
                 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
+        logger.warning(f"Gemini try-on failed: {e}. Falling back to Flux...")
+
+        # try:
+        #     prompt = "Photorealistic virtual try-on of a person wearing the given dress."
+        #     image = flux_client.text_to_image(
+        #         prompt,
+        #         model="black-forest-labs/FLUX.1-dev"
+        #     )
+        #     buf = BytesIO()
+        #     image.save(buf, format="PNG")
+        #     buf.seek(0)
+        #     return {"success": True, "notes": "Generated using Flux fallback."}, buf
+        # except Exception as flux_error:
+        #     logger.error(f"Flux fallback failed: {flux_error}", exc_info=True)
+        #     raise
+
+
+# ===============================================================
+# 🔹 SHOE IMAGE GENERATION
+# ===============================================================
+def generate_shoe_images(
+    style: str,
+    color: Optional[str] = None,
+    perspective: Optional[str] = None,
+    variations: int = 1,
+) -> Tuple[Optional[str], List[BytesIO]]:
+    """Generate photorealistic shoe images with Gemini → Flux fallback."""
+    logger.info("Generating shoe images...")
+
+    shoe_prompt = f"""
+Produce a photorealistic product image of a shoe.
+
+Details:
+- Style: {style}
+- Color: {color or 'natural tones'}
+- Perspective: {perspective or 'three-quarter'}
+- Background: plain white / studio.
+- High detail, realistic textures, clean lighting.
+"""
+
+    config = types.GenerateContentConfig(response_modalities=["Text", "Image"])
+
+    try:
+        response = genai_client.models.generate_content(
+            model="gemini-2.0-flash-preview-image-generation",
+            contents=shoe_prompt,
+            config=config
+        )
+
+        generated_text, images = None, []
+
+        for part in response.candidates[0].content.parts:
+            if part.text:
+                generated_text = part.text
+            elif getattr(part, "inline_data", None):
+                images.append(BytesIO(part.inline_data.data))
+
+        return generated_text, images
+
+    except Exception as e:
+        logger.warning(f"Gemini shoe generation failed: {e}. Falling back to Flux...")
+
+        try:
+            image = flux_client.text_to_image(
+                shoe_prompt,
+                model="black-forest-labs/FLUX.1-dev"
+            )
+            buf = BytesIO()
+            image.save(buf, format="PNG")
+            buf.seek(0)
+            return None, [buf]
+        except Exception as flux_error:
+            logger.error(f"Flux fallback failed: {flux_error}", exc_info=True)
+            raise