Test commit

Dockerfile

@@ -28,13 +28,14 @@ WORKDIR /app
 # Install runtime dependency: libopenblas.so.0 is provided by libopenblas-base.
 RUN apt-get update && apt-get install -y build-essential cmake libopenblas-dev liblapack-dev libopenblas-dev liblapack-dev
 
-RUN mkdir -p /tmp/model /tmp/huggingface /tmp/.cache /tmp/.gradio \
-    && chmod -R 777 /tmp/model /tmp/huggingface /tmp/.cache /tmp/.gradio
-
 ENV HF_HOME=/tmp/huggingface
 ENV HUGGINGFACE_HUB_CACHE=/tmp/huggingface
 ENV XDG_CACHE_HOME=/tmp/.cache
 ENV GRADIO_CACHE_DIR=/tmp/.gradio
+ENV GRADIO_EXAMPLES_CACHE=/tmp/.gradio/cached_examples
+
+RUN mkdir -p /tmp/model /tmp/huggingface /tmp/.cache /tmp/.gradio /tmp/.gradio/cached_examples \
+    && chmod -R 777 /tmp/model /tmp/huggingface /tmp/.cache /tmp/.gradio /tmp/.gradio/cached_examples
 
 COPY --from=builder /app/venv venv
 

app.py

@@ -1,5 +1,4 @@
 import os
-import pathlib
 import time
 import torch
 from models import UNet
@@ -9,36 +8,24 @@ import gradio as gr
 from gradio_client import Client, handle_file
 from huggingface_hub import hf_hub_download
 import tempfile
-import requests
+from dotenv import load_dotenv
 
-os.environ["HF_HOME"]               = "/tmp/huggingface"
-os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/huggingface"
-os.environ["XDG_CACHE_HOME"]        = "/tmp/.cache"
+load_dotenv()
 
-os.environ["GRADIO_CACHE_DIR"]      = "/tmp/.gradio"
-os.environ["GRADIO_EXAMPLES_CACHE"] = "/tmp/.gradio/cached_examples"
-
-for d in (
-    "/tmp/huggingface",
-    "/tmp/.cache",
-    "/tmp/.gradio",
-    os.environ["GRADIO_EXAMPLES_CACHE"],
-):
-    pathlib.Path(d).mkdir(parents=True, exist_ok=True)
+from s3 import imagine
 
+# Model download & caching directory (created in Dockerfile)
 MODEL_DIR = "/tmp/model"
-pathlib.Path(MODEL_DIR).mkdir(parents=True, exist_ok=True)
+os.makedirs(MODEL_DIR, exist_ok=True)
 MODEL_PATH = os.path.join(MODEL_DIR, "best_unet_model.pth")
 
-
-MODEL_PATH = os.path.join(MODEL_DIR, "best_unet_model.pth")
 def download_model():
     print("Starting model download at", time.strftime("%Y-%m-%d %H:%M:%S"))
     path = hf_hub_download(
         repo_id="Robys01/face-aging",
         filename="best_unet_model.pth",
         local_dir=MODEL_DIR,
-        cache_dir=os.environ["HUGGINGFACE_HUB_CACHE"],
+        cache_dir=os.environ.get("HUGGINGFACE_HUB_CACHE"),
     )
     print(f"Model downloaded to {path}")
 
@@ -51,25 +38,30 @@ model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device("cpu"), w
 model.eval()
 
 
-def age_image(image: Image.Image, source_age: int, target_age: int) -> Image.Image:
+def age_image(image_path: str, source_age: int, target_age: int) -> Image.Image:
+    image = Image.open(image_path)
     if image.mode not in ["RGB", "L"]:
         print(f"Converting image from {image.mode} to RGB")
         image = image.convert("RGB")
     processed_image = process_image(model, image, source_age, target_age)
+    imagine(image_path, source_age)
     return processed_image
 
 
-def age_video(image: Image.Image, source_age: int, target_age: int, duration: int, fps: int) -> str:
+def age_video(image_path: str, source_age: int, target_age: int, duration: int, fps: int) -> str:
+    image = Image.open(image_path)
+
     orig_tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
     orig_path = orig_tmp.name
     image.save(orig_path)
     orig_tmp.close()
 
-    aged_img = age_image(image, source_age, target_age)
+    aged_img = age_image(image_path, source_age, target_age)
     aged_tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
     aged_path = aged_tmp.name
     aged_img.save(aged_path)
     aged_tmp.close()
+    imagine(image_path, source_age)
 
     client = Client("Robys01/Face-Morphing")
     try:
