Add Image Generation from Prompts: Implement new functionality to generate images from prompts, including backend processing, API integration, and frontend handling with progress modal. Update settings and registry for new AI function.
This commit is contained in:
IGNY8 VPS (Salman)
@@ -6,6 +6,7 @@ from igny8_core.ai.functions.generate_ideas import GenerateIdeasFunction
|
||||
from igny8_core.ai.functions.generate_content import GenerateContentFunction
|
||||
from igny8_core.ai.functions.generate_images import GenerateImagesFunction, generate_images_core
|
||||
from igny8_core.ai.functions.generate_image_prompts import GenerateImagePromptsFunction
|
||||
from igny8_core.ai.functions.generate_images_from_prompts import GenerateImagesFromPromptsFunction
|
||||
|
||||
__all__ = [
|
||||
'AutoClusterFunction',
|
||||
@@ -14,4 +15,5 @@ __all__ = [
|
||||
'GenerateImagesFunction',
|
||||
'generate_images_core',
|
||||
'GenerateImagePromptsFunction',
|
||||
'GenerateImagesFromPromptsFunction',
|
||||
]
|
||||
|
||||
311
backend/igny8_core/ai/functions/generate_images_from_prompts.py
Normal file
311
backend/igny8_core/ai/functions/generate_images_from_prompts.py
Normal file
@@ -0,0 +1,311 @@
|
||||
"""
|
||||
Generate Images from Prompts AI Function
|
||||
Generates actual images from existing image prompts using AI
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, List, Any
|
||||
from django.db import transaction
|
||||
from igny8_core.ai.base import BaseAIFunction
|
||||
from igny8_core.modules.writer.models import Images, Content
|
||||
from igny8_core.ai.ai_core import AICore
|
||||
from igny8_core.ai.validators import validate_ids
|
||||
from igny8_core.ai.prompts import PromptRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
||||
"""Generate actual images from image prompts using AI"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return 'generate_images_from_prompts'
|
||||
|
||||
def get_metadata(self) -> Dict:
|
||||
return {
|
||||
'display_name': 'Generate Images from Prompts',
|
||||
'description': 'Generate actual images from existing image prompts',
|
||||
'phases': {
|
||||
'INIT': 'Validating image prompts...',
|
||||
'PREP': 'Preparing image generation queue...',
|
||||
'AI_CALL': 'Generating images with AI...',
|
||||
'PARSE': 'Processing image URLs...',
|
||||
'SAVE': 'Saving image URLs...',
|
||||
'DONE': 'Images generated!'
|
||||
}
|
||||
}
|
||||
|
||||
def get_max_items(self) -> int:
|
||||
return 100 # Max images per batch
|
||||
|
||||
def validate(self, payload: dict, account=None) -> Dict:
|
||||
"""Validate image IDs exist and have prompts"""
|
||||
result = validate_ids(payload, max_items=self.get_max_items())
|
||||
if not result['valid']:
|
||||
return result
|
||||
|
||||
# Check images exist and have prompts
|
||||
image_ids = payload.get('ids', [])
|
||||
if image_ids:
|
||||
queryset = Images.objects.filter(id__in=image_ids)
|
||||
if account:
|
||||
queryset = queryset.filter(account=account)
|
||||
|
||||
images = list(queryset.select_related('content', 'task'))
|
||||
|
||||
if not images:
|
||||
return {
|
||||
'valid': False,
|
||||
'error': 'No images found with provided IDs'
|
||||
}
|
||||
|
||||
# Check all images have prompts
|
||||
images_without_prompts = [img.id for img in images if not img.prompt or not img.prompt.strip()]
|
||||
if images_without_prompts:
|
||||
return {
|
||||
'valid': False,
|
||||
'error': f'Images {images_without_prompts} do not have prompts'
|
||||
}
|
||||
|
||||
# Check all images are pending
|
||||
images_not_pending = [img.id for img in images if img.status != 'pending']
|
||||
if images_not_pending:
|
||||
return {
|
||||
'valid': False,
|
||||
'error': f'Images {images_not_pending} are not in pending status'
|
||||
}
|
||||
|
||||
return {'valid': True}
|
||||
|
||||
def prepare(self, payload: dict, account=None) -> Dict:
|
||||
"""Load images and image generation settings"""
|
||||
image_ids = payload.get('ids', [])
|
||||
|
||||
queryset = Images.objects.filter(id__in=image_ids, status='pending')
|
||||
if account:
|
||||
queryset = queryset.filter(account=account)
|
||||
|
||||
images = list(queryset.select_related('content', 'task', 'account', 'site', 'sector'))
|
||||
|
||||
if not images:
|
||||
raise ValueError("No pending images found with prompts")
|
||||
|
||||
# Get image generation settings
|
||||
image_settings = {}
|
||||
if account:
|
||||
try:
|
||||
from igny8_core.modules.system.models import IntegrationSettings
|
||||
integration = IntegrationSettings.objects.get(
|
||||
account=account,
|
||||
integration_type='image_generation',
|
||||
is_active=True
|
||||
)
|
||||
image_settings = integration.config or {}
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load image generation settings: {e}")
|
||||
|
||||
# Extract settings with defaults
|
||||
provider = image_settings.get('provider') or image_settings.get('service', 'openai')
|
||||
if provider == 'runware':
|
||||
model = image_settings.get('model') or image_settings.get('runwareModel', 'runware:97@1')
|
||||
else:
|
||||
model = image_settings.get('model', 'dall-e-3')
|
||||
|
||||
# Get prompt templates
|
||||
image_prompt_template = PromptRegistry.get_image_prompt_template(account)
|
||||
negative_prompt = PromptRegistry.get_negative_prompt(account)
|
||||
|
||||
return {
|
||||
'images': images,
|
||||
'account': account,
|
||||
'provider': provider,
|
||||
'model': model,
|
||||
'image_type': image_settings.get('image_type', 'realistic'),
|
||||
'image_format': image_settings.get('image_format', 'webp'),
|
||||
'image_prompt_template': image_prompt_template,
|
||||
'negative_prompt': negative_prompt,
|
||||
}
|
||||
|
||||
def build_prompt(self, data: Dict, account=None) -> str:
|
||||
"""
|
||||
Build prompt for AI_CALL phase.
|
||||
For image generation, we return a placeholder since we process images in save_output.
|
||||
"""
|
||||
# Return placeholder - actual processing happens in save_output
|
||||
return "Image generation queue prepared"
|
||||
|
||||
def parse_response(self, response: str, step_tracker=None) -> Dict:
|
||||
"""
|
||||
Parse response from AI_CALL.
|
||||
For image generation, we process images directly in save_output, so this is a placeholder.
|
||||
"""
|
||||
return {'processed': True}
|
||||
|
||||
def save_output(
|
||||
self,
|
||||
parsed: Dict,
|
||||
original_data: Dict,
|
||||
account=None,
|
||||
progress_tracker=None,
|
||||
step_tracker=None
|
||||
) -> Dict:
|
||||
"""
|
||||
Process all images sequentially and generate them.
|
||||
This method handles the loop and makes AI calls directly.
|
||||
"""
|
||||
images = original_data.get('images', [])
|
||||
if not images:
|
||||
raise ValueError("No images to process")
|
||||
|
||||
provider = original_data.get('provider', 'openai')
|
||||
model = original_data.get('model', 'dall-e-3')
|
||||
image_type = original_data.get('image_type', 'realistic')
|
||||
image_prompt_template = original_data.get('image_prompt_template', '')
|
||||
negative_prompt = original_data.get('negative_prompt', '')
|
||||
|
||||
ai_core = AICore(account=account or original_data.get('account'))
|
||||
|
||||
total_images = len(images)
|
||||
images_generated = 0
|
||||
images_failed = 0
|
||||
errors = []
|
||||
|
||||
# Process each image sequentially
|
||||
for index, image in enumerate(images, 1):
|
||||
try:
|
||||
# Get content title
|
||||
content = image.content
|
||||
if not content:
|
||||
# Fallback to task if no content
|
||||
if image.task:
|
||||
content_title = image.task.title
|
||||
else:
|
||||
content_title = "Content"
|
||||
else:
|
||||
content_title = content.title or content.meta_title or "Content"
|
||||
|
||||
# Format prompt using template
|
||||
if image_prompt_template:
|
||||
try:
|
||||
formatted_prompt = image_prompt_template.format(
|
||||
post_title=content_title,
|
||||
image_prompt=image.prompt,
|
||||
image_type=image_type
|
||||
)
|
||||
except KeyError as e:
|
||||
logger.warning(f"Template formatting error: {e}, using simple format")
|
||||
formatted_prompt = f"Create a high-quality {image_type} image: {image.prompt}"
|
||||
else:
|
||||
# Fallback template
|
||||
formatted_prompt = f"Create a high-quality {image_type} image: {image.prompt}"
|
||||
|
||||
# Update progress: PREP phase for this image
|
||||
if progress_tracker and step_tracker:
|
||||
prep_msg = f"Generating image {index} of {total_images}: {image.image_type}"
|
||||
step_tracker.add_request_step("PREP", "success", prep_msg)
|
||||
progress_pct = 10 + int((index - 1) / total_images * 15) # 10-25% for PREP
|
||||
progress_tracker.update("PREP", progress_pct, prep_msg, meta=step_tracker.get_meta())
|
||||
|
||||
# Generate image
|
||||
if progress_tracker and step_tracker:
|
||||
ai_msg = f"Generating {image.image_type} image {index} of {total_images} with AI"
|
||||
step_tracker.add_response_step("AI_CALL", "success", ai_msg)
|
||||
progress_pct = 25 + int((index - 1) / total_images * 45) # 25-70% for AI_CALL
|
||||
progress_tracker.update("AI_CALL", progress_pct, ai_msg, meta=step_tracker.get_meta())
|
||||
|
||||
result = ai_core.generate_image(
|
||||
prompt=formatted_prompt,
|
||||
provider=provider,
|
||||
model=model,
|
||||
size='1024x1024',
|
||||
negative_prompt=negative_prompt if provider == 'runware' else None,
|
||||
function_name='generate_images_from_prompts'
|
||||
)
|
||||
|
||||
if result.get('error'):
|
||||
# Mark as failed
|
||||
with transaction.atomic():
|
||||
image.status = 'failed'
|
||||
image.save(update_fields=['status', 'updated_at'])
|
||||
|
||||
error_msg = f"Image {index} failed: {result['error']}"
|
||||
errors.append(error_msg)
|
||||
images_failed += 1
|
||||
logger.error(f"Image generation failed for image {image.id}: {result['error']}")
|
||||
|
||||
if progress_tracker and step_tracker:
|
||||
parse_msg = f"Image {index} failed: {result['error']}"
|
||||
step_tracker.add_response_step("PARSE", "error", parse_msg)
|
||||
progress_pct = 70 + int((index - 1) / total_images * 15) # 70-85% for PARSE
|
||||
progress_tracker.update("PARSE", progress_pct, parse_msg, meta=step_tracker.get_meta())
|
||||
|
||||
continue
|
||||
|
||||
image_url = result.get('url')
|
||||
if not image_url:
|
||||
# Mark as failed
|
||||
with transaction.atomic():
|
||||
image.status = 'failed'
|
||||
image.save(update_fields=['status', 'updated_at'])
|
||||
|
||||
error_msg = f"Image {index} failed: No URL returned"
|
||||
errors.append(error_msg)
|
||||
images_failed += 1
|
||||
logger.error(f"No image URL returned for image {image.id}")
|
||||
|
||||
if progress_tracker and step_tracker:
|
||||
parse_msg = f"Image {index} failed: No URL returned"
|
||||
step_tracker.add_response_step("PARSE", "error", parse_msg)
|
||||
progress_pct = 70 + int((index - 1) / total_images * 15)
|
||||
progress_tracker.update("PARSE", progress_pct, parse_msg, meta=step_tracker.get_meta())
|
||||
|
||||
continue
|
||||
|
||||
# Update progress: PARSE phase
|
||||
if progress_tracker and step_tracker:
|
||||
parse_msg = f"Image {index} of {total_images} generated successfully"
|
||||
step_tracker.add_response_step("PARSE", "success", parse_msg)
|
||||
progress_pct = 70 + int((index - 1) / total_images * 15) # 70-85% for PARSE
|
||||
progress_tracker.update("PARSE", progress_pct, parse_msg, meta=step_tracker.get_meta())
|
||||
|
||||
# Update image record
|
||||
with transaction.atomic():
|
||||
image.image_url = image_url
|
||||
image.status = 'generated'
|
||||
image.save(update_fields=['image_url', 'status', 'updated_at'])
|
||||
|
||||
images_generated += 1
|
||||
logger.info(f"Image {image.id} ({image.image_type}) generated successfully: {image_url}")
|
||||
|
||||
# Update progress: SAVE phase
|
||||
if progress_tracker and step_tracker:
|
||||
save_msg = f"Saved image {index} of {total_images}"
|
||||
step_tracker.add_request_step("SAVE", "success", save_msg)
|
||||
progress_pct = 85 + int((index - 1) / total_images * 13) # 85-98% for SAVE
|
||||
progress_tracker.update("SAVE", progress_pct, save_msg, meta=step_tracker.get_meta())
|
||||
|
||||
except Exception as e:
|
||||
# Mark as failed
|
||||
with transaction.atomic():
|
||||
image.status = 'failed'
|
||||
image.save(update_fields=['status', 'updated_at'])
|
||||
|
||||
error_msg = f"Image {index} failed: {str(e)}"
|
||||
errors.append(error_msg)
|
||||
images_failed += 1
|
||||
logger.error(f"Exception generating image {image.id}: {str(e)}", exc_info=True)
|
||||
continue
|
||||
|
||||
# Final progress update
|
||||
if progress_tracker and step_tracker:
|
||||
final_msg = f"Generated {images_generated} of {total_images} images"
|
||||
step_tracker.add_request_step("SAVE", "success", final_msg)
|
||||
progress_tracker.update("SAVE", 98, final_msg, meta=step_tracker.get_meta())
|
||||
|
||||
return {
|
||||
'count': images_generated,
|
||||
'images_generated': images_generated,
|
||||
'images_failed': images_failed,
|
||||
'total_images': total_images,
|
||||
'errors': errors if errors else None
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user