diff --git a/backend/requirements.txt b/backend/requirements.txt index ec1760e..89d5ae0 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -8,3 +8,4 @@ python-dotenv==1.0.0 numpy==2.2.0 librosa==0.10.2 soundfile==0.12.1 +google-genai==1.0.0 diff --git a/backend/src/image_gen.py b/backend/src/image_gen.py new file mode 100644 index 0000000..9ca41ec --- /dev/null +++ b/backend/src/image_gen.py @@ -0,0 +1,110 @@ +"""Image generation client using Google Gemini Image models.""" + +import os +from pathlib import Path +from typing import Optional + +from google import genai +from google.genai import types + +ARTWORK_DIR = Path(os.environ.get("ARTWORK_DIR", "/data/artwork")) +ARTWORK_DIR.mkdir(parents=True, exist_ok=True) + +# ── Configuration ────────────────────────────────────────────────────────── + +def get_image_gen_config(): + """Load image generation configuration from environment.""" + return { + "api_key": os.environ.get("IMAGE_GEN_API_KEY") or os.environ.get("LYRIA_API_KEY"), + "model": os.environ.get("IMAGE_GEN_MODEL", "gemini-3.1-flash-image-preview"), + "safety_harassment": os.environ.get("IMAGE_GEN_SAFETY_HARASSMENT", "BLOCK_ONLY_HIGH"), + "safety_hate_speech": os.environ.get("IMAGE_GEN_SAFETY_HATE_SPEECH", "BLOCK_ONLY_HIGH"), + "safety_sexually_explicit": os.environ.get("IMAGE_GEN_SAFETY_SEXUALLY_EXPLICIT", "BLOCK_ONLY_HIGH"), + "safety_dangerous": os.environ.get("IMAGE_GEN_SAFETY_DANGEROUS", "BLOCK_ONLY_HIGH"), + } + +# ── Safety Settings ──────────────────────────────────────────────────────── + +SAFETY_LEVELS = { + "off": "OFF", + "block_none": "BLOCK_NONE", + "block_few": "BLOCK_ONLY_HIGH", + "block_some": "BLOCK_MEDIUM_AND_ABOVE", + "block_most": "BLOCK_LOW_AND_ABOVE", +} + +def build_safety_settings(config: dict) -> list: + """Build safety settings from config.""" + return [ + types.SafetySetting( + category="HARM_CATEGORY_HARASSMENT", + threshold=config.get("safety_harassment", "BLOCK_ONLY_HIGH"), + ), + types.SafetySetting( + category="HARM_CATEGORY_HATE_SPEECH", + threshold=config.get("safety_hate_speech", "BLOCK_ONLY_HIGH"), + ), + types.SafetySetting( + category="HARM_CATEGORY_SEXUALLY_EXPLICIT", + threshold=config.get("safety_sexually_explicit", "BLOCK_ONLY_HIGH"), + ), + types.SafetySetting( + category="HARM_CATEGORY_DANGEROUS_CONTENT", + threshold=config.get("safety_dangerous", "BLOCK_ONLY_HIGH"), + ), + ] + +# ── Generation ────────────────────────────────────────────────────────────── + +async def generate_artwork(prompt: str, config: Optional[dict] = None) -> bytes: + """Generate album artwork from a text prompt. + + Returns the raw PNG image bytes. + """ + cfg = config or get_image_gen_config() + client = genai.Client(api_key=cfg["api_key"]) + + safety = build_safety_settings(cfg) + + response = client.models.generate_content( + model=cfg["model"], + contents=prompt, + config=types.GenerateContentConfig( + safety_settings=safety, + response_modalities=["image"], + ), + ) + + # Extract image from response + for part in response.candidates[0].content.parts: + if part.inline_data: + return part.inline_data.data + + raise RuntimeError("No image data in response") + +# ── Available Models ──────────────────────────────────────────────────────── + +AVAILABLE_IMAGE_MODELS = [ + { + "id": "gemini-3.1-flash-image-preview", + "name": "Gemini 3.1 Flash Image", + "description": "Fast, cost-efficient image generation. Best for album art.", + "status": "Preview", + }, + { + "id": "gemini-3-pro-image-preview", + "name": "Gemini 3 Pro Image", + "description": "Highest quality image generation with reasoning.", + "status": "Preview", + }, + { + "id": "gemini-2.5-flash-image", + "name": "Gemini 2.5 Flash Image", + "description": "GA image generation. Solid quality, good speed.", + "status": "GA", + }, +] + +def get_available_models() -> list: + """Return list of available image generation models.""" + return AVAILABLE_IMAGE_MODELS diff --git a/backend/src/main.py b/backend/src/main.py index af5c231..cef2e20 100644 --- a/backend/src/main.py +++ b/backend/src/main.py @@ -9,11 +9,14 @@ from dotenv import load_dotenv # Load environment variables BEFORE importing lyria_client (it reads env at module level) load_dotenv(Path(__file__).parent.parent / ".env") -from models import GenerateRequest, GenerateResponse, ClipList, RenameRequest, LYRIA_MODELS, DEFAULT_MODEL +from models import GenerateRequest, GenerateResponse, ClipList, RenameRequest, GenerateArtworkRequest, LYRIA_MODELS, DEFAULT_MODEL from database import init_db, create_clip, get_clips, get_clip, delete_clip, update_clip_status, rename_clip, set_artwork, remove_artwork, delete_all_clips, set_analysis from lyria_client import generate_clip, save_clip_file from waveform import generate_waveform_svg, get_waveform_path from analysis import analyze_clip +from image_gen import generate_artwork, get_image_gen_config, get_available_models +from analysis import analyze_clip +from image_gen import generate_artwork, get_image_gen_config, get_available_models, SAFETY_LEVELS app = FastAPI(title="Lyria Studio", version="0.1.0") @@ -243,6 +246,58 @@ async def remove_artwork_endpoint(clip_id: str): updated = await get_clip(clip_id) return updated +@app.post("/api/clips/{clip_id}/generate-artwork") +async def generate_clip_artwork(clip_id: str, request: GenerateArtworkRequest): + """Generate album artwork using Gemini Image model.""" + clip = await get_clip(clip_id) + if not clip: + raise HTTPException(status_code=404, detail="Clip not found") + + try: + # Generate image + config = get_image_gen_config() + image_bytes = await generate_artwork(request.prompt, config) + + # Save to disk + import uuid + artwork_filename = f"{clip_id}_{uuid.uuid4().hex[:8]}.png" + artwork_path = ARTWORK_DIR / artwork_filename + with open(artwork_path, "wb") as f: + f.write(image_bytes) + + # Remove old artwork if exists + if clip.get("artwork_filename"): + old_path = ARTWORK_DIR / clip["artwork_filename"] + if old_path.exists(): + old_path.unlink() + + # Update DB + await set_artwork(clip_id, artwork_filename) + updated = await get_clip(clip_id) + return updated + + except Exception as e: + raise HTTPException(status_code=500, detail=f"Image generation failed: {str(e)}") + +@app.get("/api/image-gen/models") +async def list_image_gen_models(): + """List available image generation models.""" + return {"models": get_available_models()} + +@app.get("/api/image-gen/config") +async def get_image_gen_configuration(): + """Get current image generation configuration (safe values only).""" + config = get_image_gen_config() + return { + "model": config["model"], + "safety_harassment": config["safety_harassment"], + "safety_hate_speech": config["safety_hate_speech"], + "safety_sexually_explicit": config["safety_sexually_explicit"], + "safety_dangerous": config["safety_dangerous"], + # API key is masked + "has_api_key": bool(config.get("api_key")), + } + @app.patch("/api/clips/{clip_id}") async def rename_clip_endpoint(clip_id: str, request: RenameRequest): clip = await get_clip(clip_id) diff --git a/backend/src/models.py b/backend/src/models.py index da45270..fb0ec61 100644 --- a/backend/src/models.py +++ b/backend/src/models.py @@ -61,6 +61,10 @@ class Clip(BaseModel): class RenameRequest(BaseModel): title: str +class GenerateArtworkRequest(BaseModel): + prompt: str + class ClipList(BaseModel): clips: list[Clip] + total: int total: int \ No newline at end of file diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 3f496a0..29f9064 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -4,6 +4,7 @@ import Gallery from './components/Gallery' import Player from './components/Player' import DetailPanel from './components/DetailPanel' import SettingsModal from './components/SettingsModal' +import ArtworkGenerationModal from './components/ArtworkGenerationModal' interface Clip { id: string @@ -26,6 +27,8 @@ function App() { const [isPlaying, setIsPlaying] = useState(false) const [selectedClip, setSelectedClip] = useState(null) const [settingsOpen, setSettingsOpen] = useState(false) + const [artworkGenOpen, setArtworkGenOpen] = useState(false) + const [artworkGenClip, setArtworkGenClip] = useState(null) const fetchClips = async () => { try { @@ -134,6 +137,38 @@ function App() { } } + const handleGenerateArtwork = async (clip: Clip) => { + setArtworkGenClip(clip) + setArtworkGenOpen(true) + } + + const handleArtworkGenSubmit = async (prompt: string) => { + if (!artworkGenClip) return + try { + const res = await fetch(`/api/clips/${artworkGenClip.id}/generate-artwork`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ prompt }), + }) + if (res.ok) { + const updated = await res.json() + setClips(prev => prev.map(c => c.id === updated.id ? updated : c)) + if (currentClip?.id === updated.id) { + setCurrentClip(updated) + } + if (selectedClip?.id === updated.id) { + setSelectedClip(updated) + } + } else { + const err = await res.json() + alert(err.detail || 'Failed to generate artwork') + } + } catch (err) { + console.error('Failed to generate artwork:', err) + alert('Failed to generate artwork') + } + } + const handleSelectClip = async (clip: Clip) => { setSelectedClip(clip) // Fetch analysis if not already cached @@ -238,6 +273,19 @@ function App() { onRemoveArtwork={handleRemoveArtwork} onPlay={handlePlay} onTogglePlay={() => setIsPlaying(!isPlaying)} + onGenerateArtwork={handleGenerateArtwork} + /> + )} + + {artworkGenOpen && artworkGenClip && ( + { + setArtworkGenOpen(false) + setArtworkGenClip(null) + }} + onGenerate={handleArtworkGenSubmit} + trackName={artworkGenClip.title || artworkGenClip.prompt} /> )} diff --git a/frontend/src/components/ArtworkGenerationModal.tsx b/frontend/src/components/ArtworkGenerationModal.tsx new file mode 100644 index 0000000..457e3cd --- /dev/null +++ b/frontend/src/components/ArtworkGenerationModal.tsx @@ -0,0 +1,103 @@ +import { useState } from 'react' + +interface Props { + isOpen: boolean + onClose: () => void + onGenerate: (prompt: string) => Promise + trackName: string +} + +export default function ArtworkGenerationModal({ isOpen, onClose, onGenerate, trackName }: Props) { + const [prompt, setPrompt] = useState('') + const [loading, setLoading] = useState(false) + + if (!isOpen) return null + + const handleGenerate = async () => { + if (!prompt.trim() || loading) return + setLoading(true) + try { + await onGenerate(prompt.trim()) + setPrompt('') + onClose() + } finally { + setLoading(false) + } + } + + const handleClose = () => { + if (!loading) { + setPrompt('') + onClose() + } + } + + return ( +
+ {/* Backdrop */} +
+ + {/* Modal */} +
+ {/* Header */} +
+
+

Generate Artwork

+

{trackName}

+
+ +
+ + {/* Content */} +
+
+ +