feat: initial commit of AI Comic Generator
This commit is contained in:
@@ -0,0 +1,12 @@
|
||||
import os
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
class Settings(BaseSettings):
|
||||
PROJECT_NAME: str = "AI Comic Generator"
|
||||
API_V1_STR: str = "/api/v1"
|
||||
DATABASE_URL: str = "sqlite:///./comic_app.db"
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
|
||||
settings = Settings()
|
||||
@@ -0,0 +1,15 @@
|
||||
from sqlmodel import SQLModel, create_engine, Session
|
||||
from app.core.config import settings
|
||||
|
||||
connect_args = {}
|
||||
if "sqlite" in settings.DATABASE_URL:
|
||||
connect_args["check_same_thread"] = False
|
||||
|
||||
engine = create_engine(settings.DATABASE_URL, echo=False, connect_args=connect_args)
|
||||
|
||||
def init_db():
|
||||
SQLModel.metadata.create_all(engine)
|
||||
|
||||
def get_session():
|
||||
with Session(engine) as session:
|
||||
yield session
|
||||
@@ -0,0 +1,102 @@
|
||||
# This file is used to centrally manage the core system prompt (System Prompt) of the comic generation system
|
||||
# Provided to the backend service in the form of Python variables to avoid path or encoding issues caused by directly reading text files
|
||||
|
||||
COMIC_GENERATION_SYSTEM_PROMPT = """As a 'Comic Split Generation' expert, you will play the dual role of a comic editing expert and a comic drawing expert. Your goal is to assist users in the entire process from story creation to storyboard design, and finally generate a serialized comic.
|
||||
|
||||
Purpose and Goals:
|
||||
* Expand or optimize the user-provided content into a story with a complete plot, ensuring logical self-consistency and engagement.
|
||||
* Clarify the overall comic style. You must strictly use the user-specified style (if any) and are prohibited from making decisions on your own.
|
||||
* Extract characters from the story and generate character setting cards containing core information about appearance, clothing, and personality.
|
||||
* Split the story into serialized comic storyboards, ensuring natural transitions between panels and appropriate narrative pacing.
|
||||
* Provide extremely detailed visual descriptions for each panel, precise to lighting, composition, character expression, and key actions.
|
||||
* Provide storyboard content in JSON file format, strictly stipulating that every four panels constitute an independent JSON object block for structured processing.
|
||||
|
||||
Behaviors and Rules:
|
||||
|
||||
1) Story Optimization & Character Setting (Story & Character):
|
||||
a) Receive the user's initial idea, enrich its background details, emotional ups and downs, and climax ending.
|
||||
b) Confirm Art Style: If the user provides specific "Theme" or "Style" requirements, you must follow them unconditionally. Do not modify the style based on the story content. For example, if the user requests "Cyberpunk", even if the story is set in a martial arts background, you must generate "Cyberpunk style martial arts".
|
||||
c) Before starting the storyboard, list the main characters and their physical characteristics (such as hair color, eye color, signature accessories, etc.) in detail to ensure visual consistency of characters in subsequent storyboards.
|
||||
|
||||
2) Storyboard Splitting & JSON Construction (Storyboard & JSON):
|
||||
a) Comic Global Configuration (Comic Configuration):
|
||||
- **Must be generated first**: Before starting to generate characters and storyboards, you must generate an independent 'comic_config' JSON block.
|
||||
- **Function**: Define the visual tone, typography standards, and border styles of the entire comic.
|
||||
- JSON structure should contain:
|
||||
- 'type': 'comic_config'
|
||||
- 'language': 'English',
|
||||
- 'style': '{User Specified Style}' (Must fill in the user-specified style, if not specified, default to 'Chibi/Fantasy Style')
|
||||
- 'bubble_style': { 'shape': 'Bubble Shape', 'color': 'Background Color', 'font_color': 'Font Color', 'stroke_width': 'Stroke Width' }
|
||||
- 'narration_style': { 'shape': 'Box/Rounded', 'color': 'Background Color', 'font_color': 'Font Color', 'opacity': 'Opacity' }
|
||||
- 'border_style': { 'width': 'Line Width', 'color': 'Color', 'type': 'Solid/Hand-drawn' }
|
||||
- 'gutter_style': { 'type': 'Standard Cross Split', 'color': 'White', 'width': '10px' }
|
||||
- 'layout_settings': {
|
||||
'show_panel_numbers': false, // [Switch] Whether to show panel numbers
|
||||
'panel_number_style': { 'position': 'top-left', 'bg_color': '#000000', 'text_color': '#00FF41', 'font_size': '14px' },
|
||||
'force_uniform_borders': true,
|
||||
'composition_mode': 'grid' // grid=Grid Splicing, cinematic=Cinematic Widescreen
|
||||
}
|
||||
- 'aspect_ratio': '16:9'
|
||||
|
||||
b) Character Sheet Generation (Character Sheets):
|
||||
- **All characters must have setting cards**: Before generating story storyboards, you must generate independent JSON setting blocks for all named characters appearing in the story, including protagonists, frequent supporting characters, and villains.
|
||||
- **Protagonist Design**: The protagonist's image must be designed to be extremely attractive, with distinct physical features, meeting "high aesthetic" standards.
|
||||
- **Character Deduplication**: When generating the character list, carefully identify different names for the same character (e.g., "Butler Ma" and "Old Ma" are the same person). If found to be the same person, generate only one character setting card and use the most formal or common name in the name field. Strictly prohibit generating multiple duplicate setting cards for the same character.
|
||||
- If there are multiple main characters, please generate multiple independent 'character_sheet' JSON blocks respectively, or output them in a JSON array.
|
||||
- JSON structure should contain:
|
||||
- 'type': 'character_sheet'
|
||||
- 'name': 'Character Name'
|
||||
- 'meta_info':{
|
||||
- 'language': Language based on user input
|
||||
- 'role': 'Protagonist' | 'Supporting' | 'Extra' (Must indicate character type)
|
||||
- 'personality': 'Character personality traits, e.g., Cheerful, Cold, Hot-blooded, etc., which will affect expressions and poses'
|
||||
- 'age': 'Age description, approximate range'
|
||||
- 'relationships': 'Description of relationship with protagonist or other characters' (Must indicate interpersonal relationships)
|
||||
- 'style': '{User Specified Style}',
|
||||
- 'feature': 'Explicit character features, e.g., Youthful, Plump, etc.'
|
||||
- 'aspect_ratio': '16:9'
|
||||
}
|
||||
- 'design_panels': [
|
||||
{'view': 'Front View', 'description': 'Detailed front full-body description...'},
|
||||
{'view': 'Side View', 'description': 'Detailed side view description...'},
|
||||
{'view': 'Clothing', 'description': 'Detailed clothing details...'},
|
||||
{'view': 'Accessories', 'description': 'Detailed accessories/weapon details...'}
|
||||
]
|
||||
|
||||
c) Story Storyboard Generation (Story Storyboard):
|
||||
- Decompose the optimized story into concrete, visualizable storyboard frames.
|
||||
- **Storyboard Count Mandatory Requirement**: The total number of generated panels must be determined based on the story content. The richer the story, the more panels, unless the user input requires a minimum number of panels (e.g., "at least 36 panels").
|
||||
- The total number of generated panels must be >= the minimum number required by the user.
|
||||
- The total number of generated panels must be an integer multiple of 4 (e.g., 36, 40, 44...), rounding up if not satisfied.
|
||||
|
||||
- Construct every four panels as an independent JSON code block output.
|
||||
- JSON structure should contain:
|
||||
- 'type': 'storyboard'
|
||||
- 'meta_info': {
|
||||
'style': '{User Specified Style}',
|
||||
'language': 'English',
|
||||
'volume': 'Current Volume/Total Volumes',
|
||||
'aspect_ratio': '16:9',
|
||||
}
|
||||
- 'characters': ['List of characters appearing in this group of panels']
|
||||
- 'plot_breakdown': [
|
||||
{'panel': 1, 'scene': '...', 'action': '...', 'dialogue': '...'},
|
||||
{'panel': 2, 'scene': '...', 'action': '...', 'dialogue': '...'},
|
||||
{'panel': 3, 'scene': '...', 'action': '...', 'dialogue': '...'},
|
||||
{'panel': 4, 'scene': '...', 'action': '...', 'dialogue': '...'}
|
||||
]
|
||||
|
||||
3) Quality Control (Quality Control):
|
||||
a) Ensure visual logic consistency between panels, avoiding sudden changes in characters or environment.
|
||||
b) Visual prompts should include elements such as environment, weather, shot type (e.g., close-up, panoramic), etc.
|
||||
|
||||
4) Language & Format Requirements (Language & Format):
|
||||
a) Use the same language for dialogue and narration as the user input. English is the standard.
|
||||
b) All JSON outputs must maintain a strict, parsable code block format.
|
||||
c) All JSON keys must use lowercase English (e.g., 'language', 'style').
|
||||
|
||||
|
||||
Overall Tone:
|
||||
* Professional and highly creative, demonstrating the rigor and aesthetics of a senior industry practitioner.
|
||||
* Descriptions of visual details should be precise and evocative.
|
||||
"""
|
||||
@@ -0,0 +1,35 @@
|
||||
from sqlmodel import Session, select
|
||||
from app.models.models import ModelConfig
|
||||
from app.schemas.schemas import ModelConfigCreate, ModelConfigUpdate
|
||||
from typing import List, Optional
|
||||
|
||||
def create_model_config(session: Session, config_in: ModelConfigCreate) -> ModelConfig:
|
||||
db_config = ModelConfig.model_validate(config_in)
|
||||
session.add(db_config)
|
||||
session.commit()
|
||||
session.refresh(db_config)
|
||||
return db_config
|
||||
|
||||
def get_model_configs(session: Session, skip: int = 0, limit: int = 100) -> List[ModelConfig]:
|
||||
statement = select(ModelConfig).offset(skip).limit(limit)
|
||||
return session.exec(statement).all()
|
||||
|
||||
def get_model_config(session: Session, config_id: int) -> Optional[ModelConfig]:
|
||||
return session.get(ModelConfig, config_id)
|
||||
|
||||
def update_model_config(session: Session, db_config: ModelConfig, config_in: ModelConfigUpdate) -> ModelConfig:
|
||||
config_data = config_in.model_dump(exclude_unset=True)
|
||||
for key, value in config_data.items():
|
||||
setattr(db_config, key, value)
|
||||
session.add(db_config)
|
||||
session.commit()
|
||||
session.refresh(db_config)
|
||||
return db_config
|
||||
|
||||
def delete_model_config(session: Session, db_config: ModelConfig):
|
||||
session.delete(db_config)
|
||||
session.commit()
|
||||
|
||||
def get_active_config(session: Session, model_type: str) -> Optional[ModelConfig]:
|
||||
statement = select(ModelConfig).where(ModelConfig.model_type == model_type, ModelConfig.is_active == True)
|
||||
return session.exec(statement).first()
|
||||
@@ -0,0 +1,99 @@
|
||||
from sqlmodel import Session, select
|
||||
from app.models.models import Project, Character, StoryboardItem, GlobalConfig
|
||||
from app.schemas.schemas import ProjectCreate, ProjectUpdate
|
||||
from typing import List, Optional
|
||||
|
||||
def create_project(session: Session, project_in: ProjectCreate) -> Project:
|
||||
db_project = Project.model_validate(project_in)
|
||||
session.add(db_project)
|
||||
session.commit()
|
||||
session.refresh(db_project)
|
||||
return db_project
|
||||
|
||||
def get_projects(session: Session, skip: int = 0, limit: int = 100) -> List[Project]:
|
||||
statement = select(Project).offset(skip).limit(limit).order_by(Project.updated_at.desc())
|
||||
return session.exec(statement).all()
|
||||
|
||||
def get_project(session: Session, project_id: str) -> Optional[Project]:
|
||||
return session.get(Project, project_id)
|
||||
|
||||
def update_project(session: Session, db_project: Project, project_in: ProjectUpdate) -> Project:
|
||||
project_data = project_in.model_dump(exclude_unset=True)
|
||||
for key, value in project_data.items():
|
||||
setattr(db_project, key, value)
|
||||
session.add(db_project)
|
||||
session.commit()
|
||||
session.refresh(db_project)
|
||||
return db_project
|
||||
|
||||
def delete_project(session: Session, db_project: Project):
|
||||
session.delete(db_project)
|
||||
session.commit()
|
||||
|
||||
# Helpers for sub-entities
|
||||
def create_global_config(session: Session, project_id: str, data: dict) -> GlobalConfig:
|
||||
# Check if exists
|
||||
statement = select(GlobalConfig).where(GlobalConfig.project_id == project_id)
|
||||
existing = session.exec(statement).first()
|
||||
if existing:
|
||||
existing.data = data
|
||||
session.add(existing)
|
||||
session.commit()
|
||||
session.refresh(existing)
|
||||
return existing
|
||||
|
||||
db_config = GlobalConfig(project_id=project_id, data=data)
|
||||
session.add(db_config)
|
||||
session.commit()
|
||||
session.refresh(db_config)
|
||||
return db_config
|
||||
|
||||
def save_characters(session: Session, project_id: str, characters_data: List[dict]) -> List[Character]:
|
||||
# Strategy: Clear existing or update?
|
||||
# For simplicity in this flow: Clear and Re-insert is easier if we regenerate all.
|
||||
# But user might want to edit specific ones.
|
||||
# Better: Update by name match, create if new.
|
||||
|
||||
results = []
|
||||
for char_data in characters_data:
|
||||
name = char_data.get("name")
|
||||
if not name: continue
|
||||
|
||||
statement = select(Character).where(Character.project_id == project_id, Character.name == name)
|
||||
existing = session.exec(statement).first()
|
||||
|
||||
if existing:
|
||||
existing.data = char_data
|
||||
session.add(existing)
|
||||
results.append(existing)
|
||||
else:
|
||||
new_char = Character(project_id=project_id, name=name, data=char_data)
|
||||
session.add(new_char)
|
||||
results.append(new_char)
|
||||
|
||||
session.commit()
|
||||
return results
|
||||
|
||||
def save_storyboard(session: Session, project_id: str, storyboard_data: List[dict]) -> List[StoryboardItem]:
|
||||
# Similar strategy: Clear and Re-insert is risky if we have images.
|
||||
# But storyboard is sequential.
|
||||
# Let's delete all and re-insert for now as "Regenerate JSON" usually means fresh start.
|
||||
# IF the user is just editing JSON text, we replace everything.
|
||||
|
||||
# Check if there are existing items with images we want to preserve?
|
||||
# Ideally, we should try to map them back, but it's hard if sequence changes.
|
||||
# For now: delete all items for this project and insert new.
|
||||
|
||||
statement = select(StoryboardItem).where(StoryboardItem.project_id == project_id)
|
||||
existing_items = session.exec(statement).all()
|
||||
for item in existing_items:
|
||||
session.delete(item)
|
||||
|
||||
results = []
|
||||
for i, item_data in enumerate(storyboard_data):
|
||||
new_item = StoryboardItem(project_id=project_id, sequence=i+1, data=item_data)
|
||||
session.add(new_item)
|
||||
results.append(new_item)
|
||||
|
||||
session.commit()
|
||||
return results
|
||||
@@ -0,0 +1,40 @@
|
||||
from fastapi import FastAPI
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from app.core.config import settings
|
||||
from app.core.database import init_db
|
||||
from app.routers import configs, projects, generation, export, tasks, history
|
||||
import os
|
||||
|
||||
app = FastAPI(title=settings.PROJECT_NAME)
|
||||
|
||||
# CORS
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Mount static files
|
||||
static_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "static")
|
||||
if not os.path.exists(static_dir):
|
||||
os.makedirs(static_dir)
|
||||
app.mount("/static", StaticFiles(directory=static_dir), name="static")
|
||||
|
||||
# Include routers
|
||||
app.include_router(configs.router, prefix=f"{settings.API_V1_STR}/configs", tags=["configs"])
|
||||
app.include_router(projects.router, prefix=f"{settings.API_V1_STR}/projects", tags=["projects"])
|
||||
app.include_router(generation.router, prefix=f"{settings.API_V1_STR}/generate", tags=["generation"])
|
||||
app.include_router(export.router, prefix=f"{settings.API_V1_STR}/export", tags=["export"])
|
||||
app.include_router(tasks.router, prefix=f"{settings.API_V1_STR}/tasks", tags=["tasks"])
|
||||
app.include_router(history.router, prefix=f"{settings.API_V1_STR}/history", tags=["history"])
|
||||
|
||||
@app.on_event("startup")
|
||||
def on_startup():
|
||||
init_db()
|
||||
|
||||
@app.get("/")
|
||||
def read_root():
|
||||
return {"message": "Welcome to AI Comic Generator API"}
|
||||
@@ -0,0 +1,97 @@
|
||||
from typing import Optional, List, Dict, Any
|
||||
from sqlmodel import SQLModel, Field, Relationship, Column, JSON
|
||||
from datetime import datetime
|
||||
import uuid
|
||||
|
||||
# --- Base Models ---
|
||||
|
||||
class ModelConfigBase(SQLModel):
|
||||
provider: str
|
||||
api_key: str
|
||||
base_url: Optional[str] = None
|
||||
model_name: str
|
||||
model_type: str
|
||||
is_active: bool = True
|
||||
|
||||
class ProjectBase(SQLModel):
|
||||
title: str
|
||||
description: Optional[str] = None
|
||||
story_input: Optional[str] = None
|
||||
# Generation Preferences
|
||||
theme: Optional[str] = None
|
||||
language: Optional[str] = "zh-CN"
|
||||
panel_count: Optional[int] = 16
|
||||
aspect_ratio: Optional[str] = "16:9"
|
||||
|
||||
class CharacterBase(SQLModel):
|
||||
name: str
|
||||
data: Dict = Field(default={}, sa_column=Column(JSON))
|
||||
image_url: Optional[str] = None
|
||||
|
||||
class StoryboardItemBase(SQLModel):
|
||||
sequence: int
|
||||
data: Dict = Field(default={}, sa_column=Column(JSON))
|
||||
image_url: Optional[str] = None
|
||||
|
||||
class GlobalConfigBase(SQLModel):
|
||||
data: Dict = Field(default={}, sa_column=Column(JSON))
|
||||
|
||||
class TaskBase(SQLModel):
|
||||
type: str # 'storyboard', 'image_generation', 'export'
|
||||
status: str # 'pending', 'processing', 'completed', 'failed'
|
||||
progress: int = 0 # 0-100
|
||||
message: Optional[str] = None
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
result: Dict = Field(default={}, sa_column=Column(JSON))
|
||||
|
||||
class ImageHistoryBase(SQLModel):
|
||||
entity_type: str # 'character' or 'storyboard_item'
|
||||
entity_id: int
|
||||
image_url: str
|
||||
|
||||
# --- Table Models ---
|
||||
|
||||
class ModelConfig(ModelConfigBase, table=True):
|
||||
id: Optional[int] = Field(default=None, primary_key=True)
|
||||
|
||||
class Project(ProjectBase, table=True):
|
||||
id: str = Field(default_factory=lambda: str(uuid.uuid4()), primary_key=True)
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
updated_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
|
||||
characters: List["Character"] = Relationship(back_populates="project", sa_relationship_kwargs={"cascade": "all, delete"})
|
||||
storyboard_items: List["StoryboardItem"] = Relationship(back_populates="project", sa_relationship_kwargs={"cascade": "all, delete"})
|
||||
global_config: Optional["GlobalConfig"] = Relationship(back_populates="project", sa_relationship_kwargs={"cascade": "all, delete"})
|
||||
tasks: List["Task"] = Relationship(back_populates="project", sa_relationship_kwargs={"cascade": "all, delete"})
|
||||
image_history: List["ImageHistory"] = Relationship(back_populates="project", sa_relationship_kwargs={"cascade": "all, delete"})
|
||||
|
||||
class Character(CharacterBase, table=True):
|
||||
id: Optional[int] = Field(default=None, primary_key=True)
|
||||
project_id: str = Field(foreign_key="project.id")
|
||||
project: Project = Relationship(back_populates="characters")
|
||||
|
||||
class StoryboardItem(StoryboardItemBase, table=True):
|
||||
id: Optional[int] = Field(default=None, primary_key=True)
|
||||
project_id: str = Field(foreign_key="project.id")
|
||||
project: Project = Relationship(back_populates="storyboard_items")
|
||||
|
||||
class GlobalConfig(GlobalConfigBase, table=True):
|
||||
id: Optional[int] = Field(default=None, primary_key=True)
|
||||
project_id: str = Field(foreign_key="project.id")
|
||||
project: Project = Relationship(back_populates="global_config")
|
||||
|
||||
class Task(TaskBase, table=True):
|
||||
id: str = Field(default_factory=lambda: str(uuid.uuid4()), primary_key=True)
|
||||
project_id: str = Field(foreign_key="project.id")
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
updated_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
|
||||
project: Project = Relationship(back_populates="tasks")
|
||||
|
||||
class ImageHistory(ImageHistoryBase, table=True):
|
||||
id: Optional[int] = Field(default=None, primary_key=True)
|
||||
project_id: str = Field(foreign_key="project.id")
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
|
||||
project: Project = Relationship(back_populates="image_history")
|
||||
@@ -0,0 +1,39 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from sqlmodel import Session
|
||||
from typing import List
|
||||
from app.core.database import get_session
|
||||
from app.models.models import ModelConfig
|
||||
from app.schemas.schemas import ModelConfigCreate, ModelConfigUpdate
|
||||
from app.cruds import crud_config
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@router.post("/", response_model=ModelConfig)
|
||||
def create_config(config_in: ModelConfigCreate, session: Session = Depends(get_session)):
|
||||
return crud_config.create_model_config(session, config_in)
|
||||
|
||||
@router.get("/", response_model=List[ModelConfig])
|
||||
def read_configs(skip: int = 0, limit: int = 100, session: Session = Depends(get_session)):
|
||||
return crud_config.get_model_configs(session, skip, limit)
|
||||
|
||||
@router.get("/{config_id}", response_model=ModelConfig)
|
||||
def read_config(config_id: int, session: Session = Depends(get_session)):
|
||||
config = crud_config.get_model_config(session, config_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Config not found")
|
||||
return config
|
||||
|
||||
@router.put("/{config_id}", response_model=ModelConfig)
|
||||
def update_config(config_id: int, config_in: ModelConfigUpdate, session: Session = Depends(get_session)):
|
||||
config = crud_config.get_model_config(session, config_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Config not found")
|
||||
return crud_config.update_model_config(session, config, config_in)
|
||||
|
||||
@router.delete("/{config_id}")
|
||||
def delete_config(config_id: int, session: Session = Depends(get_session)):
|
||||
config = crud_config.get_model_config(session, config_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Config not found")
|
||||
crud_config.delete_model_config(session, config)
|
||||
return {"ok": True}
|
||||
@@ -0,0 +1,80 @@
|
||||
import shutil
|
||||
import os
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from sqlmodel import Session
|
||||
from app.core.database import get_session
|
||||
from app.models.models import Project
|
||||
from app.services.image_service import split_comic_page
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@router.get("/{project_id}")
|
||||
def export_project(
|
||||
project_id: str,
|
||||
split_images: bool = False,
|
||||
session: Session = Depends(get_session)
|
||||
):
|
||||
project = session.get(Project, project_id)
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Check if any images generated
|
||||
has_images = any(item.image_url for item in project.storyboard_items) or any(c.image_url for c in project.characters)
|
||||
if not has_images:
|
||||
raise HTTPException(status_code=400, detail="No images generated yet. Cannot export.")
|
||||
|
||||
base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
project_static_dir = os.path.join(base_dir, "static", project_id)
|
||||
export_dir = os.path.join(project_static_dir, "export")
|
||||
|
||||
if os.path.exists(export_dir):
|
||||
shutil.rmtree(export_dir)
|
||||
os.makedirs(export_dir)
|
||||
|
||||
# Export Characters
|
||||
chars_dir = os.path.join(export_dir, "characters")
|
||||
os.makedirs(chars_dir)
|
||||
for char in project.characters:
|
||||
if char.image_url:
|
||||
# Resolve absolute path from relative URL
|
||||
# URL: /static/{project_id}/characters/xxx.png
|
||||
# Path: backend/static/{project_id}/characters/xxx.png
|
||||
# We can construct it directly if we know the structure, but let's parse url
|
||||
rel_path = char.image_url.lstrip("/") # static/...
|
||||
local_path = os.path.join(base_dir, rel_path.replace("/", os.sep))
|
||||
|
||||
if os.path.exists(local_path):
|
||||
shutil.copy(local_path, os.path.join(chars_dir, f"{char.name}.png"))
|
||||
|
||||
panels_dir = os.path.join(export_dir, "panels")
|
||||
if split_images:
|
||||
os.makedirs(panels_dir)
|
||||
|
||||
# Sort items
|
||||
items = sorted(project.storyboard_items, key=lambda x: x.sequence)
|
||||
|
||||
for item in items:
|
||||
if item.image_url:
|
||||
rel_path = item.image_url.lstrip("/")
|
||||
local_path = os.path.join(base_dir, rel_path.replace("/", os.sep))
|
||||
|
||||
if os.path.exists(local_path):
|
||||
shutil.copy(local_path, os.path.join(export_dir, f"comic_part_{item.sequence}.png"))
|
||||
|
||||
if split_images:
|
||||
with open(local_path, "rb") as f:
|
||||
img_bytes = f.read()
|
||||
|
||||
try:
|
||||
panels = split_comic_page(img_bytes)
|
||||
for idx, panel_bytes in enumerate(panels):
|
||||
p_name = f"panel_{item.sequence}_{idx+1}.png"
|
||||
with open(os.path.join(panels_dir, p_name), "wb") as f:
|
||||
f.write(panel_bytes)
|
||||
except Exception as e:
|
||||
print(f"Failed to split panel {item.id}: {e}")
|
||||
|
||||
zip_path_base = os.path.join(project_static_dir, "export_archive")
|
||||
shutil.make_archive(zip_path_base, 'zip', export_dir)
|
||||
|
||||
return {"download_url": f"/static/{project_id}/export_archive.zip"}
|
||||
@@ -0,0 +1,824 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks
|
||||
from sqlmodel import Session
|
||||
from app.core.database import get_session
|
||||
from app.models.models import Project, Character, StoryboardItem, Task, ImageHistory
|
||||
from app.services.ai_service import AIService
|
||||
from app.services.consistency_service import ConsistencyService
|
||||
from app.utils.json_utils import extract_json_blocks
|
||||
from app.cruds import crud_project
|
||||
import os
|
||||
import uuid
|
||||
import json
|
||||
import traceback
|
||||
|
||||
import logging
|
||||
import sys
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def save_generated_image(session, project_id, entity_type, entity_id, image_bytes):
|
||||
base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
project_static_dir = os.path.join(base_dir, "static", project_id)
|
||||
|
||||
sub_dir = "characters" if entity_type == "character" else "panels"
|
||||
target_dir = os.path.join(project_static_dir, sub_dir)
|
||||
|
||||
if not os.path.exists(target_dir):
|
||||
os.makedirs(target_dir)
|
||||
|
||||
filename = f"{entity_type}_{entity_id}_{uuid.uuid4().hex[:8]}.png"
|
||||
filepath = os.path.join(target_dir, filename)
|
||||
|
||||
with open(filepath, "wb") as f:
|
||||
f.write(image_bytes)
|
||||
|
||||
relative_url = f"/static/{project_id}/{sub_dir}/{filename}"
|
||||
|
||||
# Save History
|
||||
history = ImageHistory(
|
||||
project_id=project_id,
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
image_url=relative_url
|
||||
)
|
||||
session.add(history)
|
||||
|
||||
return relative_url
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
from app.core.prompts import COMIC_GENERATION_SYSTEM_PROMPT
|
||||
|
||||
def get_system_prompt():
|
||||
return COMIC_GENERATION_SYSTEM_PROMPT
|
||||
|
||||
# --- Background Task Functions ---
|
||||
|
||||
def generate_storyboard_task(task_id: str, project_id: str, user_input: str):
|
||||
logger.info(f"Starting storyboard generation task: {task_id} for project: {project_id}")
|
||||
# We need a fresh session for the background task
|
||||
from app.core.database import engine
|
||||
with Session(engine) as session:
|
||||
task = session.get(Task, task_id)
|
||||
if not task:
|
||||
logger.error(f"Task {task_id} not found")
|
||||
return
|
||||
|
||||
task.status = "processing"
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
try:
|
||||
project = crud_project.get_project(session, project_id)
|
||||
logger.info(f"Project found: {project.title}")
|
||||
|
||||
# Save User Input (Persist it)
|
||||
project.story_input = user_input
|
||||
session.add(project)
|
||||
session.commit()
|
||||
|
||||
ai = AIService(session)
|
||||
system_prompt = get_system_prompt()
|
||||
|
||||
# Construct Final Prompt with Preferences
|
||||
final_prompt = user_input
|
||||
|
||||
# --- Replace Placeholders in System Prompt ---
|
||||
system_prompt = get_system_prompt()
|
||||
|
||||
# Defaults
|
||||
style = "Standard"
|
||||
if project.theme: style = project.theme
|
||||
|
||||
lang = "English"
|
||||
if project.language:
|
||||
lang_map = {"zh-CN": "Simplified Chinese", "en-US": "English", "ja-JP": "Japanese"}
|
||||
lang = lang_map.get(project.language, project.language)
|
||||
|
||||
# Inject into System Prompt
|
||||
system_prompt = system_prompt.replace("{User Specified Style}", style)
|
||||
# We could also inject language if we had a placeholder, but style is the main one failing.
|
||||
# Let's add language instruction to system prompt dynamically if needed,
|
||||
# or rely on the "Language & Format" section in prompt which says "Use user input language".
|
||||
|
||||
# --- Construct User Prompt ---
|
||||
final_prompt = user_input
|
||||
|
||||
# Append preferences as normal requirements
|
||||
prefs = []
|
||||
if project.theme: prefs.append(f"Theme: {project.theme}")
|
||||
if project.language: prefs.append(f"Language: {project.language}")
|
||||
if project.panel_count: prefs.append(f"Estimated Panel Count: {project.panel_count}")
|
||||
if project.aspect_ratio: prefs.append(f"Aspect Ratio: {project.aspect_ratio}")
|
||||
|
||||
if prefs:
|
||||
final_prompt += "\n\nRequirements:\n" + "\n".join(prefs)
|
||||
|
||||
logger.info("Calling AI service for storyboard generation...")
|
||||
generated_text = ai.generate_storyboard(system_prompt, final_prompt)
|
||||
|
||||
# --- Save Generated Text to Temp File ---
|
||||
import time
|
||||
base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
temp_dir = os.path.join(base_dir, "static", project_id, "temp")
|
||||
if not os.path.exists(temp_dir):
|
||||
os.makedirs(temp_dir)
|
||||
|
||||
timestamp = int(time.time())
|
||||
temp_file = os.path.join(temp_dir, f"ai_output_{timestamp}.txt")
|
||||
try:
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
f.write(generated_text)
|
||||
logger.info(f"Saved raw AI output to {temp_file}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save temp AI output: {e}")
|
||||
# ----------------------------------------
|
||||
|
||||
logger.info("AI generation complete. Extracting JSON blocks...")
|
||||
json_blocks = extract_json_blocks(generated_text)
|
||||
|
||||
char_blocks = [b for b in json_blocks if b.get("type") == "character_sheet"]
|
||||
story_blocks = [b for b in json_blocks if b.get("type") == "storyboard"]
|
||||
|
||||
if not story_blocks:
|
||||
story_blocks = [b for b in json_blocks if b.get("type") not in ["character_sheet", "comic_config"]]
|
||||
|
||||
# --- Missing Character Check & Fix ---
|
||||
story_char_names = set()
|
||||
for block in story_blocks:
|
||||
chars = block.get("characters", [])
|
||||
if isinstance(chars, str):
|
||||
story_char_names.add(chars)
|
||||
elif isinstance(chars, list):
|
||||
for c in chars:
|
||||
if isinstance(c, str): story_char_names.add(c)
|
||||
elif isinstance(c, dict): story_char_names.add(c.get("name", ""))
|
||||
|
||||
generated_char_names = set(b.get("name") for b in char_blocks if b.get("name"))
|
||||
|
||||
# Simple fuzzy matching or direct check
|
||||
missing_chars = []
|
||||
for name in story_char_names:
|
||||
# Check if name is contained in any generated char name (e.g. "Xiao Ming" vs "Ming")
|
||||
found = False
|
||||
for g_name in generated_char_names:
|
||||
if name in g_name or g_name in name:
|
||||
found = True
|
||||
break
|
||||
if not found and name and len(name) > 1: # Ignore single chars or empty
|
||||
missing_chars.append(name)
|
||||
|
||||
if missing_chars:
|
||||
logger.info(f"Detected missing characters: {missing_chars}. Requesting AI to generate them...")
|
||||
fix_prompt = f"You missed generating character sheets for the following characters that appeared in the storyboard: {', '.join(missing_chars)}. Please generate 'character_sheet' JSON blocks for them now. Do not generate anything else."
|
||||
|
||||
try:
|
||||
fix_response = ai.generate_storyboard(system_prompt, fix_prompt) # Re-use generate method
|
||||
fix_blocks = extract_json_blocks(fix_response)
|
||||
new_chars = [b for b in fix_blocks if b.get("type") == "character_sheet"]
|
||||
if new_chars:
|
||||
logger.info(f"Successfully generated {len(new_chars)} missing characters.")
|
||||
char_blocks.extend(new_chars)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to generate missing characters: {e}")
|
||||
|
||||
config_block = next((b for b in json_blocks if b.get("type") == "comic_config"), None)
|
||||
|
||||
# If AI didn't return config, create one from project prefs
|
||||
if not config_block and (project.aspect_ratio or project.language):
|
||||
config_block = {
|
||||
"type": "comic_config",
|
||||
"style": "Standard", # Default
|
||||
"aspect_ratio": project.aspect_ratio or "16:9",
|
||||
"language": project.language or "en-US"
|
||||
}
|
||||
|
||||
if config_block:
|
||||
crud_project.create_global_config(session, project_id, config_block)
|
||||
|
||||
# --- Enforce Consistency: Update meta_info for all blocks ---
|
||||
# Re-read global config if we just created/updated it
|
||||
# Or use the config_block we have
|
||||
|
||||
if not config_block:
|
||||
# Try to fetch existing
|
||||
# But we just generated it. If None, we create a default one above.
|
||||
# Let's use the one we have.
|
||||
pass
|
||||
|
||||
if config_block:
|
||||
global_style = config_block.get("style", "")
|
||||
global_aspect = config_block.get("aspect_ratio", "16:9")
|
||||
global_lang = config_block.get("language", "en-US")
|
||||
|
||||
# Update Character Sheets
|
||||
for char in char_blocks:
|
||||
char["meta_info"] = char.get("meta_info", {})
|
||||
char["meta_info"]["language"] = global_lang
|
||||
char["meta_info"]["style"] = global_style
|
||||
|
||||
# Remove top-level redundant keys if they exist to avoid confusion
|
||||
char.pop("language", None)
|
||||
char.pop("style", None)
|
||||
|
||||
# Update Storyboard Items
|
||||
for block in story_blocks:
|
||||
meta = block.get("meta_info", {})
|
||||
meta["style"] = global_style
|
||||
meta["language"] = global_lang
|
||||
meta["aspect_ratio"] = global_aspect
|
||||
|
||||
# Also inject specific style configs if present
|
||||
if "bubble_style" in config_block: meta["bubble_style"] = config_block["bubble_style"]
|
||||
if "narration_style" in config_block: meta["narration_style"] = config_block["narration_style"]
|
||||
if "border_style" in config_block: meta["border_style"] = config_block["border_style"]
|
||||
if "gutter_style" in config_block: meta["gutter_style"] = config_block["gutter_style"]
|
||||
if "layout_settings" in config_block: meta["layout_settings"] = config_block["layout_settings"]
|
||||
|
||||
block["meta_info"] = meta
|
||||
|
||||
# Save to DB
|
||||
crud_project.save_characters(session, project_id, char_blocks)
|
||||
crud_project.save_storyboard(session, project_id, story_blocks)
|
||||
|
||||
# Consistency
|
||||
consistency = ConsistencyService(session)
|
||||
consistency.normalize_project(project_id)
|
||||
|
||||
task.status = "completed"
|
||||
task.result = {"blocks_found": len(json_blocks)}
|
||||
task.progress = 100
|
||||
session.add(task)
|
||||
session.commit()
|
||||
logger.info(f"Storyboard task {task_id} completed successfully.")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Storyboard task {task_id} failed: {e}")
|
||||
traceback.print_exc()
|
||||
task.status = "failed"
|
||||
task.message = str(e)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
def generate_all_images_task(task_id: str, project_id: str):
|
||||
logger.info(f"Starting batch image generation task: {task_id} for project: {project_id}")
|
||||
from app.core.database import engine
|
||||
with Session(engine) as session:
|
||||
task = session.get(Task, task_id)
|
||||
if not task:
|
||||
logger.error(f"Task {task_id} not found")
|
||||
return
|
||||
|
||||
task.status = "processing"
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
try:
|
||||
project = session.get(Project, project_id)
|
||||
ai = AIService(session)
|
||||
|
||||
base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
static_root = os.path.join(base_dir, "static")
|
||||
if not os.path.exists(static_root): os.makedirs(static_root)
|
||||
|
||||
# 1. Generate Characters
|
||||
total_chars = len(project.characters)
|
||||
logger.info(f"Generating {total_chars} characters...")
|
||||
for i, char in enumerate(project.characters):
|
||||
if char.image_url:
|
||||
logger.info(f"Character {char.name} already has image, skipping.")
|
||||
continue
|
||||
|
||||
logger.info(f"Generating image for character: {char.name}")
|
||||
json_prompt = json.dumps(char.data, ensure_ascii=False, indent=2)
|
||||
json_prompt += "\n\n generate a character design sheet with 4 panels: front view, side view, clothing details, accessories."
|
||||
|
||||
try:
|
||||
image_bytes = ai.generate_image(json_prompt)
|
||||
relative_url = save_generated_image(session, project_id, "character", char.id, image_bytes)
|
||||
char.image_url = relative_url
|
||||
session.add(char)
|
||||
session.commit()
|
||||
logger.info(f"Character {char.name} generated successfully.")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to generate char {char.id}: {e}")
|
||||
print(f"Failed to generate char {char.id}: {e}")
|
||||
|
||||
# Update task progress (Characters are 20% of work?)
|
||||
# Let's simple split: chars + storyboard items
|
||||
|
||||
# 2. Generate Storyboard Items (Sequential)
|
||||
# Re-fetch items to ensure order
|
||||
items = sorted(project.storyboard_items, key=lambda x: x.sequence)
|
||||
total_items = len(items)
|
||||
logger.info(f"Generating {total_items} storyboard panels...")
|
||||
|
||||
generated_history = [] # Keep track of generated images for context
|
||||
|
||||
# Populate history with existing images
|
||||
# Actually comic_generator logic builds history as it goes.
|
||||
# We should probably load existing images into history if we are resuming?
|
||||
# For "one click", let's assume we scan all items.
|
||||
|
||||
for i, item in enumerate(items):
|
||||
task.progress = int((i / total_items) * 100)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
if item.image_url:
|
||||
# Add to history
|
||||
filename = os.path.basename(item.image_url)
|
||||
# We need to find where it is stored.
|
||||
# Assuming standard structure
|
||||
# We need absolute path for history
|
||||
# item.image_url is like /static/{project_id}/panels/{filename}
|
||||
rel_path = item.image_url.lstrip("/")
|
||||
abs_path = os.path.join(base_dir, rel_path.replace("/", os.sep))
|
||||
|
||||
if os.path.exists(abs_path):
|
||||
generated_history.append(abs_path)
|
||||
continue
|
||||
|
||||
logger.info(f"Generating panel {item.sequence}...")
|
||||
|
||||
# Prepare Context
|
||||
context_images = []
|
||||
|
||||
# a) Character Sheets
|
||||
char_names = item.data.get("characters", [])
|
||||
if isinstance(char_names, str): char_names = [char_names]
|
||||
elif isinstance(char_names, list):
|
||||
names = []
|
||||
for c in char_names:
|
||||
if isinstance(c, dict): names.append(c.get("name", ""))
|
||||
elif isinstance(c, str): names.append(c)
|
||||
char_names = names
|
||||
|
||||
for name in char_names:
|
||||
for p_char in project.characters:
|
||||
if p_char.image_url and (p_char.name in name or name in p_char.name):
|
||||
# Resolve absolute path for char image
|
||||
rel_path = p_char.image_url.lstrip("/")
|
||||
abs_path = os.path.join(base_dir, rel_path.replace("/", os.sep))
|
||||
if os.path.exists(abs_path) and abs_path not in context_images:
|
||||
context_images.append(abs_path)
|
||||
|
||||
# b) Previous History (Last 3 logic)
|
||||
if len(generated_history) >= 3:
|
||||
selected = [generated_history[0]] + generated_history[-2:]
|
||||
else:
|
||||
selected = generated_history
|
||||
|
||||
for path in selected:
|
||||
if path not in context_images:
|
||||
context_images.append(path)
|
||||
|
||||
# Generate
|
||||
json_prompt = json.dumps(item.data, ensure_ascii=False, indent=2)
|
||||
json_prompt += "\n\n use json block as user input prompt to generate 2*2 grid comic image."
|
||||
|
||||
try:
|
||||
image_bytes = ai.generate_image(json_prompt, context_images)
|
||||
relative_url = save_generated_image(session, project_id, "panel", item.id, image_bytes)
|
||||
|
||||
item.image_url = relative_url
|
||||
session.add(item)
|
||||
session.commit()
|
||||
|
||||
# Add to history (absolute path for context usage)
|
||||
# We need absolute path for next context
|
||||
# save_generated_image returns relative /static/...
|
||||
# Reconstruct absolute path
|
||||
# strip leading /
|
||||
abs_path = os.path.join(base_dir, relative_url.lstrip("/").replace("/", os.sep))
|
||||
generated_history.append(abs_path)
|
||||
logger.info(f"Panel {item.sequence} generated successfully.")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to generate panel {item.id}: {e}")
|
||||
print(f"Failed to generate panel {item.id}: {e}")
|
||||
|
||||
task.status = "completed"
|
||||
task.progress = 100
|
||||
session.add(task)
|
||||
session.commit()
|
||||
logger.info(f"Batch generation task {task_id} completed successfully.")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Batch generation task {task_id} failed: {e}")
|
||||
traceback.print_exc()
|
||||
task.status = "failed"
|
||||
task.message = str(e)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
def generate_all_characters_task(task_id: str, project_id: str):
|
||||
logger.info(f"Starting batch character generation task: {task_id} for project: {project_id}")
|
||||
from app.core.database import engine
|
||||
with Session(engine) as session:
|
||||
task = session.get(Task, task_id)
|
||||
if not task:
|
||||
logger.error(f"Task {task_id} not found")
|
||||
return
|
||||
|
||||
task.status = "processing"
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
try:
|
||||
project = session.get(Project, project_id)
|
||||
ai = AIService(session)
|
||||
|
||||
total_chars = len(project.characters)
|
||||
logger.info(f"Generating {total_chars} characters...")
|
||||
|
||||
for i, char in enumerate(project.characters):
|
||||
if char.image_url:
|
||||
logger.info(f"Character {char.name} already has image, skipping.")
|
||||
continue
|
||||
|
||||
logger.info(f"Generating image for character: {char.name}")
|
||||
|
||||
# Construct Natural Language Prompt from JSON
|
||||
data = char.data
|
||||
meta = data.get("meta_info", {})
|
||||
name = data.get("name", "Unknown")
|
||||
role = meta.get("role", "")
|
||||
age = meta.get("age", "")
|
||||
personality = data.get("personality", "") or meta.get("personality", "")
|
||||
style = meta.get("style", "")
|
||||
|
||||
# Build Description from panels
|
||||
description = ""
|
||||
panels = data.get("design_panels", [])
|
||||
for p in panels:
|
||||
view = p.get("view", "")
|
||||
desc = p.get("description", "")
|
||||
description += f"- {view}: {desc}\n"
|
||||
|
||||
prompt = f"""Character Design Request:
|
||||
Name: {name}
|
||||
Role: {role}
|
||||
Age: {age}
|
||||
Personality: {personality}
|
||||
Style: {style}
|
||||
|
||||
Visual Description:
|
||||
{description}
|
||||
|
||||
Task: Generate a high-quality character reference sheet (Character Design) based on the above description.
|
||||
Include Front View, Side View, and detailed clothing/accessories.
|
||||
Ensure the character's expression and pose reflect their personality: {personality}.
|
||||
"""
|
||||
|
||||
try:
|
||||
image_bytes = ai.generate_image(prompt)
|
||||
relative_url = save_generated_image(session, project_id, "character", char.id, image_bytes)
|
||||
char.image_url = relative_url
|
||||
session.add(char)
|
||||
session.commit()
|
||||
logger.info(f"Character {char.name} generated successfully.")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to generate char {char.id}: {e}")
|
||||
|
||||
# Update progress
|
||||
progress = int(((i + 1) / total_chars) * 100)
|
||||
task.progress = progress
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
task.status = "completed"
|
||||
task.progress = 100
|
||||
session.add(task)
|
||||
session.commit()
|
||||
logger.info(f"Batch character generation task {task_id} completed successfully.")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Batch character generation task {task_id} failed: {e}")
|
||||
traceback.print_exc()
|
||||
task.status = "failed"
|
||||
task.message = str(e)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
def generate_character_task(task_id: str, character_id: int):
|
||||
logger.info(f"Starting character generation task: {task_id} for char: {character_id}")
|
||||
from app.core.database import engine
|
||||
with Session(engine) as session:
|
||||
task = session.get(Task, task_id)
|
||||
if not task:
|
||||
logger.error(f"Task {task_id} not found")
|
||||
return
|
||||
|
||||
task.status = "processing"
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
try:
|
||||
char = session.get(Character, character_id)
|
||||
if not char:
|
||||
raise ValueError("Character not found")
|
||||
|
||||
ai = AIService(session)
|
||||
|
||||
# Construct Natural Language Prompt from JSON
|
||||
data = char.data
|
||||
meta = data.get("meta_info", {})
|
||||
name = data.get("name", "Unknown")
|
||||
role = meta.get("role", "")
|
||||
age = meta.get("age", "")
|
||||
personality = data.get("personality", "") or meta.get("personality", "")
|
||||
style = meta.get("style", "")
|
||||
|
||||
# Build Description from panels
|
||||
description = ""
|
||||
panels = data.get("design_panels", [])
|
||||
for p in panels:
|
||||
view = p.get("view", "")
|
||||
desc = p.get("description", "")
|
||||
description += f"- {view}: {desc}\n"
|
||||
|
||||
prompt = f"""Character Design Request:
|
||||
Name: {name}
|
||||
Role: {role}
|
||||
Age: {age}
|
||||
Personality: {personality}
|
||||
Style: {style}
|
||||
|
||||
Visual Description:
|
||||
{description}
|
||||
|
||||
Task: Generate a high-quality character reference sheet (Character Design) based on the above description.
|
||||
Include Front View, Side View, and detailed clothing/accessories.
|
||||
Ensure the character's expression and pose reflect their personality: {personality}.
|
||||
"""
|
||||
|
||||
logger.info(f"Calling AI service for character {char.name}...")
|
||||
image_bytes = ai.generate_image(prompt)
|
||||
|
||||
relative_url = save_generated_image(session, char.project_id, "character", char.id, image_bytes)
|
||||
|
||||
char.image_url = relative_url
|
||||
session.add(char)
|
||||
|
||||
task.status = "completed"
|
||||
task.progress = 100
|
||||
session.add(task)
|
||||
session.commit()
|
||||
logger.info(f"Character task {task_id} completed successfully.")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Character task {task_id} failed: {e}")
|
||||
traceback.print_exc()
|
||||
task.status = "failed"
|
||||
task.message = str(e)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
# --- Endpoints ---
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
class StoryboardRequest(BaseModel):
|
||||
user_input: str
|
||||
|
||||
@router.post("/storyboard/{project_id}")
|
||||
def generate_storyboard(
|
||||
project_id: str,
|
||||
request: StoryboardRequest,
|
||||
background_tasks: BackgroundTasks,
|
||||
session: Session = Depends(get_session)
|
||||
):
|
||||
user_input = request.user_input
|
||||
logger.info(f"Received request to generate storyboard for project {project_id}")
|
||||
project = crud_project.get_project(session, project_id)
|
||||
if not project:
|
||||
logger.error(f"Project {project_id} not found")
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Save User Input Immediately
|
||||
project.story_input = user_input
|
||||
session.add(project)
|
||||
session.commit()
|
||||
|
||||
# Create Task
|
||||
task = Task(
|
||||
type="storyboard",
|
||||
status="pending",
|
||||
project_id=project_id,
|
||||
name="Generate Storyboard",
|
||||
description=f"Generating storyboard based on user input..."
|
||||
)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
session.refresh(task)
|
||||
logger.info(f"Task created: {task.id}")
|
||||
|
||||
background_tasks.add_task(generate_storyboard_task, task.id, project_id, user_input)
|
||||
|
||||
return {"task_id": task.id}
|
||||
|
||||
@router.post("/all-images/{project_id}")
|
||||
def generate_all_images(
|
||||
project_id: str,
|
||||
background_tasks: BackgroundTasks,
|
||||
session: Session = Depends(get_session)
|
||||
):
|
||||
logger.info(f"Received request to generate all images for project {project_id}")
|
||||
project = crud_project.get_project(session, project_id)
|
||||
if not project:
|
||||
logger.error(f"Project {project_id} not found")
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
task = Task(
|
||||
type="image_generation",
|
||||
status="pending",
|
||||
project_id=project_id,
|
||||
name="Batch Generate Images",
|
||||
description="Generating all storyboard images"
|
||||
)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
session.refresh(task)
|
||||
logger.info(f"Task created: {task.id}")
|
||||
|
||||
background_tasks.add_task(generate_all_images_task, task.id, project_id)
|
||||
|
||||
return {"task_id": task.id}
|
||||
|
||||
# Keep individual endpoints for manual control, but maybe make them async too?
|
||||
# User asked for "Back task" for "generation". Usually implies the bulk actions.
|
||||
# Single panel generation is usually fast enough (5-10s), but can be async if desired.
|
||||
# For now, let's keep single endpoints sync for immediate feedback, or make them async if user insists "All generation".
|
||||
# The prompt says "Generate text/image takes long time".
|
||||
# Let's keep single endpoints sync for simplicity of interaction (user waits 5s is ok),
|
||||
# but "One Click" and "Storyboard" are definitely async.
|
||||
|
||||
@router.post("/all-characters/{project_id}")
|
||||
def generate_all_characters(
|
||||
project_id: str,
|
||||
background_tasks: BackgroundTasks,
|
||||
session: Session = Depends(get_session)
|
||||
):
|
||||
logger.info(f"Received request to generate all characters for project {project_id}")
|
||||
project = crud_project.get_project(session, project_id)
|
||||
if not project:
|
||||
logger.error(f"Project {project_id} not found")
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
task = Task(
|
||||
type="character_generation",
|
||||
status="pending",
|
||||
project_id=project_id,
|
||||
name="Batch Generate Characters",
|
||||
description="Generating all character design sheets"
|
||||
)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
session.refresh(task)
|
||||
logger.info(f"Task created: {task.id}")
|
||||
|
||||
background_tasks.add_task(generate_all_characters_task, task.id, project_id)
|
||||
|
||||
return {"task_id": task.id}
|
||||
|
||||
@router.post("/character/{character_id}")
|
||||
def generate_character(
|
||||
character_id: int,
|
||||
background_tasks: BackgroundTasks,
|
||||
session: Session = Depends(get_session)
|
||||
):
|
||||
char = session.get(Character, character_id)
|
||||
if not char:
|
||||
raise HTTPException(status_code=404, detail="Character not found")
|
||||
|
||||
task = Task(
|
||||
type="character_generation",
|
||||
status="pending",
|
||||
project_id=char.project_id,
|
||||
name=f"Draw Character: {char.name}",
|
||||
description=f"Drawing design sheet for character {char.name}"
|
||||
)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
session.refresh(task)
|
||||
|
||||
background_tasks.add_task(generate_character_task, task.id, character_id)
|
||||
|
||||
return {"task_id": task.id}
|
||||
|
||||
def generate_panel_task(task_id: str, item_id: int):
|
||||
logger.info(f"Starting panel generation task: {task_id} for item: {item_id}")
|
||||
from app.core.database import engine
|
||||
with Session(engine) as session:
|
||||
task = session.get(Task, task_id)
|
||||
if not task:
|
||||
logger.error(f"Task {task_id} not found")
|
||||
return
|
||||
|
||||
task.status = "processing"
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
try:
|
||||
item = session.get(StoryboardItem, item_id)
|
||||
if not item:
|
||||
raise ValueError("Storyboard item not found")
|
||||
|
||||
project = item.project
|
||||
ai = AIService(session)
|
||||
json_prompt = json.dumps(item.data, ensure_ascii=False, indent=2)
|
||||
json_prompt += "\n\n use json block as user input prompt to generate 2*2 grid comic image."
|
||||
|
||||
context_images = []
|
||||
base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
# 1. Find Character Images
|
||||
char_names = item.data.get("characters", [])
|
||||
if isinstance(char_names, str): char_names = [char_names]
|
||||
elif isinstance(char_names, list):
|
||||
names = []
|
||||
for c in char_names:
|
||||
if isinstance(c, dict): names.append(c.get("name", ""))
|
||||
elif isinstance(c, str): names.append(c)
|
||||
char_names = names
|
||||
|
||||
if project.characters:
|
||||
for name in char_names:
|
||||
for p_char in project.characters:
|
||||
if p_char.image_url and (p_char.name in name or name in p_char.name):
|
||||
rel_path = p_char.image_url.lstrip("/")
|
||||
abs_path = os.path.join(base_dir, rel_path.replace("/", os.sep))
|
||||
if os.path.exists(abs_path) and abs_path not in context_images:
|
||||
context_images.append(abs_path)
|
||||
|
||||
# 2. Previous Panels
|
||||
prev_items = sorted([i for i in project.storyboard_items if i.sequence < item.sequence and i.image_url], key=lambda x: x.sequence)
|
||||
if prev_items:
|
||||
selected = []
|
||||
if len(prev_items) >= 3:
|
||||
selected = [prev_items[0]] + prev_items[-2:]
|
||||
else:
|
||||
selected = prev_items
|
||||
|
||||
for prev in selected:
|
||||
rel_path = prev.image_url.lstrip("/")
|
||||
abs_path = os.path.join(base_dir, rel_path.replace("/", os.sep))
|
||||
if os.path.exists(abs_path) and abs_path not in context_images:
|
||||
context_images.append(abs_path)
|
||||
|
||||
# Style Consistency
|
||||
meta_style = item.data.get("meta_info", {}).get("style", "")
|
||||
if not meta_style and project.global_config:
|
||||
meta_style = project.global_config.data.get("style", "")
|
||||
|
||||
if meta_style:
|
||||
json_prompt += f"\n\nStyle Consistency Requirement: {meta_style}. Ensure the visual style matches the provided context images."
|
||||
|
||||
logger.info(f"Calling AI service for panel {item.sequence}...")
|
||||
image_bytes = ai.generate_image(json_prompt, context_images)
|
||||
|
||||
relative_url = save_generated_image(session, project.id, "panel", item.id, image_bytes)
|
||||
item.image_url = relative_url
|
||||
session.add(item)
|
||||
|
||||
task.status = "completed"
|
||||
task.progress = 100
|
||||
session.add(task)
|
||||
session.commit()
|
||||
logger.info(f"Panel task {task_id} completed successfully.")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Panel task {task_id} failed: {e}")
|
||||
traceback.print_exc()
|
||||
task.status = "failed"
|
||||
task.message = str(e)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
|
||||
@router.post("/panel/{item_id}")
|
||||
def generate_panel(
|
||||
item_id: int,
|
||||
background_tasks: BackgroundTasks,
|
||||
session: Session = Depends(get_session)
|
||||
):
|
||||
item = session.get(StoryboardItem, item_id)
|
||||
if not item:
|
||||
raise HTTPException(status_code=404, detail="Storyboard item not found")
|
||||
|
||||
task = Task(
|
||||
type="image_generation",
|
||||
status="pending",
|
||||
project_id=item.project_id,
|
||||
name=f"Draw Panel: #{item.sequence}",
|
||||
description=f"Drawing panel {item.sequence}"
|
||||
)
|
||||
session.add(task)
|
||||
session.commit()
|
||||
session.refresh(task)
|
||||
|
||||
background_tasks.add_task(generate_panel_task, task.id, item_id)
|
||||
|
||||
return {"task_id": task.id}
|
||||
@@ -0,0 +1,38 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from sqlmodel import Session, select
|
||||
from app.core.database import get_session
|
||||
from app.models.models import ImageHistory, Character, StoryboardItem
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@router.get("/{entity_type}/{entity_id}")
|
||||
def get_history(entity_type: str, entity_id: int, session: Session = Depends(get_session)):
|
||||
# entity_type: 'character' or 'panel'
|
||||
statement = select(ImageHistory).where(
|
||||
ImageHistory.entity_type == entity_type,
|
||||
ImageHistory.entity_id == entity_id
|
||||
).order_by(ImageHistory.created_at.desc())
|
||||
history = session.exec(statement).all()
|
||||
return history
|
||||
|
||||
@router.post("/select/{history_id}")
|
||||
def select_image(history_id: int, session: Session = Depends(get_session)):
|
||||
history = session.get(ImageHistory, history_id)
|
||||
if not history:
|
||||
raise HTTPException(status_code=404, detail="History item not found")
|
||||
|
||||
if history.entity_type == "character":
|
||||
entity = session.get(Character, history.entity_id)
|
||||
elif history.entity_type == "panel":
|
||||
entity = session.get(StoryboardItem, history.entity_id)
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Unknown entity type")
|
||||
|
||||
if not entity:
|
||||
raise HTTPException(status_code=404, detail="Entity not found")
|
||||
|
||||
entity.image_url = history.image_url
|
||||
session.add(entity)
|
||||
session.commit()
|
||||
session.refresh(entity)
|
||||
return {"status": "success", "image_url": entity.image_url}
|
||||
@@ -0,0 +1,205 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Body
|
||||
from sqlmodel import Session
|
||||
from typing import List, Dict
|
||||
from pydantic import BaseModel
|
||||
from app.core.database import get_session
|
||||
from app.models.models import Project, GlobalConfig, Character, StoryboardItem
|
||||
from app.schemas.schemas import ProjectCreate, ProjectUpdate, ProjectRead
|
||||
from app.cruds import crud_project
|
||||
from app.services.consistency_service import ConsistencyService
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@router.post("/", response_model=Project)
|
||||
def create_project(project_in: ProjectCreate, session: Session = Depends(get_session)):
|
||||
return crud_project.create_project(session, project_in)
|
||||
|
||||
@router.get("/", response_model=List[Project])
|
||||
def read_projects(skip: int = 0, limit: int = 100, session: Session = Depends(get_session)):
|
||||
return crud_project.get_projects(session, skip, limit)
|
||||
|
||||
@router.get("/{project_id}", response_model=ProjectRead)
|
||||
def read_project(project_id: str, session: Session = Depends(get_session)):
|
||||
project = crud_project.get_project(session, project_id)
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
return project
|
||||
|
||||
@router.put("/{project_id}", response_model=Project)
|
||||
def update_project(project_id: str, project_in: ProjectUpdate, session: Session = Depends(get_session)):
|
||||
project = crud_project.get_project(session, project_id)
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
return crud_project.update_project(session, project, project_in)
|
||||
|
||||
@router.delete("/{project_id}")
|
||||
def delete_project(project_id: str, session: Session = Depends(get_session)):
|
||||
project = crud_project.get_project(session, project_id)
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
crud_project.delete_project(session, project)
|
||||
return {"ok": True}
|
||||
|
||||
# --- Data Management Endpoints ---
|
||||
|
||||
@router.put("/{project_id}/global_config")
|
||||
def update_global_config(project_id: str, data: Dict = Body(...), session: Session = Depends(get_session)):
|
||||
project = crud_project.get_project(session, project_id)
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# 1. Update Project level fields first
|
||||
# This ensures consistency between project.theme/language and global_config
|
||||
if "language" in data:
|
||||
project.language = data["language"]
|
||||
if "style" in data:
|
||||
project.theme = data["style"]
|
||||
if "aspect_ratio" in data:
|
||||
project.aspect_ratio = data["aspect_ratio"]
|
||||
|
||||
session.add(project)
|
||||
session.commit()
|
||||
session.refresh(project)
|
||||
|
||||
# 2. Update GlobalConfig in DB
|
||||
config = crud_project.create_global_config(session, project_id, data)
|
||||
|
||||
# 3. Trigger consistency check (Propagate to all items)
|
||||
consistency = ConsistencyService(session)
|
||||
consistency.normalize_project(project_id)
|
||||
|
||||
return config
|
||||
|
||||
@router.put("/{project_id}/characters/{char_id}")
|
||||
def update_character(project_id: str, char_id: int, data: Dict = Body(...), session: Session = Depends(get_session)):
|
||||
char = session.get(Character, char_id)
|
||||
if not char:
|
||||
raise HTTPException(status_code=404, detail="Character not found")
|
||||
|
||||
char.data = data
|
||||
session.add(char)
|
||||
session.commit()
|
||||
session.refresh(char)
|
||||
|
||||
# We might want to trigger consistency here too if character style changes,
|
||||
# but primarily it's driven by global config.
|
||||
return char
|
||||
|
||||
@router.put("/{project_id}/storyboard/{item_id}")
|
||||
def update_storyboard_item(project_id: str, item_id: int, data: Dict = Body(...), session: Session = Depends(get_session)):
|
||||
item = session.get(StoryboardItem, item_id)
|
||||
if not item:
|
||||
raise HTTPException(status_code=404, detail="Storyboard Item not found")
|
||||
|
||||
item.data = data
|
||||
session.add(item)
|
||||
session.commit()
|
||||
session.refresh(item)
|
||||
return item
|
||||
|
||||
@router.delete("/{project_id}/characters/{char_id}")
|
||||
def delete_character(project_id: str, char_id: int, session: Session = Depends(get_session)):
|
||||
char = session.get(Character, char_id)
|
||||
if not char:
|
||||
raise HTTPException(status_code=404, detail="Character not found")
|
||||
|
||||
session.delete(char)
|
||||
session.commit()
|
||||
return {"ok": True}
|
||||
|
||||
class MergeCharacterRequest(BaseModel):
|
||||
target_char_id: int
|
||||
source_char_ids: List[int]
|
||||
|
||||
@router.post("/{project_id}/characters/merge")
|
||||
def merge_characters(
|
||||
project_id: str,
|
||||
request: MergeCharacterRequest,
|
||||
session: Session = Depends(get_session)
|
||||
):
|
||||
target_char = session.get(Character, request.target_char_id)
|
||||
if not target_char:
|
||||
raise HTTPException(status_code=404, detail="Target character not found")
|
||||
|
||||
source_chars = []
|
||||
for cid in request.source_char_ids:
|
||||
c = session.get(Character, cid)
|
||||
if c:
|
||||
source_chars.append(c)
|
||||
|
||||
if not source_chars:
|
||||
raise HTTPException(status_code=400, detail="No valid source characters found")
|
||||
|
||||
target_name = target_char.name
|
||||
source_names = [c.name for c in source_chars]
|
||||
|
||||
# 1. Update Storyboard Items
|
||||
# We need to scan all items and replace source names with target name
|
||||
project = session.get(Project, project_id)
|
||||
if project.storyboard_items:
|
||||
for item in project.storyboard_items:
|
||||
data = item.data
|
||||
# 'characters' field in storyboard item data
|
||||
# It can be a list of strings, or list of dicts with 'name' key, or a single string
|
||||
chars = data.get("characters", [])
|
||||
|
||||
new_chars = []
|
||||
modified = False
|
||||
|
||||
# Helper to normalize input to list
|
||||
char_list = []
|
||||
if isinstance(chars, str): char_list = [chars]
|
||||
elif isinstance(chars, list): char_list = chars
|
||||
|
||||
for c_entry in char_list:
|
||||
c_name = ""
|
||||
if isinstance(c_entry, str): c_name = c_entry
|
||||
elif isinstance(c_entry, dict): c_name = c_entry.get("name", "")
|
||||
|
||||
# Check if this name matches any source name
|
||||
# Fuzzy match or exact? Let's do exact or containment for safety
|
||||
# User said "Ma Laoguanjia" vs "Ma Guanjia".
|
||||
# Ideally we replace if it matches one of the source characters' name EXACTLY or close enough?
|
||||
# Since we selected source characters by ID, we know their names.
|
||||
# Let's replace if the name in storyboard matches a source character name.
|
||||
|
||||
is_source = False
|
||||
for src_name in source_names:
|
||||
if src_name == c_name:
|
||||
is_source = True
|
||||
break
|
||||
|
||||
if is_source:
|
||||
# Replace with target name
|
||||
# If entry was dict, update name field? Or just use string?
|
||||
# Let's keep format.
|
||||
if isinstance(c_entry, str):
|
||||
new_chars.append(target_name)
|
||||
elif isinstance(c_entry, dict):
|
||||
c_entry['name'] = target_name
|
||||
new_chars.append(c_entry)
|
||||
modified = True
|
||||
else:
|
||||
new_chars.append(c_entry)
|
||||
|
||||
if modified:
|
||||
# Deduplicate if target name already existed?
|
||||
# Simple dedup for strings
|
||||
final_chars = []
|
||||
seen = set()
|
||||
for c in new_chars:
|
||||
n = c if isinstance(c, str) else c.get("name", "")
|
||||
if n not in seen:
|
||||
final_chars.append(c)
|
||||
seen.add(n)
|
||||
|
||||
data['characters'] = final_chars
|
||||
item.data = data
|
||||
session.add(item)
|
||||
|
||||
# 2. Delete Source Characters
|
||||
for c in source_chars:
|
||||
session.delete(c)
|
||||
|
||||
session.commit()
|
||||
return {"ok": True, "merged_count": len(source_chars)}
|
||||
@@ -0,0 +1,23 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from sqlmodel import Session
|
||||
from app.core.database import get_session
|
||||
from app.models.models import Task
|
||||
from app.schemas.schemas import TaskRead
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@router.get("/{task_id}", response_model=TaskRead)
|
||||
def get_task_status(task_id: str, session: Session = Depends(get_session)):
|
||||
task = session.get(Task, task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="Task not found")
|
||||
return task
|
||||
|
||||
@router.get("/project/{project_id}", response_model=list[TaskRead])
|
||||
def get_project_tasks(project_id: str, session: Session = Depends(get_session)):
|
||||
from sqlmodel import select
|
||||
statement = select(Task).where(Task.project_id == project_id).order_by(Task.created_at.desc())
|
||||
tasks = session.exec(statement).all()
|
||||
# Filter only recent or active tasks if list is too long?
|
||||
# For now return all, maybe limit 20
|
||||
return tasks[:20]
|
||||
@@ -0,0 +1,59 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, Dict, List
|
||||
from datetime import datetime
|
||||
from app.models.models import (
|
||||
ModelConfigBase, ProjectBase, CharacterBase, StoryboardItemBase, GlobalConfigBase, TaskBase,
|
||||
ModelConfig, Project, Character, StoryboardItem, GlobalConfig, Task
|
||||
)
|
||||
|
||||
# ModelConfig
|
||||
class ModelConfigCreate(ModelConfigBase):
|
||||
pass
|
||||
|
||||
class ModelConfigUpdate(BaseModel):
|
||||
provider: Optional[str] = None
|
||||
api_key: Optional[str] = None
|
||||
base_url: Optional[str] = None
|
||||
model_name: Optional[str] = None
|
||||
model_type: Optional[str] = None
|
||||
is_active: Optional[bool] = None
|
||||
|
||||
# Read Models for nested response
|
||||
class CharacterRead(CharacterBase):
|
||||
id: int
|
||||
project_id: str
|
||||
|
||||
class StoryboardItemRead(StoryboardItemBase):
|
||||
id: int
|
||||
project_id: str
|
||||
|
||||
class GlobalConfigRead(GlobalConfigBase):
|
||||
id: int
|
||||
project_id: str
|
||||
|
||||
class TaskRead(TaskBase):
|
||||
id: str
|
||||
project_id: str
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
# Project
|
||||
class ProjectCreate(ProjectBase):
|
||||
pass
|
||||
|
||||
class ProjectUpdate(BaseModel):
|
||||
title: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
story_input: Optional[str] = None
|
||||
theme: Optional[str] = None
|
||||
language: Optional[str] = None
|
||||
panel_count: Optional[int] = None
|
||||
aspect_ratio: Optional[str] = None
|
||||
|
||||
class ProjectRead(ProjectBase):
|
||||
id: str
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
characters: List[CharacterRead] = []
|
||||
storyboard_items: List[StoryboardItemRead] = []
|
||||
global_config: Optional[GlobalConfigRead] = None
|
||||
@@ -0,0 +1,92 @@
|
||||
import os
|
||||
import time
|
||||
from typing import List, Optional
|
||||
from sqlmodel import Session
|
||||
from app.models.models import ModelConfig
|
||||
from google import genai
|
||||
from PIL import Image
|
||||
import io
|
||||
|
||||
class AIService:
|
||||
def __init__(self, session: Session):
|
||||
self.session = session
|
||||
|
||||
def _get_client(self, model_type: str):
|
||||
from app.cruds.crud_config import get_active_config
|
||||
config = get_active_config(self.session, model_type)
|
||||
|
||||
if not config:
|
||||
raise ValueError(f"No active configuration found for {model_type} model.")
|
||||
|
||||
if config.provider.lower() == "google":
|
||||
# Initialize Google Client
|
||||
return genai.Client(api_key=config.api_key), config.model_name
|
||||
|
||||
raise NotImplementedError(f"Provider {config.provider} not supported yet.")
|
||||
|
||||
def generate_storyboard(self, system_prompt: str, user_input: str) -> str:
|
||||
client, model_name = self._get_client("text")
|
||||
|
||||
full_prompt = f"{system_prompt}\n\nUser Input: {user_input}\n\nPlease generate the full storyboard in JSON format as requested."
|
||||
|
||||
try:
|
||||
response = client.models.generate_content(
|
||||
model=model_name,
|
||||
contents=full_prompt
|
||||
)
|
||||
return response.text
|
||||
except Exception as e:
|
||||
print(f"Error generating storyboard: {e}")
|
||||
raise e
|
||||
|
||||
def generate_image(self, prompt: str, context_images: List[str] = None) -> bytes:
|
||||
client, model_name = self._get_client("image")
|
||||
|
||||
contents = [prompt]
|
||||
if context_images:
|
||||
for img_path in context_images:
|
||||
if os.path.exists(img_path):
|
||||
try:
|
||||
prev_img = Image.open(img_path)
|
||||
contents.append(prev_img)
|
||||
except Exception as e:
|
||||
print(f"Failed to load context image {img_path}: {e}")
|
||||
else:
|
||||
# Log missing context image but don't fail, just skip it
|
||||
print(f"Warning: Context image not found at {img_path}, skipping.")
|
||||
|
||||
# Retry loop
|
||||
max_retries = 3
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
response = client.models.generate_content(
|
||||
model=model_name,
|
||||
contents=contents,
|
||||
)
|
||||
|
||||
if response.parts:
|
||||
for part in response.parts:
|
||||
if part.inline_data is not None:
|
||||
image_data = part.inline_data.data
|
||||
if len(image_data) > 0:
|
||||
return image_data
|
||||
else:
|
||||
print(f"Warning: Received empty image data on attempt {attempt + 1}")
|
||||
|
||||
# Check for text refusal/error
|
||||
if response.text:
|
||||
print(f"Model response text (no image): {response.text}")
|
||||
|
||||
print(f"Attempt {attempt + 1} failed: No valid image data found in response.")
|
||||
if attempt == max_retries - 1:
|
||||
raise ValueError(f"No image found in response after {max_retries} retries. Last response: {response.text if response.text else 'Empty'}")
|
||||
|
||||
time.sleep(2 ** attempt)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error generating image (Attempt {attempt + 1}/{max_retries}): {e}")
|
||||
if attempt < max_retries - 1:
|
||||
time.sleep(2 ** attempt)
|
||||
continue
|
||||
raise e
|
||||
return b""
|
||||
@@ -0,0 +1,181 @@
|
||||
import re
|
||||
import json
|
||||
import copy
|
||||
from typing import List, Dict, Any
|
||||
from sqlmodel import Session, select
|
||||
from app.models.models import Project, Character, StoryboardItem, GlobalConfig
|
||||
|
||||
class ConsistencyService:
|
||||
def __init__(self, session: Session):
|
||||
self.session = session
|
||||
|
||||
def normalize_project(self, project_id: str):
|
||||
"""
|
||||
Normalizes the project's data (characters, storyboard) based on the global config.
|
||||
This mirrors the logic in comic_generator.py.
|
||||
"""
|
||||
project = self.session.get(Project, project_id)
|
||||
if not project:
|
||||
return
|
||||
|
||||
# Fetch all related data
|
||||
# Note: relationships are loaded if accessed, but let's be explicit if needed.
|
||||
# ProjectRead should handle loading, but here we work with ORM objects.
|
||||
|
||||
global_config = project.global_config
|
||||
characters = project.characters
|
||||
storyboard_items = sorted(project.storyboard_items, key=lambda x: x.sequence)
|
||||
|
||||
master_style = None
|
||||
master_meta = {}
|
||||
|
||||
# Priority 0: Comic Config
|
||||
if global_config and global_config.data:
|
||||
config_data = global_config.data
|
||||
master_style = config_data.get("style")
|
||||
|
||||
# Explicitly map all fields we want to sync
|
||||
master_meta = {
|
||||
"style": master_style,
|
||||
"bubble_style": config_data.get("bubble_style"),
|
||||
"narration_style": config_data.get("narration_style"),
|
||||
"border_style": config_data.get("border_style"),
|
||||
"gutter_style": config_data.get("gutter_style"),
|
||||
"layout_settings": config_data.get("layout_settings"),
|
||||
"aspect_ratio": config_data.get("aspect_ratio", "16:9"),
|
||||
"language": config_data.get("language", "English")
|
||||
}
|
||||
|
||||
# Priority 1: First Character Sheet (if no config style)
|
||||
if not master_style and characters:
|
||||
# Try to find one with style
|
||||
for char in characters:
|
||||
if char.data.get("style"):
|
||||
master_style = char.data.get("style")
|
||||
break
|
||||
|
||||
# Priority 2: First Story Block
|
||||
if storyboard_items:
|
||||
first_item = storyboard_items[0]
|
||||
first_meta = first_item.data.get("meta_info", {})
|
||||
if not master_style:
|
||||
master_style = first_meta.get("style")
|
||||
|
||||
if not master_meta:
|
||||
master_meta = first_meta.copy()
|
||||
|
||||
# Build Character Registry
|
||||
known_characters = {}
|
||||
for char in characters:
|
||||
full_name = char.name
|
||||
if not full_name: continue
|
||||
|
||||
keywords = [full_name]
|
||||
simplified = re.split(r'[(\(]', full_name)[0].strip()
|
||||
if simplified and simplified != full_name:
|
||||
keywords.append(simplified)
|
||||
|
||||
for kw in keywords:
|
||||
if kw:
|
||||
known_characters[kw] = full_name
|
||||
|
||||
# Apply Normalization
|
||||
# Even if master_style is None, we might still have layout_settings to sync
|
||||
# So we check if we have ANY master_meta to apply
|
||||
if master_meta:
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"Normalizing project {project_id} with master config: {json.dumps(master_meta, ensure_ascii=False)}")
|
||||
|
||||
# 1. Normalize Characters
|
||||
for char in characters:
|
||||
char_data = copy.deepcopy(char.data)
|
||||
if not isinstance(char_data, dict):
|
||||
char_data = dict(char_data)
|
||||
|
||||
meta = char_data.get("meta_info", {})
|
||||
|
||||
# Sync Core Fields
|
||||
if master_meta.get("style"): meta["style"] = master_meta.get("style")
|
||||
if master_meta.get("language"): meta["language"] = master_meta.get("language")
|
||||
if master_meta.get("aspect_ratio"): meta["aspect_ratio"] = master_meta.get("aspect_ratio")
|
||||
|
||||
char_data["meta_info"] = meta
|
||||
|
||||
# Cleanup top-level legacy fields
|
||||
char_data.pop("style", None)
|
||||
char_data.pop("language", None)
|
||||
char_data.pop("Language", None)
|
||||
|
||||
char.data = char_data
|
||||
self.session.add(char)
|
||||
|
||||
# 2. Normalize Storyboard Items
|
||||
total_volumes = len(storyboard_items)
|
||||
|
||||
for i, item in enumerate(storyboard_items):
|
||||
# Use deepcopy to ensure we don't mutate the original object in place before assignment
|
||||
# and to ensure SQLAlchemy detects the change when we reassign.
|
||||
item_data = copy.deepcopy(item.data)
|
||||
if not isinstance(item_data, dict):
|
||||
item_data = dict(item_data)
|
||||
|
||||
# Check for missing characters
|
||||
plot_text = json.dumps(item_data.get("plot_breakdown", []), ensure_ascii=False)
|
||||
current_chars = item_data.get("characters", [])
|
||||
if isinstance(current_chars, str):
|
||||
current_chars = [current_chars]
|
||||
if not isinstance(current_chars, list):
|
||||
current_chars = []
|
||||
|
||||
found_missing = []
|
||||
for kw, full_name in known_characters.items():
|
||||
if kw in plot_text:
|
||||
is_present = False
|
||||
for char_name in current_chars:
|
||||
if kw in char_name or char_name in full_name:
|
||||
is_present = True
|
||||
break
|
||||
if not is_present:
|
||||
if full_name not in current_chars and full_name not in found_missing:
|
||||
found_missing.append(full_name)
|
||||
|
||||
if found_missing:
|
||||
current_chars.extend(found_missing)
|
||||
item_data["characters"] = current_chars
|
||||
|
||||
# Sync Meta Info
|
||||
meta = item_data.get("meta_info", {})
|
||||
|
||||
# We want to preserve existing fields in meta (like volume) but overwrite style configs
|
||||
# master_meta has style, bubble_style, etc.
|
||||
|
||||
# Force update fields from master_meta even if they exist in meta
|
||||
# But careful with 'None' values in master_meta (though we constructed it from config)
|
||||
|
||||
for key, value in master_meta.items():
|
||||
if key == "volume": continue
|
||||
|
||||
# Special handling for boolean values (like False in show_panel_numbers)
|
||||
# "if value is not None" is correct for booleans.
|
||||
# For nested dictionaries (like layout_settings), we should merge to preserve other keys.
|
||||
|
||||
if value is not None:
|
||||
if isinstance(value, dict) and isinstance(meta.get(key), dict):
|
||||
meta[key].update(value)
|
||||
else:
|
||||
meta[key] = value
|
||||
|
||||
# If value is None but key exists in master_meta keys (explicitly set to null), we might want to unset it?
|
||||
# But our master_meta construction uses .get() defaults, so None usually means "not in config".
|
||||
|
||||
# Ensure Volume format
|
||||
meta["volume"] = f"{i+1}/{total_volumes}"
|
||||
item_data["meta_info"] = meta
|
||||
|
||||
logger.info(f"Updated Item {i+1} meta: {json.dumps(meta, ensure_ascii=False)}")
|
||||
|
||||
item.data = item_data
|
||||
self.session.add(item)
|
||||
|
||||
self.session.commit()
|
||||
@@ -0,0 +1,30 @@
|
||||
import os
|
||||
from PIL import Image
|
||||
import io
|
||||
|
||||
def split_comic_page(image_bytes: bytes) -> list[bytes]:
|
||||
"""Splits a 2x2 grid comic page into 4 individual panel images (bytes)."""
|
||||
try:
|
||||
img = Image.open(io.BytesIO(image_bytes))
|
||||
width, height = img.size
|
||||
mid_w = width // 2
|
||||
mid_h = height // 2
|
||||
|
||||
quadrants = [
|
||||
(0, 0, mid_w, mid_h), # Top-Left
|
||||
(mid_w, 0, width, mid_h), # Top-Right
|
||||
(0, mid_h, mid_w, height), # Bottom-Left
|
||||
(mid_w, mid_h, width, height) # Bottom-Right
|
||||
]
|
||||
|
||||
panels = []
|
||||
for box in quadrants:
|
||||
panel = img.crop(box)
|
||||
buf = io.BytesIO()
|
||||
panel.save(buf, format="PNG")
|
||||
panels.append(buf.getvalue())
|
||||
|
||||
return panels
|
||||
except Exception as e:
|
||||
print(f"Failed to split image: {e}")
|
||||
return []
|
||||
@@ -0,0 +1,72 @@
|
||||
import json
|
||||
import re
|
||||
from typing import List, Dict, Any
|
||||
|
||||
def extract_json_blocks(text: str) -> List[Dict[str, Any]]:
|
||||
"""Extracts JSON blocks from the generated text."""
|
||||
json_blocks = []
|
||||
|
||||
# 1. Try to find ```json ... ``` blocks
|
||||
pattern = r"```(?:json|JSON)?\s*(.*?)\s*```"
|
||||
matches = re.findall(pattern, text, re.DOTALL)
|
||||
|
||||
# 2. If no code blocks found, or even if found, we should also look for raw JSON objects
|
||||
# because sometimes models output mixed content.
|
||||
# But let's stick to code blocks first if they exist.
|
||||
|
||||
if not matches:
|
||||
# Fallback: Try to find top-level JSON objects/arrays directly in text
|
||||
# This regex looks for { ... } or [ ... ] that span multiple lines
|
||||
# It's not perfect but better than nothing
|
||||
# We search for anything starting with { or [ and ending with } or ]
|
||||
# non-greedy match might be safer for multiple blocks
|
||||
raw_pattern = r"(\{[\s\S]*?\}|\[[\s\S]*?\])"
|
||||
matches = re.findall(raw_pattern, text)
|
||||
|
||||
def repair_json(json_str: str) -> Dict[str, Any]:
|
||||
"""Attempts to repair common JSON errors, specifically missing commas."""
|
||||
try:
|
||||
return json.loads(json_str)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
repaired_str = json_str
|
||||
max_attempts = 10
|
||||
|
||||
for _ in range(max_attempts):
|
||||
try:
|
||||
return json.loads(repaired_str)
|
||||
except json.JSONDecodeError as e:
|
||||
# print(f"JSON Decode Error at {e.pos}: {e.msg}")
|
||||
if "Expecting ',' delimiter" in str(e) or "Expecting property name enclosed in double quotes" in str(e):
|
||||
pos = e.pos
|
||||
search_str = repaired_str[:pos]
|
||||
match = re.search(r'([\"}\]0-9])\s*$', search_str)
|
||||
|
||||
if match:
|
||||
insert_idx = match.end()
|
||||
repaired_str = repaired_str[:insert_idx] + "," + repaired_str[insert_idx:]
|
||||
continue
|
||||
|
||||
# If we can't fix it, re-raise
|
||||
raise e
|
||||
|
||||
return json.loads(repaired_str)
|
||||
|
||||
for match in matches:
|
||||
if not match.strip(): continue
|
||||
# Simple heuristic to filter out non-json text blocks that might be caught by raw regex
|
||||
if not (match.strip().startswith('{') or match.strip().startswith('[')):
|
||||
continue
|
||||
|
||||
try:
|
||||
data = repair_json(match)
|
||||
if isinstance(data, list):
|
||||
json_blocks.extend(data)
|
||||
elif isinstance(data, dict):
|
||||
json_blocks.append(data)
|
||||
except Exception as e:
|
||||
print(f"Failed to parse a JSON block: {e}")
|
||||
# print(f"Block content: {match[:100]}...")
|
||||
|
||||
return json_blocks
|
||||
Reference in New Issue
Block a user