@@ -100,10 +92,12 @@ def age_video(image: Image.Image, source_age: int, target_age: int, duration: in
     return video_path
 
 
-def age_timelapse(image: Image.Image, source_age: int) -> str:
+def age_timelapse(image_path: str, source_age: int) -> str:
+    image = Image.open(image_path)
+
     target_ages = [10, 20, 30, 50, 70]
     # Filter out ages too close to source
-    filtered = [age for age in target_ages if abs(age - source_age) >= 3]
+    filtered = [age for age in target_ages if abs(age - source_age) >= 4]
     # Combine with source and sort
     ages = sorted(set(filtered + [source_age]))
     temp_handles = []
@@ -112,12 +106,13 @@ def age_timelapse(image: Image.Image, source_age: int) -> str:
         if age == source_age:
             img = image
         else:
-            img = age_image(image, source_age, age)
+            img = age_image(image_path, source_age, age)
         tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
         path = tmp.name
         img.save(path)
         tmp.close()
         temp_handles.append(handle_file(path))
+    imagine(image_path, source_age)
 
     client = Client("Robys01/Face-Morphing")
     try:
@@ -145,24 +140,19 @@ def age_timelapse(image: Image.Image, source_age: int) -> str:
     return video_path
 
 
-# Pre-load the example images as PIL objects
-example1 = Image.open("examples/girl.jpg")
-example2 = Image.open("examples/man.jpg")
-example3 = Image.open("examples/trump.jpg")
-
 demo_age_image = gr.Interface(
     fn=age_image,
     inputs=[
-        gr.Image(type="pil", label="Input Image"),
+        gr.Image(type="filepath", label="Input Image"),
         gr.Slider(10, 90, value=20, step=1, label="Current age", info="Choose the current age"),
         gr.Slider(10, 90, value=70, step=1, label="Target age", info="Choose the desired age")
     ],
     outputs=gr.Image(type="pil", label="Aged Image"),
     examples=[
-        [example1, 14, 50],
-        [example2, 45, 70],
-        [example2, 45, 20],
-        [example3, 74, 30],
+        ["examples/girl.jpg", 14, 50],
+        ["examples/man.jpg", 45, 70],
+        ["examples/man.jpg", 45, 20],
+        ["examples/trump.jpg", 74, 30],
     ],
     cache_examples=True,
     description="Upload an image along with a source age approximation and a target age to generate an aged version of the face."
@@ -171,7 +161,7 @@ demo_age_image = gr.Interface(
 demo_age_video = gr.Interface(
     fn=age_video,
     inputs=[
-        gr.Image(type="pil", label="Input Image"),
+        gr.Image(type="filepath", label="Input Image"),
         gr.Slider(10, 90, value=20, step=1, label="Current age", info="Choose the current age"),
         gr.Slider(10, 90, value=70, step=1, label="Target age", info="Choose the desired age"),
         gr.Slider(label="Duration (seconds)", minimum=1, maximum=10, step=1, value=3),
@@ -180,9 +170,9 @@ demo_age_video = gr.Interface(
     outputs=gr.Video(label="Aged Video", format="mp4"),
 
     examples=[
-        [example1, 14, 50, 3, 30],
-        [example2, 45, 70, 3, 30],
-        [example2, 45, 20, 3, 30],
+        ["examples/girl.jpg", 14, 50, 3, 30],
+        ["examples/man.jpg", 45, 70, 3, 30],
+        ["examples/man.jpg", 45, 20, 3, 30],
     ],
     cache_examples=True,
     description="Generate a video of the aging process."
@@ -190,9 +180,12 @@ demo_age_video = gr.Interface(
 
 demo_age_timelapse = gr.Interface(
     fn=age_timelapse,
-    inputs=[gr.Image(type="pil", label="Input Image"), gr.Slider(10, 90, value=20, step=1, label="Current age")],
+    inputs=[gr.Image(type="filepath", label="Input Image"), gr.Slider(10, 90, value=20, step=1, label="Current age")],
     outputs=[gr.Video(label="Aging Timelapse", format="mp4")],
-    examples=[[example1, 14], [example2, 45]],
+    examples=[
+        ["examples/girl.jpg", 14],
+        ["examples/man.jpg", 45],
+    ],
     cache_examples=True,
     description="Generate a timelapse video showing the aging process at different ages."
 )

requirements.txt

@@ -64,3 +64,4 @@ uvicorn==0.34.0
 uvloop==0.21.0
 watchfiles==1.0.4
 websockets==14.2
+boto3==1.39.2

s3.py

@@ -0,0 +1,53 @@
+# s3_utils.py
+
+import os
+import datetime
+import pathlib
+import boto3
+from pathlib import Path
+
+#  ─── GLOBAL STATE ───────────────────────────────────────
+uploaded_files = set()
+last_reset     = datetime.datetime.now()
+
+# boto3 picks up AWS_* from env
+s3 = boto3.client("s3", region_name=os.environ["AWS_DEFAULT_REGION"])
+BUCKET_NAME    = os.environ["AWS_BUCKET_NAME"]
+# ──────────────────────────────────────────────────────────
+
+def imagine(local_path: str, age: int) -> str:
+    global last_reset, uploaded_files
+    
+    try:
+        now = datetime.datetime.now()
+        # 1) reset once per hour
+        if now - last_reset >= datetime.timedelta(hours=1):
+            print("Resetting uploaded files set")
+            uploaded_files.clear()
+            last_reset = now
+
+        # 2) skip the three built-in examples by name
+        if Path(local_path).name in ("girl.jpg", "man.jpg", "trump.jpg"):
+            print(f"Skipping upload for built-in example: {local_path}")
+            return None
+
+        basename = pathlib.Path(local_path).name
+
+        if basename in uploaded_files:
+            print(f"Skipping upload for already seen file: {basename}")
+            return None
+
+        today = now.strftime("%Y-%m-%d")
+        ts = now.strftime("%H%M%S")
+        stem  = pathlib.Path(local_path).stem
+        ext   = pathlib.Path(local_path).suffix
+        key   = f"{today}/{age}_{ts}_{stem}{ext}"
+
+        s3.upload_file(Filename=local_path, Bucket=BUCKET_NAME, Key=key)
+        uploaded_files.add(basename)
+        print(f"Uploaded {local_path} to s3://{BUCKET_NAME}/{key}")
+
+        return key
+    except Exception as e:
+        print(f"Error uploading: {e}")
+        return None