Merge branch 'feature/imagen-upscale' of github.com:GoogleCloudPlatform/gcc-creative-studio into feature/gallery-enhancements-and-fixes

Author: MauroCominotti

backend/alembic/versions/0bd50a4bf20c_add_original_gcs_uris.py

@@ -0,0 +1,37 @@
+"""add_original_gcs_uris
+
+Revision ID: 0bd50a4bf20c
+Revises: 9393a3d298c6
+Create Date: 2026-01-29 12:53:26.493393
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '0bd50a4bf20c'
+down_revision: Union[str, None] = '9393a3d298c6'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+    
+    media_columns = [c['name'] for c in inspector.get_columns('media_items')]
+    if 'original_gcs_uris' not in media_columns:
+        op.add_column('media_items', sa.Column('original_gcs_uris', sa.ARRAY(sa.String()), nullable=True))
+        
+    source_columns = [c['name'] for c in inspector.get_columns('source_assets')]
+    if 'original_gcs_uri' not in source_columns:
+        op.add_column('source_assets', sa.Column('original_gcs_uri', sa.String(), nullable=True))
+
+
+def downgrade() -> None:
+    op.drop_column('source_assets', 'original_gcs_uri')
+    op.drop_column('media_items', 'original_gcs_uris')
+

backend/main.py

@@ -18,6 +18,8 @@
 setup_logging()
 
 import logging
+import os
+from dotenv import load_dotenv
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from os import getenv
@@ -54,15 +56,15 @@
 
 # Get a logger instance for use in this file. It will inherit the root setup.
 logger = logging.getLogger(__name__)
-
+load_dotenv()
 
 def configure_cors(app):
     """Configures CORS middleware based on the environment."""
-    environment = getenv("ENVIRONMENT")
+    environment = os.getenv("ENVIRONMENT")
     allowed_origins = []
 
     if environment == "production":
-        frontend_url = getenv("FRONTEND_URL")
+        frontend_url = os.getenv("FRONTEND_URL")
         if not frontend_url:
             raise ValueError(
                 "FRONTEND_URL environment variable not set in production"

backend/src/brand_guidelines/brand_guideline_controller.py

@@ -36,7 +36,7 @@
     GenerateUploadUrlResponseDto,
 )
 from src.workspaces.repository.workspace_repository import WorkspaceRepository
-from src.workspaces.workspace_auth_guard import workspace_auth_service
+from src.workspaces.workspace_auth_guard import WorkspaceAuth
 from src.users.user_model import UserModel, UserRoleEnum
 
 MAX_UPLOAD_SIZE_BYTES = 500 * 1024 * 1024  # 500 MB
@@ -62,18 +62,17 @@ async def generate_upload_url(
     request_dto: GenerateUploadUrlDto,
     current_user: UserModel = Depends(get_current_user),
     service: BrandGuidelineService = Depends(),
-    workspace_repo: WorkspaceRepository = Depends(),
+    workspace_auth: WorkspaceAuth = Depends(),
 ):
     """
     Generates a secure, short-lived URL that the client can use to upload a
     brand guideline PDF directly to Google Cloud Storage.
     """
     # If a workspace ID is provided, ensure the user has access to it.
     if request_dto.workspace_id:
-        await workspace_auth_service.authorize(
+        await workspace_auth.authorize(
             workspace_id=request_dto.workspace_id,
             user=current_user,
-            workspace_repo=workspace_repo,
         )
 
     if not request_dto.content_type == "application/pdf":

backend/src/common/base_dto.py

@@ -43,6 +43,7 @@ class GenerationModelEnum(str, Enum):
     """Enum representing the available generation models."""
 
     # Image-Specific Models
+    IMAGEN_4_UPSCALE_PREVIEW = "imagen-4.0-upscale-preview"
     IMAGEN_4_001 = "imagen-4.0-generate-001"
     IMAGEN_4_ULTRA = "imagen-4.0-ultra-generate-001"
     IMAGEN_4_ULTRA_PREVIEW = "imagen-4.0-ultra-generate-preview-06-06"

backend/src/common/schema/media_item_model.py

@@ -165,6 +165,7 @@ class MediaItem(Base):
     source_media_items: Mapped[Optional[List[dict]]] = mapped_column(JSONB, nullable=True)
 
     gcs_uris: Mapped[List[str]] = mapped_column(ARRAY(String), default=[])
+    original_gcs_uris: Mapped[Optional[List[str]]] = mapped_column(ARRAY(String), nullable=True, default=[])
 
     # Video specific
     duration_seconds: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
@@ -253,6 +254,14 @@ class MediaItemModel(BaseDocument):
             description="A list of public URLs for the media to be displayed (e.g., video or image).",
         ),
     ]
+    original_gcs_uris: Annotated[
+        Optional[List[str]],
+        Field(
+            default=None,
+            min_length=0, 
+            description="A list of public URLs (original / non-upscaled) for the media to be displayed (e.g., video or image).",
+        ),
+    ]
 
     # Video specific
     duration_seconds: Optional[float] = None

backend/src/common/storage_service.py

@@ -220,8 +220,28 @@ def store_to_gcs(
                 f"Blob '{destination_blob_name}' not found in bucket '{self.bucket_name}'."
             )
             return None
-        except exceptions.GoogleAPICallError as e:
-            logger.error(
-                f"Failed to download '{destination_blob_name}' from GCS: {e}"
-            )
-            return None
+
+    def check_file_exists(self, gcs_uri: str) -> bool:
+        """
+        Checks if a file exists in GCS.
+
+        Args:
+            gcs_uri: The full GCS URI (e.g., "gs://bucket-name/path/to/blob").
+
+        Returns:
+            True if the file exists, False otherwise.
+        """
+        try:
+            if not gcs_uri:
+                return False
+            if not gcs_uri.startswith("gs://"):
+                logger.warning(f"Invalid GCS URI format: {gcs_uri}")
+                return False
+
+            bucket_name, blob_name = gcs_uri.replace("gs://", "").split("/", 1)
+            bucket = self.client.bucket(bucket_name)
+            blob = bucket.blob(blob_name)
+            return blob.exists()
+        except Exception as e:
+            logger.error(f"Error checking file existence for {gcs_uri}: {e}")
+            return False

backend/src/config/config_service.py

@@ -155,4 +155,4 @@ def IMAGE_BUCKET(self) -> str:
 
 
 # Create a single, cached instance of the settings to be used throughout the app.
-config_service = ConfigService()
+config_service = ConfigService()

backend/src/galleries/dto/gallery_response_dto.py

@@ -49,6 +49,7 @@ class MediaItemResponse(MediaItemModel):
     """
 
     presigned_urls: List[str] = []
