wifix199 commited on
Commit
93e97ff
·
verified ·
1 Parent(s): 9e3d03f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -1,10 +1,10 @@
1
  import gradio as gr
2
  import torch
3
- from diffusers import FluxPipeline
4
 
5
  # Load the model
6
- pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
7
- pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
8
 
9
  def generate_image(prompt):
10
  image = pipe(
@@ -28,4 +28,4 @@ ui = gr.Interface(
28
  )
29
 
30
  # Launch the UI
31
- ui.launch()
 
1
  import gradio as gr
2
  import torch
3
+ from diffusers import DiffusionPipeline # Use DiffusionPipeline instead
4
 
5
  # Load the model
6
+ pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
7
+ pipe.enable_model_cpu_offload() # save some VRAM by offloading the model to CPU
8
 
9
  def generate_image(prompt):
10
  image = pipe(
 
28
  )
29
 
30
  # Launch the UI
31
+ ui.launch()