V3

backend/app/models/image.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Literal, Optional
 
 from pydantic import BaseModel, Field
 
@@ -12,6 +12,10 @@ class ImageGenerationRequest(BaseModel):
     project_id: str = Field(
         description="The project ID to associate with this image pair."
     )
+    type: Literal["generate", "edit"] = Field(
+        default="generate",
+        description="The type of operation: 'generate' for new images or 'edit' for modifying existing images.",
+    )
 
 
 class ImageGenerationResponse(BaseModel):

backend/app/services/image.py

@@ -6,6 +6,7 @@
 from PIL import Image
 
 from app.models.image import ImageGenerationRequest, ImageGenerationResponse
+from app.utils.prompts import EDIT_PROMPT, GENERATE_PROMPT
 
 log = logging.getLogger(__name__)
 
@@ -27,12 +28,11 @@ async def generate_image(
         Returns:
             ImageGenerationResponse containing the generated image data
         """
-        log.info(f"Generating image with prompt: {input.prompt}")
+        log.info(f"Generating image with type '{input.type}' and prompt: {input.prompt}")
 
-        prompt = f"""Generate a drawing image based on the following prompt and the reference image.
-        The image background should match the reference image background. You are just drawing diagrams, so make sure you are not over verbose.
-        But if there are paragraphs in the image, you should keep them there.
-        Prompt: {input.prompt}"""
+        # Select prompt template based on operation type
+        prompt_template = GENERATE_PROMPT if input.type == "generate" else EDIT_PROMPT
+        prompt = prompt_template.format(user_prompt=input.prompt)
 
         # Prepare reference image if provided
         reference_image = None

backend/app/utils/prompts.py

@@ -0,0 +1,17 @@
+"""
+Prompt templates for image generation and editing.
+"""
+
+GENERATE_PROMPT = """Generate a drawing image based on the following reference image of a drawing and the explanation by the user.
+The image background should match the reference image background. You are just drawing diagrams, so make sure you are not over verbose.
+But if there are paragraphs in the image, you should keep them there.
+
+The main goal should be to create a diagram. If the components or items of the diagram is explained in the user voice explanation, you should item or component name them in the diagram.
+User voice explanation: {user_prompt}"""
+
+EDIT_PROMPT = """Edit the provided drawing image based on the user's explanation and instructions.
+Maintain the overall structure and background of the original image while incorporating the requested changes.
+Focus on modifying or enhancing specific elements as described by the user.
+
+User edit instructions: {user_prompt}"""
+