multimodalart HF staff commited on
Commit
3a748d4
1 Parent(s): f864900

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -7
app.py CHANGED
@@ -5,19 +5,23 @@ import logging
5
  import torch
6
  from PIL import Image
7
  import spaces
8
- from diffusers import DiffusionPipeline
 
9
  from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
10
  import copy
11
  import random
12
  import time
13
 
 
 
14
  # Load LoRAs from JSON file
15
  with open('loras.json', 'r') as f:
16
  loras = json.load(f)
17
 
18
  # Initialize the base model
19
- base_model = "black-forest-labs/FLUX.1-dev"
20
- pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 
21
 
22
  MAX_SEED = 2**32-1
23
 
@@ -60,19 +64,23 @@ def update_selection(evt: gr.SelectData, width, height):
60
  @spaces.GPU(duration=70)
61
  def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
62
  pipe.to("cuda")
 
 
63
  generator = torch.Generator(device="cuda").manual_seed(seed)
64
 
65
  with calculateDuration("Generating image"):
66
- # Generate image
67
- image = pipe(
68
  prompt=prompt_mash,
 
69
  num_inference_steps=steps,
70
- guidance_scale=cfg_scale,
71
  width=width,
72
  height=height,
 
73
  generator=generator,
 
74
  joint_attention_kwargs={"scale": lora_scale},
75
- ).images[0]
 
76
  return image
77
 
78
  def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
 
5
  import torch
6
  from PIL import Image
7
  import spaces
8
+ from diffusers import DiffusionPipeline, AutoencoderTiny
9
+ from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
10
  from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
11
  import copy
12
  import random
13
  import time
14
 
15
+ dtype = torch.bfloat16
16
+
17
  # Load LoRAs from JSON file
18
  with open('loras.json', 'r') as f:
19
  loras = json.load(f)
20
 
21
  # Initialize the base model
22
+ taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
23
+ pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
24
+ torch.cuda.empty_cache()
25
 
26
  MAX_SEED = 2**32-1
27
 
 
64
  @spaces.GPU(duration=70)
65
  def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
66
  pipe.to("cuda")
67
+ if randomize_seed:
68
+ seed = random.randint(0, MAX_SEED)
69
  generator = torch.Generator(device="cuda").manual_seed(seed)
70
 
71
  with calculateDuration("Generating image"):
72
+ for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
 
73
  prompt=prompt_mash,
74
+ guidance_scale=guidance_scale,
75
  num_inference_steps=steps,
 
76
  width=width,
77
  height=height,
78
+ guidance_scale=cfg_scale,
79
  generator=generator,
80
+ output_type="pil",
81
  joint_attention_kwargs={"scale": lora_scale},
82
+ ):
83
+ yield img, seed
84
  return image
85
 
86
  def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):