Axel-Student commited on
Commit
1fc5ea4
·
1 Parent(s): c9fdad1
Files changed (1) hide show
  1. app.py +5 -0
app.py CHANGED
@@ -4,13 +4,18 @@ import torch
4
  from fastapi import FastAPI, HTTPException
5
  from fastapi.responses import StreamingResponse
6
  from diffusers import FluxPipeline # type: ignore
 
7
 
8
  app = FastAPI()
9
 
10
  # Récupération du token et authentification
11
  token = os.getenv("HF_TOKEN")
 
12
  login(token=token)
13
 
 
 
 
14
  def generate_image(prompt: str):
15
  image = pipe(
16
  prompt,
 
4
  from fastapi import FastAPI, HTTPException
5
  from fastapi.responses import StreamingResponse
6
  from diffusers import FluxPipeline # type: ignore
7
+ from huggingface_hub import login
8
 
9
  app = FastAPI()
10
 
11
  # Récupération du token et authentification
12
  token = os.getenv("HF_TOKEN")
13
+ token = os.getenv("HF_TOKEN")
14
  login(token=token)
15
 
16
+ pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
17
+ pipe.enable_model_cpu_offload()
18
+
19
  def generate_image(prompt: str):
20
  image = pipe(
21
  prompt,