diff --git a/backend/app/main.py b/backend/app/main.py index e69c7bd..f408e8f 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -5,7 +5,7 @@ from contextlib import asynccontextmanager from app.config import settings from app.database import init_db -from app.routers import projects, edits, images, patches +from app.routers import projects, edits, images, patches, generate @asynccontextmanager @@ -36,6 +36,7 @@ app.include_router(projects.router) app.include_router(edits.router) app.include_router(images.router) app.include_router(patches.router) +app.include_router(generate.router) @app.get("/") diff --git a/backend/app/routers/generate.py b/backend/app/routers/generate.py new file mode 100644 index 0000000..b247991 --- /dev/null +++ b/backend/app/routers/generate.py @@ -0,0 +1,176 @@ +from fastapi import APIRouter, Depends, HTTPException, Form +from sqlalchemy.orm import Session +from typing import Optional +from PIL import Image +from io import BytesIO +import os + +from app.database import get_db +from app.models.project import Project +from app.schemas import TextToImageRequest, TextToImageResponse +from app.services.ai_provider import get_ai_provider +from app.services.edit_service import EditService +from app.config import settings + +router = APIRouter(prefix="/generate", tags=["generate"]) + + +@router.post("/text-to-image", response_model=TextToImageResponse) +async def text_to_image( + prompt: str = Form(...), + width: int = Form(1024), + height: int = Form(1024), + negative_prompt: Optional[str] = Form(None), + ai_provider: Optional[str] = Form(None), + ai_model: Optional[str] = Form(None), + create_project: bool = Form(True), + project_name: Optional[str] = Form(None), + db: Session = Depends(get_db) +): + """ + Generate an image from text prompt + + Args: + prompt: Text description of desired image + width: Image width (default 1024) + height: Image height (default 1024) + negative_prompt: What to avoid in generation + ai_provider: Override default AI provider + ai_model: Specific model to use + create_project: Whether to create a new project with the result + project_name: Name for the new project (if create_project=True) + + Returns: + Generated image info and optionally project details + """ + + # Validate dimensions + if width < 256 or width > 2048 or height < 256 or height > 2048: + raise HTTPException( + status_code=400, + detail="Width and height must be between 256 and 2048" + ) + + # Get AI provider + provider = get_ai_provider(ai_provider, ai_model) + + try: + # Generate image + image_bytes = await provider.text_to_image( + prompt=prompt, + width=width, + height=height, + model=ai_model, + negative_prompt=negative_prompt + ) + + project_id = None + image_url = None + + if create_project: + # Create a new project + project = Project( + name=project_name or f"Generated: {prompt[:50]}", + user_id=None # TODO: Add authentication + ) + db.add(project) + db.commit() + db.refresh(project) + + project_id = project.id + + # Save image as both original and current + edit_service = EditService() + edit_service.ensure_project_dir(project_id) + + original_path = edit_service.get_original_image_path(project_id) + current_path = edit_service.get_current_image_path(project_id) + + # Save image + img = Image.open(BytesIO(image_bytes)) + img.save(original_path, 'PNG') + img.save(current_path, 'PNG') + + image_url = f"/projects/{project_id}/current" + + return TextToImageResponse( + status="success", + prompt=prompt, + width=width, + height=height, + project_id=project_id, + image_url=image_url, + ai_provider=ai_provider or settings.ai_provider, + ai_model=ai_model + ) + + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post("/layer/text-to-image", response_model=TextToImageResponse) +async def text_to_image_layer( + project_id: int = Form(...), + prompt: str = Form(...), + width: int = Form(512), + height: int = Form(512), + x: int = Form(0), + y: int = Form(0), + negative_prompt: Optional[str] = Form(None), + ai_provider: Optional[str] = Form(None), + ai_model: Optional[str] = Form(None), + db: Session = Depends(get_db) +): + """ + Generate an image as a new layer in an existing project + + This generates a smaller image that can be placed as a layer + on top of the current project canvas. + """ + + # Verify project exists + project = db.query(Project).filter(Project.id == project_id).first() + if not project: + raise HTTPException(status_code=404, detail="Project not found") + + # Get AI provider + provider = get_ai_provider(ai_provider, ai_model) + + try: + # Generate image + image_bytes = await provider.text_to_image( + prompt=prompt, + width=width, + height=height, + model=ai_model, + negative_prompt=negative_prompt + ) + + # Save as temporary layer file + edit_service = EditService() + layers_dir = edit_service.get_project_dir(project_id) / "layers" + layers_dir.mkdir(exist_ok=True) + + # Generate unique layer filename + import time + layer_filename = f"generated_{int(time.time())}.png" + layer_path = layers_dir / layer_filename + + # Save layer image + with open(layer_path, 'wb') as f: + f.write(image_bytes) + + return TextToImageResponse( + status="success", + prompt=prompt, + width=width, + height=height, + project_id=project_id, + image_url=f"/projects/{project_id}/layers/{layer_filename}", + layer_position={"x": x, "y": y}, + ai_provider=ai_provider or settings.ai_provider, + ai_model=ai_model + ) + + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) diff --git a/backend/app/schemas.py b/backend/app/schemas.py index b362f80..2943c69 100644 --- a/backend/app/schemas.py +++ b/backend/app/schemas.py @@ -107,6 +107,30 @@ class PatchApply(BaseModel): feather_px: int = 5 +# Text-to-Image schemas +class TextToImageRequest(BaseModel): + prompt: str + width: int = 1024 + height: int = 1024 + negative_prompt: Optional[str] = None + ai_provider: Optional[str] = None + ai_model: Optional[str] = None + create_project: bool = True + project_name: Optional[str] = None + + +class TextToImageResponse(BaseModel): + status: str + prompt: str + width: int + height: int + project_id: Optional[int] = None + image_url: Optional[str] = None + layer_position: Optional[Dict[str, int]] = None + ai_provider: str + ai_model: Optional[str] = None + + # Generic responses class StatusResponse(BaseModel): status: str diff --git a/backend/app/services/ai_provider.py b/backend/app/services/ai_provider.py index 6030b5f..1ca4587 100644 --- a/backend/app/services/ai_provider.py +++ b/backend/app/services/ai_provider.py @@ -36,6 +36,30 @@ class AIProvider(ABC): """ pass + @abstractmethod + async def text_to_image( + self, + prompt: str, + width: int = 1024, + height: int = 1024, + model: Optional[str] = None, + negative_prompt: Optional[str] = None + ) -> bytes: + """ + Generate an image from text prompt + + Args: + prompt: Text description of desired image + width: Image width in pixels + height: Image height in pixels + model: Optional specific model to use + negative_prompt: What to avoid in the generation + + Returns: + Generated image as bytes + """ + pass + class OpenAIProvider(AIProvider): """OpenAI DALL-E 2 based image editing (NOTE: Lower quality than DALL-E 3)""" @@ -88,6 +112,43 @@ class OpenAIProvider(AIProvider): return image_response.content + async def text_to_image( + self, + prompt: str, + width: int = 1024, + height: int = 1024, + model: Optional[str] = None, + negative_prompt: Optional[str] = None + ) -> bytes: + """Generate image using OpenAI DALL-E""" + + async with httpx.AsyncClient(timeout=60.0) as client: + data = { + 'prompt': prompt, + 'n': 1, + 'size': f'{width}x{height}' if width == height else '1024x1024' + } + + headers = { + 'Authorization': f'Bearer {self.api_key}' + } + + response = await client.post( + f"{self.base_url}/images/generations", + json=data, + headers=headers + ) + + response.raise_for_status() + result = response.json() + + # Download the generated image + image_url = result['data'][0]['url'] + image_response = await client.get(image_url) + image_response.raise_for_status() + + return image_response.content + class StabilityAIProvider(AIProvider): """Stability AI based image editing (SDXL Inpainting)""" @@ -154,6 +215,55 @@ class StabilityAIProvider(AIProvider): image_data = result['artifacts'][0]['base64'] return base64.b64decode(image_data) + async def text_to_image( + self, + prompt: str, + width: int = 1024, + height: int = 1024, + model: Optional[str] = None, + negative_prompt: Optional[str] = None + ) -> bytes: + """Generate image using Stability AI SDXL""" + + # Select model + model_key = model or self.default_model + engine_id = self.MODELS.get(model_key, self.MODELS['sdxl']) + + async with httpx.AsyncClient(timeout=120.0) as client: + # Build prompts array + data = { + 'text_prompts[0][text]': prompt, + 'text_prompts[0][weight]': '1.0', + 'cfg_scale': '7', + 'samples': '1', + 'steps': '50', + 'height': str(height), + 'width': str(width), + } + + # Add negative prompt if provided + if negative_prompt: + data['text_prompts[1][text]'] = negative_prompt + data['text_prompts[1][weight]'] = '-1.0' + + headers = { + 'Authorization': f'Bearer {self.api_key}', + 'Accept': 'application/json' + } + + response = await client.post( + f"{self.base_url}/generation/{engine_id}/text-to-image", + data=data, + headers=headers + ) + + response.raise_for_status() + result = response.json() + + # Decode base64 image + image_data = result['artifacts'][0]['base64'] + return base64.b64decode(image_data) + class ReplicateProvider(AIProvider): """Replicate API with multiple model support""" @@ -278,6 +388,77 @@ class ReplicateProvider(AIProvider): raise Exception("Replicate prediction timed out") + async def text_to_image( + self, + prompt: str, + width: int = 1024, + height: int = 1024, + model: Optional[str] = None, + negative_prompt: Optional[str] = None + ) -> bytes: + """Generate image using Replicate SDXL""" + + # Use SDXL for text-to-image + model_version = 'stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b' + + async with httpx.AsyncClient(timeout=120.0) as client: + # Create prediction + prediction_data = { + "version": model_version, + "input": { + "prompt": prompt, + "width": width, + "height": height, + "num_outputs": 1, + "guidance_scale": 7.5, + "num_inference_steps": 50, + } + } + + # Add negative prompt if provided + if negative_prompt: + prediction_data["input"]["negative_prompt"] = negative_prompt + + headers = { + 'Authorization': f'Bearer {self.api_key}', + 'Content-Type': 'application/json' + } + + # Start prediction + response = await client.post( + f"{self.base_url}/predictions", + json=prediction_data, + headers=headers + ) + response.raise_for_status() + prediction = response.json() + + # Poll for completion + prediction_url = prediction['urls']['get'] + max_attempts = 60 + attempt = 0 + + while attempt < max_attempts: + await asyncio.sleep(2) + + status_response = await client.get(prediction_url, headers=headers) + status_response.raise_for_status() + status_data = status_response.json() + + if status_data['status'] == 'succeeded': + # Download result image + output_url = status_data['output'][0] + image_response = await client.get(output_url) + image_response.raise_for_status() + return image_response.content + + elif status_data['status'] == 'failed': + raise Exception(f"Replicate prediction failed: {status_data.get('error')}") + + attempt += 1 + + raise Exception("Replicate text-to-image timed out") + class MockAIProvider(AIProvider): """Mock provider for testing (returns original patch)""" @@ -294,6 +475,30 @@ class MockAIProvider(AIProvider): """Return the original patch (for testing)""" return patch_image_bytes + async def text_to_image( + self, + prompt: str, + width: int = 1024, + height: int = 1024, + model: Optional[str] = None, + negative_prompt: Optional[str] = None + ) -> bytes: + """Generate a placeholder image (for testing)""" + from PIL import Image, ImageDraw, ImageFont + + # Create a simple placeholder image + img = Image.new('RGB', (width, height), color='lightgray') + draw = ImageDraw.Draw(img) + + # Draw text + text = f"Mock Image\n{width}x{height}\n{prompt[:50]}" + draw.text((width//4, height//2), text, fill='black') + + # Convert to bytes + buffer = BytesIO() + img.save(buffer, format='PNG') + return buffer.getvalue() + def get_ai_provider(provider_name: Optional[str] = None, model: Optional[str] = None) -> AIProvider: """