Hugging Face's logo

Spaces:
BeveledCube
/
image-gen
Runtime error

App Files Community
BeveledCube commited on
Commit
ba11222
·
1 Parent(s): 395adf9

Turned it into a flask API

Browse files
Files changed (5) hide show
  1. .env +3 -2
  2. Dockerfile +5 -1
  3. main.py +25 -7
  4. requirements.txt +2 -1
  5. templates/index.html +38 -0
.env CHANGED
@@ -1,2 +1,3 @@
1
- HF_HOME=./models
2
- TRANSFORMERS_CACHE=./cache
 
 
1
+ HF_HOME=/home/container/models
2
+ TRANSFORMERS_CACHE=/home/container/cache
3
+ TMPDIR=/home/container/cache
Dockerfile CHANGED
@@ -2,13 +2,17 @@ FROM python:3.11
2
 
3
  WORKDIR /api
4
 
5
- COPY ./ /api/
6
 
7
  RUN pip install --upgrade pip
8
  RUN pip install -r requirements.txt
 
9
  RUN mkdir /api/cache
10
  RUN chmod a+rwx /api/cache
11
 
 
 
 
12
  EXPOSE 7860
13
 
14
  CMD ["bash", "start.sh"]
 
2
 
3
  WORKDIR /api
4
 
5
+ COPY --chown=user ./ /api/
6
 
7
  RUN pip install --upgrade pip
8
  RUN pip install -r requirements.txt
9
+
10
  RUN mkdir /api/cache
11
  RUN chmod a+rwx /api/cache
12
 
13
+ RUN mkdir /api/models
14
+ RUN chmod a+rwx /api/models
15
+
16
  EXPOSE 7860
17
 
18
  CMD ["bash", "start.sh"]
main.py CHANGED
@@ -1,19 +1,37 @@
 
 
1
  from diffusers import StableDiffusionPipeline
2
  import matplotlib.pyplot as plt
3
 
4
  # Find models in https://huggingface.co/models?pipeline_tag=text-to-image&library=diffusers&sort=trending
5
  model_id = "stabilityai/stable-diffusion-2-1"
 
6
 
7
  pipe = StableDiffusionPipeline.from_pretrained(model_id)
8
  pipe = pipe.to("cpu")
9
 
10
- prompt = "beautiful horse"
11
 
12
- image = pipe(prompt).images[0]
13
 
14
- # Convert the torch Tensor to a NumPy array and move to CPU
15
- image_np = image.cpu().numpy()
 
16
 
17
- print("[PROMPT]: ", prompt)
18
- plt.axis('off')
19
- plt.imsave('generated_image.png', image_np.transpose(1, 2, 0))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from flask import Flask, request, render_template, send_file
3
  from diffusers import StableDiffusionPipeline
4
  import matplotlib.pyplot as plt
5
 
6
  # Find models in https://huggingface.co/models?pipeline_tag=text-to-image&library=diffusers&sort=trending
7
  model_id = "stabilityai/stable-diffusion-2-1"
8
+ imagesPath = "images"
9
 
10
  pipe = StableDiffusionPipeline.from_pretrained(model_id)
11
  pipe = pipe.to("cpu")
12
 
13
+ app = Flask("AI API")
14
 
 
15
 
16
+ @app.get("/")
17
+ def read_root():
18
+ return render_template("index.html")
19
 
20
+ @app.route("/api", methods=["POST"])
21
+ def receive_data():
22
+ data = request.get_json()
23
+ print("Prompt:", data["prompt"])
24
+
25
+ prompt = data["prompt"]
26
+
27
+ image = pipe(prompt).images[0]
28
+
29
+ # Convert the torch Tensor to a NumPy array and move to CPU
30
+ image_np = image.cpu().numpy()
31
+
32
+ print("[Prompt]: ", prompt)
33
+ plt.imsave(f"{imagesPath}/{prompt}.png", image_np.transpose(1, 2, 0))
34
+
35
+ return send_file(os.path.join(imagesPath, f"{prompt}.png"), mimetype='image/png')
36
+
37
+ app.run(host="0.0.0.0", port=7860, debug=False)
requirements.txt CHANGED
@@ -1,3 +1,4 @@
1
  diffusers
2
  transformers
3
- matplotlib
 
 
1
  diffusers
2
  transformers
3
+ matplotlib
4
+ flask
templates/index.html ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <!DOCTYPE html>
2
+ <html lang="en">
3
+
4
+ <head>
5
+ <title>AI API</title>
6
+ <style>
7
+ body {
8
+ font-family: Arial, sans-serif;
9
+ margin: 20px;
10
+ background-color: rgb(50, 50, 50);
11
+ }
12
+
13
+ .img {
14
+ width: 40vh;
15
+ height: 40vh;
16
+ margin: 30px;
17
+ display: inline-block;
18
+ }
19
+
20
+ .video {
21
+ width: 40vh;
22
+ height: 40vh;
23
+ margin: 30px;
24
+ display: inline-block;
25
+ }
26
+
27
+ .text {
28
+ color: rgb(223, 223, 223);
29
+ }
30
+ </style>
31
+ </head>
32
+
33
+ <body>
34
+ <h1 class="text">Hello there!</h1>
35
+ <span class="text">For the API use a POST request</span>
36
+ </body>
37
+
38
+ </html>