diff --git a/actions/generate_storyboard.py b/actions/generate_storyboard.py index 81666f99..8b7028ad 100644 --- a/actions/generate_storyboard.py +++ b/actions/generate_storyboard.py @@ -285,19 +285,25 @@ def execute( # Group images by product ID to build context products: dict[str, Product] = {} + valid_product_image_combinations = {} for img_obj in images: product_id = str(img_obj.get(Dimension.PRODUCT_ID.value, "1")) - product_description = str(img_obj.get("product_description", None)) + desc = img_obj.get("product_description") + product_description = str(desc) if desc is not None else "" if product_id not in products: products[product_id] = Product(product_id, product_description) + if product_id not in valid_product_image_combinations: + valid_product_image_combinations[product_id] = set() + image_id = str(img_obj.get(Dimension.IMAGE_ID.value, "1")) product = products[product_id] image = Image( - id=str(img_obj.get(Dimension.IMAGE_ID.value, "1")), + id=image_id, uri=gcs.get_uri(img_obj[Key.FILE.value]), ) product.images.append(image) + valid_product_image_combinations[product_id].add(image_id) # Add products and images prompt_parts.append(genai.types.Part.from_text(text="### Products & Images:\n\n")) @@ -373,41 +379,71 @@ def execute( ], ) - start_time = time.time() - response = client.models.generate_content( - model=gemini_model, - contents=prompt_parts, - config=config, - ) - end_time = time.time() - logger.info( - "Gemini API request completed in %.2f seconds.", end_time - start_time - ) - - segments = [] - if ( - response.candidates - and response.candidates[0].content - and response.candidates[0].content.parts - ): - parts = list(response.candidates[0].content.parts) - for part in parts: - if part.text: - segments.append(part.text) - storyboard_text = "".join(segments) - if not storyboard_text: - raise ValueError( - "The model returned no storyboard text. This can happen when the" - " model reaches the output token limit before emitting a response." + result = None + for _ in range(4): # 1 initial call + 3 retries + start_time = time.time() + response = client.models.generate_content( + model=gemini_model, + contents=prompt_parts, + config=config, + ) + end_time = time.time() + logger.info( + "Gemini API request completed in %.2f seconds.", end_time - start_time ) - json_result = json.loads(storyboard_text) - if not json_result.get("storyboard"): - raise ValueError("No storyboard found in the response.") + segments = [] + if ( + response.candidates + and response.candidates[0].content + and response.candidates[0].content.parts + ): + parts = list(response.candidates[0].content.parts) + for part in parts: + if part.text: + segments.append(part.text) + storyboard_text = "".join(segments) + if not storyboard_text: + raise ValueError( + "The model returned no storyboard text. This can happen when the" + " model reaches the output token limit before emitting a response." + ) + candidate_result = json.loads(storyboard_text) + + if not candidate_result.get("storyboard"): + raise ValueError("No storyboard found in the response.") + + is_valid = True + for scene in candidate_result["storyboard"]: + product_id = scene.get(Dimension.PRODUCT_ID.value) + image_id = scene.get(Dimension.IMAGE_ID.value) + video_prompt = scene.get("video_prompt") + scene_name = scene.get("scene_name") + if ( + product_id not in valid_product_image_combinations + or image_id + not in valid_product_image_combinations.get(product_id, set()) + or not video_prompt + or not scene_name + ): + is_valid = False + logger.warning( + "Invalid scene from Gemini: product_id=%s, image_id=%s", + product_id, + image_id, + ) + break # This scene is invalid, so the whole result is. + if is_valid: + result = candidate_result + break + if not result: + raise RuntimeError( + "Gemini repeatedly scripted invalid image/product combinations" + ) valid_scenes = [] - for scene in json_result["storyboard"]: + for scene in result["storyboard"]: image_id = scene.get(Dimension.IMAGE_ID.value) product_id = scene.get(Dimension.PRODUCT_ID.value) video_prompt = scene.get("video_prompt") diff --git a/actions/test/test_generate_storyboard.py b/actions/test/test_generate_storyboard.py index 59dfc32f..863751bf 100644 --- a/actions/test/test_generate_storyboard.py +++ b/actions/test/test_generate_storyboard.py @@ -27,6 +27,235 @@ class TestGenerateStoryboard(unittest.TestCase): + def _make_mock_response(self, storyboard_dict): + mock_response = MagicMock() + mock_candidate = MagicMock() + mock_part = MagicMock() + mock_part.text = json.dumps(storyboard_dict) + mock_candidate.content.parts = [mock_part] + mock_response.candidates = [mock_candidate] + return mock_response + + @patch("actions.generate_storyboard.genai.Client") + def test_execute_hallucinated_product_id_triggers_retry( + self, mock_genai_client_class + ): + mock_client = MagicMock() + mock_genai_client_class.return_value = mock_client + invalid_resp = self._make_mock_response({ + "storyboard": [{ + "image_id": "i1", + "product_id": "p_invalid", + "scene_name": "Scene 1", + "video_prompt": "Prompt 1", + }] + }) + valid_resp = self._make_mock_response({ + "storyboard": [{ + "image_id": "i1", + "product_id": "p1", + "scene_name": "Scene 1", + "video_prompt": "Prompt 1", + }] + }) + mock_client.models.generate_content.side_effect = [invalid_resp, valid_resp] + + result = generate_storyboard.execute( + self.mock_gcs, + self.mock_params, + self.mock_images, + self.mock_user_prompt, + self.gemini_model, + self.gemini_model_location, + ) + + self.assertEqual(mock_client.models.generate_content.call_count, 2) + self.assertEqual( + result, + {"storyboard": [{Key.FILE.value: "gs://bucket/storyboard.json"}]}, + ) + + @patch("actions.generate_storyboard.genai.Client") + def test_execute_hallucinated_image_id_triggers_retry( + self, mock_genai_client_class + ): + mock_client = MagicMock() + mock_genai_client_class.return_value = mock_client + invalid_resp = self._make_mock_response({ + "storyboard": [{ + "image_id": "i_invalid", + "product_id": "p1", + "scene_name": "Scene 1", + "video_prompt": "Prompt 1", + }] + }) + valid_resp = self._make_mock_response({ + "storyboard": [{ + "image_id": "i1", + "product_id": "p1", + "scene_name": "Scene 1", + "video_prompt": "Prompt 1", + }] + }) + mock_client.models.generate_content.side_effect = [invalid_resp, valid_resp] + + result = generate_storyboard.execute( + self.mock_gcs, + self.mock_params, + self.mock_images, + self.mock_user_prompt, + self.gemini_model, + self.gemini_model_location, + ) + + self.assertEqual(mock_client.models.generate_content.call_count, 2) + self.assertEqual( + result, + {"storyboard": [{Key.FILE.value: "gs://bucket/storyboard.json"}]}, + ) + + @patch("actions.generate_storyboard.genai.Client") + def test_execute_empty_scene_fields_trigger_retry( + self, mock_genai_client_class + ): + for field in ("video_prompt", "scene_name"): + with self.subTest(field=field): + mock_client = MagicMock() + mock_genai_client_class.return_value = mock_client + scene = { + "image_id": "i1", + "product_id": "p1", + "scene_name": "Scene 1", + "video_prompt": "Prompt 1", + } + empty_resp = self._make_mock_response( + {"storyboard": [{**scene, field: ""}]} + ) + valid_resp = self._make_mock_response({"storyboard": [scene]}) + mock_client.models.generate_content.side_effect = [ + empty_resp, + valid_resp, + ] + self.mock_gcs.store.reset_mock() + + generate_storyboard.execute( + self.mock_gcs, + self.mock_params, + self.mock_images, + self.mock_user_prompt, + self.gemini_model, + self.gemini_model_location, + ) + + self.assertEqual(mock_client.models.generate_content.call_count, 2) + stored = json.loads(self.mock_gcs.store.call_args.args[0]) + self.assertEqual(len(stored["storyboard"]), 1) + + @patch("actions.generate_storyboard.genai.Client") + def test_execute_exhausting_retries_raises_runtime_error( + self, mock_genai_client_class + ): + mock_client = MagicMock() + mock_genai_client_class.return_value = mock_client + invalid_resp = self._make_mock_response({ + "storyboard": [{ + "image_id": "i_invalid", + "product_id": "p1", + "scene_name": "Scene 1", + "video_prompt": "Prompt 1", + }] + }) + mock_client.models.generate_content.return_value = invalid_resp + + with self.assertRaises(RuntimeError) as cm: + generate_storyboard.execute( + self.mock_gcs, + self.mock_params, + self.mock_images, + self.mock_user_prompt, + self.gemini_model, + self.gemini_model_location, + ) + + self.assertEqual(mock_client.models.generate_content.call_count, 4) + self.assertIn("Gemini repeatedly scripted invalid image/product combinations", str(cm.exception)) + + @patch("actions.generate_storyboard.genai.Client") + def test_execute_valid_ids_pass_through_unchanged( + self, mock_genai_client_class + ): + mock_client = MagicMock() + mock_genai_client_class.return_value = mock_client + valid_resp = self._make_mock_response({ + "storyboard": [{ + "image_id": "i1", + "product_id": "p1", + "scene_name": "Scene 1", + "video_prompt": "Prompt 1", + }] + }) + mock_client.models.generate_content.return_value = valid_resp + + generate_storyboard.execute( + self.mock_gcs, + self.mock_params, + self.mock_images, + self.mock_user_prompt, + self.gemini_model, + self.gemini_model_location, + ) + + self.mock_gcs.store.assert_called_once() + stored_json = json.loads(self.mock_gcs.store.call_args[0][0]) + self.assertEqual(stored_json["storyboard"][0]["product_id"], "p1") + self.assertEqual(stored_json["storyboard"][0]["image_id"], "i1") + + @patch("actions.generate_storyboard.genai.Client") + def test_execute_missing_product_description_yields_empty_string_not_none( + self, mock_genai_client_class + ): + mock_client = MagicMock() + mock_genai_client_class.return_value = mock_client + mock_client.models.generate_content.return_value = self._make_mock_response({ + "storyboard": [{ + "image_id": "i1", + "product_id": "p1", + "scene_name": "Scene 1", + "video_prompt": "Prompt 1", + }] + }) + + # Test with both omitted and explicitly None product_description + images = [ + { + Dimension.PRODUCT_ID.value: "p1", + Dimension.IMAGE_ID.value: "i1", + Key.FILE.value: "path/to/img1.jpg", + # product_description omitted + }, + { + Dimension.PRODUCT_ID.value: "p2", + Dimension.IMAGE_ID.value: "i2", + Key.FILE.value: "path/to/img2.jpg", + "product_description": None, + }, + ] + + generate_storyboard.execute( + self.mock_gcs, + self.mock_params, + images, + self.mock_user_prompt, + self.gemini_model, + self.gemini_model_location, + ) + + _, kwargs = mock_client.models.generate_content.call_args + prompt_parts = kwargs["contents"] + all_text = "".join(p.text for p in prompt_parts if getattr(p, "text", None)) + self.assertNotIn("'None'", all_text) + self.assertNotIn("Product description:", all_text) + def setUp(self): self.mock_gcs = MagicMock() self.mock_params = {Key.GCP_PROJECT.value: "test-project"} diff --git a/ui/src/app/services/remix-engine/remix-engine-mediated.spec.ts b/ui/src/app/services/remix-engine/remix-engine-mediated.spec.ts index 75ceea02..29637199 100644 --- a/ui/src/app/services/remix-engine/remix-engine-mediated.spec.ts +++ b/ui/src/app/services/remix-engine/remix-engine-mediated.spec.ts @@ -2187,6 +2187,67 @@ describe('RemixEngineService (mediated)', () => { }); }); + it('gives each scene its own reference image object', async () => { + const preview = { + path: 'thumbnail/product.jpg', + url: 'https://signed.example/thumbnail/product.jpg', + }; + const product = { + id: 1, + name: 'Product 1', + images: [ + { + path: 'input/product.jpg', + url: 'https://signed.example/input/product.jpg', + preview, + }, + ], + }; + vi.spyOn(service, 'startStoryboardWorkflow').mockResolvedValue( + of({executionId: 'storyboard-exec'}) as any, + ); + vi.spyOn(service, 'pollWorkflow').mockResolvedValue({ + sink: { + output: { + '0': { + storyboard: [{file: 'storyboard.json'}], + outpainted_images: [ + {product_id: 1, image_id: 1, file: 'input/product.jpg'}, + ], + }, + }, + }, + } as any); + const sceneJson = {product_id: 1, image_id: 1, video_prompt: 'P'}; + mediaServiceMock.getBlob.mockResolvedValue({ + text: () => + Promise.resolve( + JSON.stringify({ + storyboard: [ + {...sceneJson, scene_name: 'Scene 1'}, + {...sceneJson, scene_name: 'Scene 2'}, + ], + }), + ), + }); + mediaServiceMock.signUrl.mockResolvedValue( + 'https://signed.example/input/product.jpg', + ); + + const result = await service.generateStoryboard( + [product as any], + 'briefing', + 'none', + ); + + const [first, second] = result!; + expect(first.referenceImage).toEqual(second.referenceImage); + expect(first.referenceImage).not.toBe(second.referenceImage); + expect(first.referenceImage?.preview).not.toBe( + second.referenceImage?.preview, + ); + }); + it('does not carry a product preview when the workflow changes the path', async () => { const product = { id: 1, @@ -2289,6 +2350,53 @@ describe('RemixEngineService (mediated)', () => { ); expect(result?.[0].referenceImage?.path).toBe('out.jpg'); }); + + it('does not throw when model returns an unmatched product or image id', async () => { + vi.spyOn(service, 'startStoryboardWorkflow').mockResolvedValue( + of({executionId: 'storyboard-exec'}) as any, + ); + vi.spyOn(service, 'pollWorkflow').mockResolvedValue({ + sink: { + output: { + '0': { + storyboard: [{file: 'storyboard.json'}], + outpainted_images: [ + {product_id: 1, image_id: 1, file: 'out.jpg'}, + ], + }, + }, + }, + } as any); + mediaServiceMock.getBlob.mockResolvedValue({ + text: () => + Promise.resolve( + JSON.stringify({ + storyboard: [ + { + product_id: 999, + image_id: 888, + scene_name: 'Scene 1', + video_prompt: 'prompt', + }, + ], + }), + ), + }); + mediaServiceMock.signUrl.mockResolvedValue( + 'https://signed.example/out.jpg', + ); + + const result = await service.generateStoryboard( + [{id: 1, name: 'Product 1'}] as any, + 'briefing', + 'outpaint', + ); + + expect(result).toBeDefined(); + expect(result?.length).toBe(1); + expect(result?.[0].referenceImage).toBeUndefined(); + expect(matSnackBarMock.open).not.toHaveBeenCalled(); + }); }); describe('global config guard (E1)', () => { diff --git a/ui/src/app/services/remix-engine/remix-engine.ts b/ui/src/app/services/remix-engine/remix-engine.ts index b8631fae..3c011c65 100644 --- a/ui/src/app/services/remix-engine/remix-engine.ts +++ b/ui/src/app/services/remix-engine/remix-engine.ts @@ -2018,9 +2018,9 @@ export class RemixEngineService { return storyboardJson['storyboard'].map((s: StoryboardItem) => { const referenceImage = - productsToOutpaintedImages[s.product_id][s.image_id]; + productsToOutpaintedImages[s.product_id]?.[s.image_id]; - return { + const scene = { id: (sceneIdCounter++).toString(), type: 'generated', name: s.scene_name, @@ -2030,14 +2030,16 @@ export class RemixEngineService { numberOfCandidates: this.configService.globalConfig.value()!.numberOfCandidates, generateAudio: this.configService.globalConfig.value()!.generateAudio, - referenceImage: { - url: referenceImage.url, - path: referenceImage.path, + } as GeneratedScene; + if (referenceImage) { + scene.referenceImage = { + ...referenceImage, ...(referenceImage.preview - ? {preview: referenceImage.preview} + ? {preview: {...referenceImage.preview}} : {}), - }, - } as GeneratedScene; + }; + } + return scene; }); } catch (error) { if (error instanceof ProjectChangedError) { diff --git a/ui/src/app/setup/setup.spec.ts b/ui/src/app/setup/setup.spec.ts index 5ee3f4e5..150f82cb 100644 --- a/ui/src/app/setup/setup.spec.ts +++ b/ui/src/app/setup/setup.spec.ts @@ -26,20 +26,177 @@ import {ComponentFixture, TestBed} from '@angular/core/testing'; import {MatSelectHarness} from '@angular/material/select/testing'; import {MatSlideToggle} from '@angular/material/slide-toggle'; import {MatSlider} from '@angular/material/slider'; +import {MatDialog} from '@angular/material/dialog'; import {MatSnackBar} from '@angular/material/snack-bar'; import {By} from '@angular/platform-browser'; -import {provideRouter} from '@angular/router'; +import {ActivatedRoute, Router, provideRouter} from '@angular/router'; import {RouterTestingHarness} from '@angular/router/testing'; +import {of} from 'rxjs'; import {beforeEach, describe, expect, it, vi} from 'vitest'; import {routes} from '../app.routes'; import {ClientMediaService} from '../services/client-media/client-media'; -import {ConfigService, ProjectConfig} from '../services/config/config'; +import { + ConfigService, + GeneratedScene, + ProjectConfig, +} from '../services/config/config'; import {ImagePreviewService} from '../services/image-preview/image-preview'; import {ImageImportService} from '../services/image-import/image-import'; import {RemixEngineService} from '../services/remix-engine/remix-engine'; import {TemplatesService} from '../services/templates/templates'; import {Setup} from './setup'; +describe('Setup storyboard dispatch validation', () => { + let component: Setup; + let projectConfigSignal: WritableSignal; + let configMock: { + projectConfig: {value: WritableSignal}; + setupInputsLoaded: () => boolean; + updateProjectConfig: ReturnType; + }; + let dialogMock: {open: ReturnType}; + let remixMock: {generateCandidates: ReturnType}; + let snackBarMock: {open: ReturnType}; + let navigate: ReturnType; + let dialogResult: GeneratedScene[] | undefined; + + const validScene = (): GeneratedScene => ({ + id: 'scene-1', + type: 'generated', + name: 'Scene 1', + prompt: 'A product scene', + referenceImage: { + path: 'input/product.jpg', + url: 'https://signed.example/input/product.jpg', + }, + }); + + beforeEach(async () => { + projectConfigSignal = signal({ + id: 'project-1', + name: 'Test Project', + aspectRatio: '16:9', + resolution: '720p', + candidateDurationSeconds: 5, + generateAudio: false, + numberOfCandidates: 1, + model: 'veo-default', + inputConfig: {products: [], composition: '', style: '', audience: ''}, + storyboard: [], + audioTracks: [], + visualOverlays: [], + }); + configMock = { + projectConfig: {value: projectConfigSignal}, + setupInputsLoaded: () => true, + updateProjectConfig: vi.fn((partial: Partial) => + projectConfigSignal.update(current => ({...current, ...partial})), + ), + }; + dialogResult = undefined; + dialogMock = { + open: vi.fn(() => ({afterClosed: () => of(dialogResult)})), + }; + remixMock = {generateCandidates: vi.fn()}; + snackBarMock = {open: vi.fn()}; + navigate = vi.fn(); + + TestBed.configureTestingModule({ + imports: [Setup], + providers: [ + {provide: ConfigService, useValue: configMock}, + {provide: MatDialog, useValue: dialogMock}, + {provide: RemixEngineService, useValue: remixMock}, + {provide: MatSnackBar, useValue: snackBarMock}, + {provide: Router, useValue: {navigate}}, + { + provide: ActivatedRoute, + useValue: {snapshot: {paramMap: {get: () => 'project-1'}}}, + }, + {provide: ClientMediaService, useValue: {convertImage: vi.fn()}}, + { + provide: ImageImportService, + useValue: { + importText: vi.fn(), + imageFilesFromDataTransfer: vi.fn().mockReturnValue([]), + imageUrlFromDataTransfer: vi.fn().mockReturnValue(null), + isEditableTarget: vi.fn().mockReturnValue(false), + }, + }, + { + provide: ImagePreviewService, + useValue: {create: vi.fn()}, + }, + { + provide: TemplatesService, + useValue: {templates: {value: () => undefined}}, + }, + ], + }); + TestBed.overrideComponent(Setup, {set: {template: ''}}); + TestBed.overrideProvider(MatDialog, {useValue: dialogMock}); + TestBed.overrideProvider(MatSnackBar, {useValue: snackBarMock}); + await TestBed.compileComponents(); + + const fixture = TestBed.createComponent(Setup); + component = fixture.componentInstance; + fixture.detectChanges(); + component.creationMode.set('ai'); + }); + + it('rejects a valid scene followed by a legacy invalid reference before dispatch', () => { + dialogResult = [ + validScene(), + { + ...validScene(), + id: 'scene-2', + referenceImage: {url: undefined, path: undefined}, + } as unknown as GeneratedScene, + ]; + + component.generateStoryboard(); + + expect(snackBarMock.open).toHaveBeenCalledWith( + 'Cannot generate videos: one or more storyboard scenes are missing a product image.', + 'Dismiss', + {panelClass: ['error-snackbar']}, + ); + expect(configMock.updateProjectConfig).not.toHaveBeenCalled(); + expect(remixMock.generateCandidates).not.toHaveBeenCalled(); + expect(navigate).not.toHaveBeenCalled(); + }); + + it('rejects a scene with no reference mapping before dispatch', () => { + dialogResult = [{...validScene(), referenceImage: undefined}]; + + component.generateStoryboard(); + + expect(snackBarMock.open).toHaveBeenCalled(); + expect(configMock.updateProjectConfig).not.toHaveBeenCalled(); + expect(remixMock.generateCandidates).not.toHaveBeenCalled(); + expect(navigate).not.toHaveBeenCalled(); + }); + + it('stores and dispatches a fresh valid storyboard result', () => { + const scene = validScene(); + dialogResult = [scene]; + + component.generateStoryboard(); + + expect(configMock.updateProjectConfig).toHaveBeenCalledWith({ + storyboard: [scene], + }); + expect(remixMock.generateCandidates).toHaveBeenCalledWith(scene, { + durationSeconds: 5, + generateAudio: false, + model: 'veo-default', + resolution: '720p', + }); + expect(navigate).toHaveBeenCalledWith(['project-1', 'storyboard']); + expect(snackBarMock.open).not.toHaveBeenCalled(); + }); +}); + describe('Setup', () => { let component: Setup; diff --git a/ui/src/app/setup/setup.ts b/ui/src/app/setup/setup.ts index 43db63fa..da3c27b6 100644 --- a/ui/src/app/setup/setup.ts +++ b/ui/src/app/setup/setup.ts @@ -811,6 +811,20 @@ export class Setup { .afterClosed() .subscribe((result: GeneratedScene[] | undefined) => { if (result) { + if ( + result.some( + scene => + typeof scene.referenceImage?.path !== 'string' || + scene.referenceImage.path.trim() === '', + ) + ) { + this.snackBar.open( + 'Cannot generate videos: one or more storyboard scenes are missing a product image.', + 'Dismiss', + {panelClass: ['error-snackbar']}, + ); + return; + } this.config.updateProjectConfig({storyboard: [...result]}); for (const scene of result) { void this.remixEngineService.generateCandidates(scene, {