MohamedRashad commited on
Commit
9200fea
·
1 Parent(s): 699c5a8

Update image generation pipeline in app.py to use new model and enable sequential CPU offload

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -24,9 +24,10 @@ from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_imag
24
  llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
25
 
26
  device = "cuda" if torch.cuda.is_available() else "cpu"
27
- pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
28
  pipe.vae.enable_tiling()
29
  pipe.vae.enable_slicing()
 
30
  torch.cuda.empty_cache()
31
 
32
  pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
@@ -87,7 +88,7 @@ def generate_item_image(object_t2i_prompt):
87
  for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
88
  prompt=object_t2i_prompt,
89
  guidance_scale=3.5,
90
- num_inference_steps=28,
91
  width=1024,
92
  height=1024,
93
  generator=torch.Generator("cpu").manual_seed(0),
 
24
  llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
25
 
26
  device = "cuda" if torch.cuda.is_available() else "cpu"
27
+ pipe = FluxPipeline.from_pretrained("sayakpaul/FLUX.1-merged", torch_dtype=torch.bfloat16).to(device)
28
  pipe.vae.enable_tiling()
29
  pipe.vae.enable_slicing()
30
+ pipe.enable_sequential_cpu_offload() # offloads modules to CPU on a submodule level (rather than model level)
31
  torch.cuda.empty_cache()
32
 
33
  pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
88
  for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
89
  prompt=object_t2i_prompt,
90
  guidance_scale=3.5,
91
+ num_inference_steps=8,
92
  width=1024,
93
  height=1024,
94
  generator=torch.Generator("cpu").manual_seed(0),