Jesivn commited on
Commit
edf0756
·
verified ·
1 Parent(s): 918f81c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -4,15 +4,15 @@ from diffusers import DiffusionPipeline, FluxPipeline
4
  import torch
5
  import time
6
 
7
- device = "cuda" if torch.cuda.is_available() else "cpu"
8
  # Load models
9
  translator = pipeline("translation", model="Helsinki-NLP/opus-mt-dra-en")
10
  summarizer = pipeline("summarization", model="Falconsai/text_summarization")
11
- image_pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device)
12
  # image_pipe = fluxPipeline("image-generation", model=DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device))
13
 
14
- # image_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(device)
15
- # image_pipe.enable_model_cpu_offload()
16
 
17
 
18
  # Functions for each task
 
4
  import torch
5
  import time
6
 
7
+ device = "cpu"
8
  # Load models
9
  translator = pipeline("translation", model="Helsinki-NLP/opus-mt-dra-en")
10
  summarizer = pipeline("summarization", model="Falconsai/text_summarization")
11
+ # image_pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device)
12
  # image_pipe = fluxPipeline("image-generation", model=DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device))
13
 
14
+ image_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(device)
15
+ image_pipe.enable_model_cpu_offload()
16
 
17
 
18
  # Functions for each task