fix: requirements.txt, app.py

app.py

@@ -4,32 +4,35 @@ import vtracer
 import tempfile
 import cairosvg
 import re
+import uvicorn
 from PIL import Image
 from datetime import datetime
-
-from flask import Flask, request, jsonify, send_from_directory
-from flask_cors import CORS
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import FileResponse, JSONResponse, Response
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional
 
 from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
-
 import torchvision.transforms as transforms
 from model import Generator 
 
+SVG_DIR = os.path.join(os.getcwd(), 'generated_svgs')
+THUMBNAIL_DIR = os.path.join(os.getcwd(), 'thumbnails')
+SKETCH_MODEL_WEIGHTS = 'checkpoints/netG_A_latest.pth'
+
 def setup_directories():
+    """Creates necessary directories if they don't exist."""
     os.makedirs(SVG_DIR, exist_ok=True)
     os.makedirs(THUMBNAIL_DIR, exist_ok=True)
     print(f"Directories '{SVG_DIR}' and '{THUMBNAIL_DIR}' are ready.")
 
-def sanitize_filename(prompt):
+def sanitize_filename(prompt: str) -> str:
     """Removes characters that are invalid for filenames."""
-
     s = re.sub(r'[\\/*?:"<>|]', "", prompt)
+    return s.strip()[:100]
 
