fix: missing torch

app.py

@@ -1,3 +1,4 @@
+from diffusers import AutoPipelineForText2Image
 from flask import (
     Flask,
     jsonify,
@@ -5,8 +6,7 @@ from flask import (
     send_file,
 )
 from io import BytesIO
-from diffusers import AutoPipelineForText2Image
-
+from torch import float16
 
 app = Flask("ImageAPI")
 
@@ -17,11 +17,15 @@ def main():
 
     if not prompt:
         return jsonify({
-            "status": "ERROR",
-            "message": "Prompt is not found!",
+            "status": 200,
+            "message": "API is Alive!",
         })
 
-    pipeline = AutoPipelineForText2Image.from_pretrained(pretrained_model_or_path="Lykon/DreamShaper")
+    pipeline = AutoPipelineForText2Image.from_pretrained(
+        pretrained_model_or_path="Lykon/DreamShaper",
+        torch_dtype=float16,
+        variant="fp16",
+    )
     pipeline.to("cuda")
     
     image = pipeline(

requirements.txt

@@ -1,4 +1,5 @@
 diffusers
 flask
 gevent
-gunicorn
+gunicorn
+torch