Switch to Gemini + update requirements

app.py

@@ -1,7 +1,6 @@
 import os
 import time
 from datetime import datetime
-from typing import List
 
 from fastapi import FastAPI, UploadFile, Form
 from fastapi.responses import FileResponse, JSONResponse
@@ -14,13 +13,9 @@ import uvicorn
 
 from dotenv import load_dotenv
 
-# Load environment variables from .env
+# Load environment variables
 load_dotenv()
-
-# Confirm API key is loaded
-print("Gemini API Key Loaded:", os.getenv("GOOGLE_API_KEY") is not None)
-
-
+print("Gemini API Key Loaded:", os.getenv("GEMINI_API_KEY") is not None)
 
 # ---------------------------
 # CONFIG
@@ -34,7 +29,7 @@ LIFETIME = 24 * 60 * 60  # 24 hours
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(RESULTS_DIR, exist_ok=True)
 
-# Gemini API key
+# Gemini setup
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 model = genai.GenerativeModel("gemini-1.5-flash")
 
@@ -53,28 +48,54 @@ def check_size(filepath):
         os.remove(filepath)
         raise ValueError(f"File too large! Max {MAX_SIZE_MB}MB allowed.")
 
-def replace_background(input_path, bg_choice):
-    """Replace background with selected file"""
-    check_size(input_path)
-    input_img = Image.open(input_path).convert("RGBA")
-    fg = remove(input_img)
+def process_image(input_img, bg_choice, bg_upload=None, brand_color="#FFFFFF", logo_upload=None, logo_opacity=80, logo_position="bottom-right"):
+    """Main image processor with custom bg, brand color, and logo support"""
+    if input_img is None:
+        return []
+
+    temp_path = os.path.join(UPLOAD_DIR, f"upload_{int(time.time())}.png")
+    input_img.save(temp_path)
+
+    fg = remove(Image.open(temp_path).convert("RGBA"))
 
-    bg_path = os.path.join(BG_DIR, bg_choice)
-    bg = Image.open(bg_path).convert("RGBA").resize(fg.size)
+    # --- Background selection ---
+    if bg_choice and bg_choice in os.listdir(BG_DIR):
+        bg = Image.open(os.path.join(BG_DIR, bg_choice)).convert("RGBA").resize(fg.size)
+    elif bg_choice == "Custom Upload" and bg_upload is not None:
+        bg = bg_upload.convert("RGBA").resize(fg.size)
+    elif bg_choice == "Solid Brand Color" and brand_color:
+        bg = Image.new("RGBA", fg.size, brand_color)
+    else:
+        bg = Image.new("RGBA", fg.size, (255, 255, 255, 255))
 
     result = Image.alpha_composite(bg, fg)
