Add text-to-image generation support
Implemented complete text-to-image functionality across all AI providers: Backend additions: - Added text_to_image() method to AIProvider abstract class - Implemented for all providers: * OpenAI: DALL-E generations API * Stability AI: SDXL text-to-image with negative prompts * Replicate: SDXL with full parameter control * Mock: Placeholder image generation for testing New API endpoints (/generate): - POST /generate/text-to-image * Generate image from prompt * Optional: create new project automatically * Configurable width/height (256-2048px) * Negative prompt support * Provider and model selection - POST /generate/layer/text-to-image * Generate image as layer in existing project * Smaller dimensions for layer composition * Position control (x, y coordinates) * Saves to project layers directory Features: - Full provider support (OpenAI, Stability, Replicate, Mock) - Negative prompts for better control - Auto-project creation option - Layer-based generation for compositing - Dimension validation (256-2048px range) - Model selection per request Use cases: - Create new images from scratch - Generate elements to add as layers - Quick ideation and iteration - Base image creation for further editing Next: Advanced canvas UI with layers and real-time preview
This commit is contained in:
Claude
+2
-1
@@ -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("/")
|
||||
|
||||
@@ -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))
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user