-    return s[:100]
-
-SVG_DIR = os.path.join(os.getcwd(), 'generated_svgs')
-THUMBNAIL_DIR = os.path.join(os.getcwd(), 'thumbnails')
-SKETCH_MODEL_WEIGHTS = 'checkpoints/netG_A_latest.pth'
 
 class ImageToSvgPipeline:
     """
@@ -46,7 +49,7 @@ class ImageToSvgPipeline:
     def _initialize_rinna_model(self):
         print("Loading Rinna Stable Diffusion model...")
         model_id = "rinna/japanese-stable-diffusion"
-
+        
         self.rinna_pipe = StableDiffusionPipeline.from_pretrained(
             model_id,
             torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
@@ -101,17 +104,16 @@ class ImageToSvgPipeline:
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
             image.save(tmp_file.name)
             tmp_path = tmp_file.name
-
+        
+        svg_output_path = tmp_path.replace(".png", ".svg")
         try:
-            svg_output_path = tmp_path.replace(".png", ".svg")
             vtracer.convert_image_to_svg_py(tmp_path, svg_output_path)
-
             with open(svg_output_path, 'r', encoding='utf-8') as f:
                 svg_data = f.read()
         finally:
             if os.path.exists(tmp_path): os.remove(tmp_path)
-            if 'svg_output_path' in locals() and os.path.exists(svg_output_path): os.remove(svg_output_path)
-
+            if os.path.exists(svg_output_path): os.remove(svg_output_path)
+        
         print("SVG extraction complete.")
         return svg_data
 
@@ -121,28 +123,37 @@ class ImageToSvgPipeline:
         svg_content = self._extract_svg(sketch_image)
         return svg_content
 
-app = Flask(__name__)
+setup_directories()
+pipeline = ImageToSvgPipeline(sketch_model_path=SKETCH_MODEL_WEIGHTS)
 
-CORS(app, resources={r"/*": {"origins": "*"}})
+app = FastAPI()
 
-pipeline = ImageToSvgPipeline(sketch_model_path=SKETCH_MODEL_WEIGHTS)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"], # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"], # Allows all methods
+    allow_headers=["*"], # Allows all headers
+)
 
-def sanitize_filename(text):
-    text = re.sub(r'[\\/*?:"<>|]', "", text)
-    return text.strip()
+class GenerateRequest(BaseModel):
+    prompt: str
 
-@app.route('/generate', methods=['POST'])
-def generate_svg():
-    data = request.json
-    prompt = data.get('prompt')
-    if not prompt: return jsonify({"error": "Prompt is required"}), 400
+@app.post("/generate")
+async def generate_svg(item: GenerateRequest):
+    """
+    Receives a prompt, generates an SVG, saves it, and returns the SVG content.
+    """
+    if not item.prompt:
+        raise HTTPException(status_code=400, detail="Prompt is required")
 
     negative_prompt = "低品質、最悪の品質、下手な手、指が6本、指が4本、奇形、醜い、ぼやけている、ぼやけた、ウォーターマーク、署名、テキスト"
     try:
-        svg_result = pipeline.process(prompt, negative_prompt)
+        svg_result = pipeline.process(item.prompt, negative_prompt)
 
+        # Save the SVG and its thumbnail
         timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-        safe_prompt = sanitize_filename(prompt)[:50]
+        safe_prompt = sanitize_filename(item.prompt)[:50]
         filename = f"{timestamp}_{safe_prompt}.svg"
 
         svg_path = os.path.join(SVG_DIR, filename)
@@ -152,17 +163,20 @@ def generate_svg():
         thumbnail_path = os.path.join(THUMBNAIL_DIR, filename.replace('.svg', '.png'))
         cairosvg.svg2png(bytestring=svg_result.encode('utf-8'), write_to=thumbnail_path, output_width=256, output_height=256)
 
-        return svg_result, 200, {'Content-Type': 'image/svg+xml'}
+        # Return the SVG data directly in the response
+        return Response(content=svg_result, media_type="image/svg+xml")
+
     except Exception as e:
         print(f"An error occurred during generation: {e}")
-        return jsonify({"error": str(e)}), 500
+        raise HTTPException(status_code=500, detail=str(e))
 
-@app.route('/gallery', methods=['GET'])
-def get_gallery():
-    try:
-        page = int(request.args.get('page', 1))
-        limit = int(request.args.get('limit', 8))
 
+@app.get("/gallery")
+def get_gallery(page: int = 1, limit: int = 8):
+    """
+    Returns a paginated list of generated drawings.
+    """
+    try:
         svg_files = sorted([f for f in os.listdir(SVG_DIR) if f.endswith('.svg')], reverse=True)
 
         start_index = (page - 1) * limit
@@ -180,31 +194,40 @@ def get_gallery():
             })
 
         has_more = end_index < len(svg_files)
-        return jsonify({"drawings": drawings, "hasMore": has_more})
+        return {"drawings": drawings, "hasMore": has_more}
     except Exception as e:
         print(f"Error fetching gallery: {e}")
-        return jsonify({"error": "Failed to fetch gallery"}), 500
+        raise HTTPException(status_code=500, detail="Failed to fetch gallery")
 
-@app.route('/svgs/<path:filename>')
-def get_svg(filename):
-    return send_from_directory(SVG_DIR, filename)
 
-@app.route('/thumbnails/<path:filename>')
-def get_thumbnail(filename):
-    return send_from_directory(THUMBNAIL_DIR, filename)
-
-@app.route('/drawings/<path:filename>', methods=['DELETE'])
-def delete_drawing_file(filename):
+@app.delete("/drawings/{filename}")
+def delete_drawing_file(filename: str):
+    """
+    Deletes a specific SVG and its corresponding thumbnail.
+    """
     try:
+        # Sanitize filename to prevent directory traversal
+        if ".." in filename or "/" in filename:
+             raise HTTPException(status_code=400, detail="Invalid filename")
+
         svg_path = os.path.join(SVG_DIR, filename)
         thumb_path = os.path.join(THUMBNAIL_DIR, filename.replace('.svg', '.png'))
+
+        if not os.path.exists(svg_path):
+             raise HTTPException(status_code=404, detail="File not found")
+
         if os.path.exists(svg_path): os.remove(svg_path)
         if os.path.exists(thumb_path): os.remove(thumb_path)
-        return jsonify({"message": f"Successfully deleted {filename}"})
+        
+        return JSONResponse(content={"message": f"Successfully deleted {filename}"})
     except Exception as e:
         print(f"Error deleting file: {e}")
-        return jsonify({"error": "Failed to delete file"}), 500
+        raise HTTPException(status_code=500, detail="Failed to delete file")
+
+app.mount("/svgs", StaticFiles(directory=SVG_DIR), name="svgs")
+app.mount("/thumbnails", StaticFiles(directory=THUMBNAIL_DIR), name="thumbnails")
+
 
 if __name__ == '__main__':
-    print("Starting Flask server...")
-    app.run(host='0.0.0.0', port=5000)
+    print("Starting FastAPI server...")
+    uvicorn.run(app, host='0.0.0.0', port=5000)

requirements.txt

@@ -1,5 +1,7 @@
 flask
 Flask-Cors
+fastapi
+uvicorn[standard]
 torch
 diffusers==0.35.1
 transformers==4.56.2
@@ -11,3 +13,4 @@ torchvision==0.23.0
 sentencepiece==0.2.0
 scipy
 numpy
+python-multipart