+
+    # --- Logo overlay ---
+    if logo_upload is not None:
+        logo = logo_upload.convert("RGBA")
+        logo_size = int(result.width * 0.15)
+        logo.thumbnail((logo_size, logo_size))
+        alpha = logo.split()[3].point(lambda p: int(p * (logo_opacity / 100)))
+        logo.putalpha(alpha)
+
+        if logo_position == "top-left":
+            pos = (20, 20)
+        elif logo_position == "top-right":
+            pos = (result.width - logo.width - 20, 20)
+        elif logo_position == "bottom-left":
+            pos = (20, result.height - logo.height - 20)
+        elif logo_position == "center":
+            pos = ((result.width - logo.width) // 2, (result.height - logo.height) // 2)
+        else:  # default bottom-right
+            pos = (result.width - logo.width - 20, result.height - logo.height - 20)
+
+        result.alpha_composite(logo, pos)
+
     timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
     result_path = os.path.join(RESULTS_DIR, f"result_{timestamp}.png")
     result.save(result_path)
     cleanup_old_files(RESULTS_DIR)
-    return result_path
 
-def process_image(input_img, bg_choice):
-    if input_img is None:
-        return []
-    temp_path = os.path.join(UPLOAD_DIR, f"upload_{int(time.time())}.png")
-    input_img.save(temp_path)
-    result_path = replace_background(temp_path, bg_choice)
     return [result_path]
 
 def generate_caption(prompt="Promote my product"):
@@ -94,10 +115,9 @@ def generate_caption(prompt="Promote my product"):
 # ---------------------------
 app = FastAPI(title="SnapLift API")
 
-# Allow CORS for mobile app access
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Change later for security
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -109,9 +129,8 @@ async def process_image_api(file: UploadFile, bg_choice: str = Form(...)):
         input_path = os.path.join(UPLOAD_DIR, file.filename)
         with open(input_path, "wb") as f:
             f.write(await file.read())
-
-        result_path = replace_background(input_path, bg_choice)
-        return FileResponse(result_path)
+        result_path = process_image(Image.open(input_path), bg_choice)
+        return FileResponse(result_path[0])
     except Exception as e:
         return JSONResponse(content={"error": str(e)}, status_code=400)
 
@@ -127,56 +146,50 @@ with gr.Blocks(css="footer {display:none !important}") as demo:
     gr.Markdown("# ✨ SnapLift – AI Social Media Booster")
     gr.Markdown("Upload your product photo, replace background, and auto-generate marketing captions + hashtags!")
 
-    # Image Editor Tab
-
-        # Image Editor Tab
+    # --- Image Editor ---
     with gr.Tab("📸 Image Editor"):
         with gr.Row():
-            input_img = gr.Image(type="pil", label="Upload Photo")
+            input_img = gr.Image(type="pil", label="Upload Main Photo")
 
             with gr.Column():
                 bg_choices = gr.Dropdown(
-                    choices=os.listdir(BG_DIR),
+                    choices=os.listdir(BG_DIR) + ["Custom Upload", "Solid Brand Color"],
                     value=os.listdir(BG_DIR)[0] if os.listdir(BG_DIR) else None,
                     label="Choose Background"
                 )
-                bg_preview = gr.Image(
-                    label="Background Preview",
-                    type="filepath"
-                )
+                bg_preview = gr.Image(label="Background Preview", type="filepath")
+                bg_upload = gr.Image(type="pil", label="Upload Custom Background")
+                brand_color = gr.ColorPicker(label="Pick Brand Colour", value="#FFFFFF")
+
+        logo_upload = gr.Image(type="pil", label="Upload Brand Logo (Optional)")
+        logo_opacity = gr.Slider(minimum=0, maximum=100, value=80, step=5, label="Logo Transparency (%)")
+        logo_position = gr.Radio(
+            choices=["top-left", "top-right", "bottom-left", "bottom-right", "center"],
+            value="bottom-right",
+            label="Logo Position"
+        )
 
-        # Update preview whenever background changes
         def update_preview(bg_choice):
-            return os.path.join(BG_DIR, bg_choice) if bg_choice else None
+            return os.path.join(BG_DIR, bg_choice) if bg_choice in os.listdir(BG_DIR) else None
 
         bg_choices.change(fn=update_preview, inputs=bg_choices, outputs=bg_preview)
 
         btn = gr.Button("✨ Generate New Photo")
         output_imgs = gr.Gallery(label="Generated Image", elem_id="gallery", columns=1, rows=1)
-        btn.click(fn=process_image, inputs=[input_img, bg_choices], outputs=output_imgs)
-
-    
-    '''
-    with gr.Tab("📸 Image Editor"):
-        with gr.Row():
-            input_img = gr.Image(type="pil", label="Upload Photo")
-            bg_choices = gr.Dropdown(
-                choices=os.listdir(BG_DIR),
-                value=os.listdir(BG_DIR)[0] if os.listdir(BG_DIR) else None,
-                label="Choose Background"
-            )
-        btn = gr.Button("✨ Generate New Photo")
-        output_imgs = gr.Gallery(label="Generated Image", elem_id="gallery", columns=1, rows=1)
-        btn.click(fn=process_image, inputs=[input_img, bg_choices], outputs=output_imgs)
+        btn.click(
+            fn=process_image,
+            inputs=[input_img, bg_choices, bg_upload, brand_color, logo_upload, logo_opacity, logo_position],
+            outputs=output_imgs
+        )
 
-    # Caption Generator Tab
+    # --- Caption Generator ---
     with gr.Tab("✍️ Caption Generator"):
         prompt = gr.Textbox(label="Enter product/promotion text", value="Promote my skincare product")
         btn2 = gr.Button("💡 Suggest Captions + Hashtags")
         caption_box = gr.Textbox(label="Suggested Posts (multi-platform)", lines=12)
         btn2.click(fn=generate_caption, inputs=[prompt], outputs=[caption_box])
 
-    '''
+# ---------------------------
 # START SERVER
 # ---------------------------
 if __name__ == "__main__":
@@ -187,7 +200,3 @@ if __name__ == "__main__":
 
     threading.Thread(target=run_gradio).start()
     uvicorn.run(app, host="0.0.0.0", port=8000)
-
-threading.Thread(target=start_fastapi, daemon=True).start()
-
-demo.launch(server_name="0.0.0.0", server_port=7860)