Update api_run.py

api_run.py

@@ -1,38 +1,35 @@
-from fastapi import FastAPI, File, UploadFile, Form
-from fastapi.responses import StreamingResponse
+import gradio as gr
 from deoldify import device
 from deoldify.device_id import DeviceId
 from deoldify.visualize import get_image_colorizer
-import io
-import os
-from tempfile import NamedTemporaryFile
-
-
-device.set(device=DeviceId.GPU0)
+from PIL import Image
 
+device.set(device=DeviceId.CPU)  # Dùng CPU cho an toàn
 colorizer = get_image_colorizer(artistic=True)
 
-app = FastAPI(title="DeOldify Colorization API")
-
-@app.post("/colorize")
-async def colorize(file: UploadFile = File(...), render_factor: int = Form(35)):
-    # Lưu file upload tạm
-    suffix = os.path.splitext(file.filename)[1]
-    with NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
-        tmp_name = tmp.name
-        content = await file.read()
-        tmp.write(content)
+def colorize_image(image, render_factor=35):
+    # Lưu tạm
+    image.save("temp_input.jpg")
     
-    # Gọi DeOldify để tô màu ảnh
+    # Tô màu
     result_path = colorizer.plot_transformed_image(
-        path=tmp_name,
+        path="temp_input.jpg",
         render_factor=render_factor,
-        display_render_factor=False,
-        figsize=(8,8),
         watermarked=False
     )
+    
+    return Image.open(result_path)
+
+# Tạo giao diện Gradio
+demo = gr.Interface(
+    fn=colorize_image,
+    inputs=[
+        gr.Image(type="pil", label="Ảnh cần tô màu"),
+        gr.Slider(10, 45, value=35, label="Render Factor")
+    ],
+    outputs=gr.Image(type="pil", label="Ảnh đã tô màu"),
+    title="DeOldify - Tô màu ảnh cũ",
+    description="Upload ảnh đen trắng để tô màu"
+)
 
-    # Đọc ảnh kết quả trả về client
-    with open(result_path, "rb") as f:
-        img_bytes = f.read()
-    return StreamingResponse(io.BytesIO(img_bytes), media_type="image/png")
+demo.launch()