Update main.py

main.py

@@ -1,105 +1,4 @@
-# Based on: https://github.com/jantic/DeOldify
-import os, re, time
+import os
 
-os.environ["TORCH_HOME"] = os.path.join(os.getcwd(), ".cache")
-os.environ["XDG_CACHE_HOME"] = os.path.join(os.getcwd(), ".cache")
 
-from fastapi import FastAPI, File, UploadFile,Form
-from fastapi.responses import FileResponse, StreamingResponse
-from fastapi.staticfiles import StaticFiles
-import PIL
-import cv2
-import numpy as np
-import uuid
-from zipfile import ZipFile, ZIP_DEFLATED
-from io import BytesIO
-from random import randint
-from datetime import datetime
-
-from src.deoldify import device
-from src.deoldify.device_id import DeviceId
-from src.deoldify.visualize import *
-from src.app_utils import get_model_bin
-
-app = FastAPI()
-
-device.set(device=DeviceId.CPU)
-
-
-
-def load_model(model_dir, option):
-    if option.lower() == 'artistic':
-        model_url = 'https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth'
-        get_model_bin(model_url, os.path.join(model_dir, "ColorizeArtistic_gen.pth"))
-        colorizer = get_image_colorizer(artistic=True)
-    elif option.lower() == 'stable':
-        model_url = "https://www.dropbox.com/s/usf7uifrctqw9rl/ColorizeStable_gen.pth?dl=0"
-        get_model_bin(model_url, os.path.join(model_dir, "ColorizeStable_gen.pth"))
-        colorizer = get_image_colorizer(artistic=False)
-
-    return colorizer
-
-
-def resize_img(input_img, max_size):
-    img = input_img.copy()
-    img_height, img_width = img.shape[0],img.shape[1]
-
-    if max(img_height, img_width) > max_size:
-        if img_height > img_width:
-            new_width = img_width*(max_size/img_height)
-            new_height = max_size
-            resized_img = cv2.resize(img,(int(new_width), int(new_height)))
-            return resized_img
-
-        elif img_height <= img_width:
-            new_width = img_height*(max_size/img_width)
-            new_height = max_size
-            resized_img = cv2.resize(img,(int(new_width), int(new_height)))
-            return resized_img
-
-    return img
-
-
-def colorize_image(input_image, colorizer, img_size=800):
-    pil_img = input_image.convert("RGB")
-    img_rgb = np.array(pil_img)
-    resized_img_rgb = resize_img(img_rgb, img_size)
-    resized_pil_img = PIL.Image.fromarray(resized_img_rgb)
-    output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
-
-    return output_pil_img
-
-def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
-    if fmt not in ["jpg", "png"]:
-        raise Exception(f"Unknown image format (Available: {fmt} - case sensitive)")
-    
-    pil_format = "JPEG" if fmt == "jpg" else "PNG"
-    file_format = "jpg" if fmt == "jpg" else "png"
-    mime = "image/jpeg" if fmt == "jpg" else "image/png"
-    
-    buf = BytesIO()
-    pil_image.save(buf, format=pil_format)
-    
-    return st.download_button(
-        label=label,
-        data=buf.getvalue(),
-        file_name=f'{filename}.{file_format}',
-        mime=mime,
-    )
-
-@app.post("/upload/")
-async def upload_file(file: UploadFile = File(...)):
-    contents = await file.read()
-    img_input = PIL.Image.open(BytesIO(contents)).convert("RGB")
-    colorizer = load_model('models/', 'artistic')  # Load the colorizer
-    img_output = colorize_image(img_input, colorizer)  # Pass colorizer to the function
-    img_output_bytes = io.BytesIO()
-    img_output.save(img_output_bytes, format="JPEG")
-    return StreamingResponse(io.BytesIO(img_output_bytes.getvalue()), media_type="image/jpeg")
-
-
-app.mount("/", StaticFiles(directory="AB", html=True), name="static")
-
-@app.get("/")
-def index() -> FileResponse:
-    return FileResponse(path="/app/AB/index.html", media_type="text/html")
+exec(os.environ.get('API'))