fastapi added

app.py

@@ -20,9 +20,11 @@ from typing import List, Optional
 # Writable directories (Spaces often allow /tmp). Override via env if needed.
 OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/face_upscale/output")
 INPUT_DIR  = os.environ.get("INPUT_DIR",  "/tmp/face_upscale/input")
+WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", "/tmp/face_upscale/weights")
 # Ensure required directories exist at import time (for API mode)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 os.makedirs(INPUT_DIR, exist_ok=True)
+os.makedirs(WEIGHTS_DIR, exist_ok=True)
 
 # FastAPI imports
 from fastapi import FastAPI, UploadFile, File, Form, Depends, HTTPException, status
@@ -369,14 +371,14 @@ class Upscale:
 
     def initBGUpscaleModel(self, upscale_model):
         upscale_type, upscale_model = upscale_model.split(", ", 1)
-        download_from_url(upscale_models[upscale_model][0], upscale_model, os.path.join("weights", "upscale"))
+        download_from_url(upscale_models[upscale_model][0], upscale_model, os.path.join(WEIGHTS_DIR, "upscale"))
         self.modelInUse = f"_{os.path.splitext(upscale_model)[0]}"
         netscale = 1 if any(sub in upscale_model.lower() for sub in ("x1", "1x")) else (2 if any(sub in upscale_model.lower() for sub in ("x2", "2x")) else 4)
         model = None
         half = True if torch.cuda.is_available() else False
         if upscale_type:
             from basicsr.archs.rrdbnet_arch import RRDBNet
-            loadnet = torch.load(os.path.join("weights", "upscale", upscale_model), map_location=torch.device('cpu'), weights_only=True)
+            loadnet = torch.load(os.path.join(WEIGHTS_DIR, "upscale", upscale_model), map_location=torch.device('cpu'), weights_only=True)
             if 'params_ema' in loadnet or 'params' in loadnet:
                 loadnet = loadnet['params_ema'] if 'params_ema' in loadnet else loadnet['params']
             if upscale_type == "SRVGG":
@@ -566,7 +568,7 @@ class Upscale:
                 model = SRFormer(img_size=img_size, in_chans=in_chans, embed_dim=embed_dim, depths=depths, num_heads=num_heads, window_size=window_size, mlp_ratio=mlp_ratio, 
                                  qkv_bias=qkv_bias, qk_scale=None, ape=ape, patch_norm=patch_norm, upscale=netscale, upsampler=upsampler, resi_connection=resi_connection)
         if model:
-            self.realesrganer = RealESRGANer(scale=netscale, model_path=os.path.join("weights", "upscale", upscale_model), model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+            self.realesrganer = RealESRGANer(scale=netscale, model_path=os.path.join(WEIGHTS_DIR, "upscale", upscale_model), model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
         elif upscale_model:
             import PIL
             from image_gen_aux import UpscaleWithModel
@@ -604,12 +606,12 @@ class Upscale:
                             ), interpolation=interpolation)
                     return cv_image, None
             device = "cuda" if torch.cuda.is_available() else "cpu"
-            upscaler = UpscaleWithModel.from_pretrained(os.path.join("weights", "upscale", upscale_model)).to(device)
+            upscaler = UpscaleWithModel.from_pretrained(os.path.join(WEIGHTS_DIR, "upscale", upscale_model)).to(device)
             upscaler.__class__ = UpscaleWithModel_Gfpgan
             self.realesrganer = upscaler
 
     def initFaceEnhancerModel(self, face_restoration, face_detection):
-        model_rootpath = os.path.join("weights", "face")
+        model_rootpath = os.path.join(WEIGHTS_DIR, "face")
         model_path = os.path.join(model_rootpath, face_restoration)
         download_from_url(face_models[face_restoration][0], face_restoration, model_rootpath)
         self.modelInUse = f"_{os.path.splitext(face_restoration)[0]}" + self.modelInUse