+    original_presigned_urls: Optional[List[str]] = None
     presigned_thumbnail_urls: Optional[List[str]] = []
     enriched_source_assets: Optional[List[SourceAssetLinkResponse]] = None
     enriched_source_media_items: Optional[List[SourceMediaItemLinkResponse]] = (

backend/src/galleries/gallery_controller.py

@@ -21,7 +21,7 @@
 from src.galleries.gallery_service import GalleryService
 from src.users.user_model import UserModel, UserRoleEnum
 from src.workspaces.repository.workspace_repository import WorkspaceRepository
-from src.workspaces.workspace_auth_guard import workspace_auth_service
+from src.workspaces.workspace_auth_guard import WorkspaceAuth
 
 router = APIRouter(
     prefix="/api/gallery",
@@ -48,7 +48,7 @@ async def search_gallery_items(
     search_dto: GallerySearchDto,
     current_user: UserModel = Depends(get_current_user),
     service: GalleryService = Depends(),
-    workspace_repo: WorkspaceRepository = Depends(),
+    workspace_auth: WorkspaceAuth = Depends(),
 ):
     """
     Performs a paginated search for media items within a specific workspace.
@@ -58,10 +58,9 @@ async def search_gallery_items(
     """
     # This dependency call acts as a gatekeeper. If the user is not authorized
     # for the workspace_id inside search_dto, it will raise an exception.
-    await workspace_auth_service.authorize(
+    await workspace_auth.authorize(
         workspace_id=search_dto.workspace_id,
         user=current_user,
-        workspace_repo=workspace_repo,
     )
 
     return await service.get_paginated_gallery(

backend/src/galleries/gallery_service.py

@@ -32,15 +32,22 @@
     SourceAssetLinkResponse,
     SourceMediaItemLinkResponse,
 )
+from src.common.base_dto import (
+    AspectRatioEnum,
+    GenerationModelEnum,
+    MimeTypeEnum,
+)
 from src.galleries.dto.gallery_search_dto import GallerySearchDto
+from src.images.dto.upscale_imagen_dto import UpscaleImagenDto
+from src.images.imagen_service import ImagenService
 from src.images.repository.media_item_repository import MediaRepository
 from src.source_assets.repository.source_asset_repository import (
     SourceAssetRepository,
 )
 from src.users.user_model import UserModel, UserRoleEnum
 from src.workspaces.repository.workspace_repository import WorkspaceRepository
 from src.workspaces.schema.workspace_model import WorkspaceScopeEnum
-from src.workspaces.workspace_auth_guard import workspace_auth_service
+from src.workspaces.workspace_auth_guard import WorkspaceAuth
 
 logger = logging.getLogger(__name__)
 
@@ -56,12 +63,16 @@ def __init__(
         source_asset_repo: SourceAssetRepository = Depends(),
         workspace_repo: WorkspaceRepository = Depends(),
         iam_signer_credentials: IamSignerCredentials = Depends(),
+        workspace_auth: WorkspaceAuth = Depends(),
+        imagen_service: ImagenService = Depends(),
     ):
         """Initializes the service with its dependencies."""
         self.media_repo = media_repo
         self.source_asset_repo = source_asset_repo
         self.workspace_repo = workspace_repo
         self.iam_signer_credentials = iam_signer_credentials
+        self.workspace_auth = workspace_auth
+        self.imagen_service = imagen_service 
 
     async def _enrich_source_asset_link(
         self, link: SourceAssetLink
@@ -70,7 +81,7 @@ async def _enrich_source_asset_link(
         Fetches the source asset document and generates a presigned URL for it.
         """
         asset_doc = await self.source_asset_repo.get_by_id(link.asset_id)
-
+        
         if not asset_doc:
             return None
 
@@ -171,6 +182,16 @@ async def _create_gallery_response(
             if uri
         ]
 
+        # 1.5 Create tasks for original media URLs
+        all_original_gcs_uris = item.original_gcs_uris or []
+        original_url_tasks = [
+            asyncio.to_thread(
+                self.iam_signer_credentials.generate_presigned_url, uri
+            )
+            for uri in all_original_gcs_uris
+            if uri
+        ]
+
         # 2. Create tasks for thumbnail URLs
         thumbnail_tasks = [
             asyncio.to_thread(
@@ -199,11 +220,13 @@ async def _create_gallery_response(
         # 5. Gather all results concurrently
         (
             presigned_urls,
+            original_presigned_urls,
             presigned_thumbnail_urls,
             enriched_source_assets_with_nones,
             enriched_source_media_items_with_nones,
         ) = await asyncio.gather(
             asyncio.gather(*main_url_tasks),
+            asyncio.gather(*original_url_tasks),
             asyncio.gather(*thumbnail_tasks),
             asyncio.gather(*source_asset_tasks),
             asyncio.gather(*source_media_item_tasks),
@@ -220,6 +243,7 @@ async def _create_gallery_response(
         return MediaItemResponse(
             **item.model_dump(exclude={"source_assets"}),
             presigned_urls=presigned_urls,
+            original_presigned_urls=original_presigned_urls,
             presigned_thumbnail_urls=presigned_thumbnail_urls,
             enriched_source_assets=enriched_source_assets or None,
             enriched_source_media_items=enriched_source_media_items or None,
@@ -259,7 +283,7 @@ async def get_paginated_gallery(
         )
 
     async def get_media_by_id(
-        self, item_id: int, current_user: UserModel
+        self, item_id: str, current_user: UserModel
     ) -> Optional[MediaItemResponse]:
         """
         Retrieves a single media item, performs an authorization check,
@@ -282,10 +306,13 @@ async def get_media_by_id(
             )
 
         # Use the centralized authorization logic
-        await workspace_auth_service.authorize(
+        await self.workspace_auth.authorize(
             workspace_id=item.workspace_id,
             user=current_user,
-            workspace_repo=self.workspace_repo,
         )
 
         return await self._create_gallery_response(item)
+    
+
+
+