tecuts commited on
Commit
2efdf09
·
verified ·
1 Parent(s): ea17dd4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -8
app.py CHANGED
@@ -71,6 +71,8 @@ os.makedirs(DOWNLOADS_DIR, exist_ok=True)
71
  # Mount the downloads directory
72
  app.mount("/files", StaticFiles(directory=DOWNLOADS_DIR), name="files")
73
 
 
 
74
  class DownloadRequest(BaseModel):
75
  url: str
76
 
@@ -79,14 +81,10 @@ class DownloadResponse(BaseModel):
79
  message: str
80
  filename: str = None
81
  download_url: str = None
 
82
 
83
 
84
 
85
- def verify_file_size(file_path: str, expected_size: int):
86
- actual_size = os.path.getsize(file_path)
87
- if actual_size < expected_size:
88
- raise Exception(f"Downloaded file size mismatch: expected {expected_size} bytes, got {actual_size} bytes")
89
-
90
  @app.post("/download", response_model=DownloadResponse)
91
  async def download_file(request: DownloadRequest):
92
  try:
@@ -136,8 +134,8 @@ async def download_file(request: DownloadRequest):
136
  downloaded_file = files[0]
137
  logger.info(f"Downloaded file: {downloaded_file}")
138
 
139
- # Verify the size of the downloaded file (expected size in bytes, e.g., 7000000 for 7MB)
140
- verify_file_size(downloaded_file, expected_size=7000000)
141
 
142
  # Generate the download URL
143
  space_url = os.getenv("SPACE_URL", "http://localhost:7860")
@@ -151,7 +149,8 @@ async def download_file(request: DownloadRequest):
151
  success=True,
152
  message="File downloaded successfully",
153
  filename=downloaded_file,
154
- download_url=download_url
 
155
  )
156
 
157
  except subprocess.CalledProcessError as e:
@@ -172,6 +171,7 @@ async def download_file(request: DownloadRequest):
172
  os.chdir(original_dir)
173
 
174
 
 
175
  @app.get("/test")
176
  async def test():
177
  """Test endpoint to verify setup"""
 
71
  # Mount the downloads directory
72
  app.mount("/files", StaticFiles(directory=DOWNLOADS_DIR), name="files")
73
 
74
+
75
+
76
  class DownloadRequest(BaseModel):
77
  url: str
78
 
 
81
  message: str
82
  filename: str = None
83
  download_url: str = None
84
+ file_size: int = None # Include file size in the response
85
 
86
 
87
 
 
 
 
 
 
88
  @app.post("/download", response_model=DownloadResponse)
89
  async def download_file(request: DownloadRequest):
90
  try:
 
134
  downloaded_file = files[0]
135
  logger.info(f"Downloaded file: {downloaded_file}")
136
 
137
+ # Get the file size
138
+ file_size = os.path.getsize(downloaded_file)
139
 
140
  # Generate the download URL
141
  space_url = os.getenv("SPACE_URL", "http://localhost:7860")
 
149
  success=True,
150
  message="File downloaded successfully",
151
  filename=downloaded_file,
152
+ download_url=download_url,
153
+ file_size=file_size # Include the file size in the response
154
  )
155
 
156
  except subprocess.CalledProcessError as e:
 
171
  os.chdir(original_dir)
172
 
173
 
174
+
175
  @app.get("/test")
176
  async def test():
177
  """Test endpoint to verify setup"""