SeedOfEvil commited on
Commit
75460c3
·
verified ·
1 Parent(s): 8ccbda4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +33 -23
app.py CHANGED
@@ -7,8 +7,9 @@ from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, Autoen
7
  from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
8
  from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
9
 
10
- dtype = torch.bfloat16
11
  device = "cuda" if torch.cuda.is_available() else "cpu"
 
12
 
13
  taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
14
  good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
@@ -21,33 +22,42 @@ MAX_IMAGE_SIZE = 2048
21
  # Bind the live preview helper function to the pipe
22
  pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
23
 
24
- @spaces.GPU(duration=75)
 
25
  def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
26
- # --- Fix: Truncate prompt to allowed token length ---
27
- # Use the pipeline's tokenizer to tokenize the prompt.
28
- tokenizer = pipe.tokenizer
29
- max_length = tokenizer.model_max_length # Typically 77 for CLIP
30
- tokens = tokenizer.tokenize(prompt)
31
- if len(tokens) > max_length:
32
- prompt = tokenizer.decode(tokenizer.encode(prompt, truncation=True, max_length=max_length), skip_special_tokens=True)
33
- # -------------------------------------------------------
34
-
 
 
 
 
35
  if randomize_seed:
36
  seed = random.randint(0, MAX_SEED)
37
  generator = torch.Generator().manual_seed(seed)
38
 
39
- for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
40
- prompt=prompt,
41
- guidance_scale=guidance_scale,
42
- num_inference_steps=num_inference_steps,
43
- width=width,
44
- height=height,
45
- generator=generator,
46
- output_type="pil",
47
- good_vae=good_vae,
48
- ):
49
- yield img, seed
50
-
 
 
 
 
51
  examples = [
52
  "a tiny astronaut hatching from an egg on the moon",
53
  "a cat holding a sign that says hello world",
 
7
  from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
8
  from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
9
 
10
+ # Use float16 on CUDA to avoid potential dtype issues
11
  device = "cuda" if torch.cuda.is_available() else "cpu"
12
+ dtype = torch.float16 if device == "cuda" else torch.float32
13
 
14
  taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
15
  good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
 
22
  # Bind the live preview helper function to the pipe
23
  pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
24
 
25
+ # Increased GPU duration to 120 seconds
26
+ @spaces.GPU(duration=120)
27
  def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
28
+ # --- Fix: Truncate the prompt to maximum token length ---
29
+ try:
30
+ tokenizer = pipe.tokenizer
31
+ max_length = tokenizer.model_max_length # Typically 77 tokens for CLIP
32
+ # Tokenize and encode with truncation enabled
33
+ encoded = tokenizer(prompt, truncation=True, max_length=max_length, return_tensors="pt")
34
+ # Decode back to string (skipping special tokens)
35
+ prompt = tokenizer.decode(encoded.input_ids[0], skip_special_tokens=True)
36
+ except Exception as e:
37
+ print("Error during prompt truncation:", e)
38
+ raise e
39
+ # -----------------------------------------------------------
40
+
41
  if randomize_seed:
42
  seed = random.randint(0, MAX_SEED)
43
  generator = torch.Generator().manual_seed(seed)
44
 
45
+ try:
46
+ for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
47
+ prompt=prompt,
48
+ guidance_scale=guidance_scale,
49
+ num_inference_steps=num_inference_steps,
50
+ width=width,
51
+ height=height,
52
+ generator=generator,
53
+ output_type="pil",
54
+ good_vae=good_vae,
55
+ ):
56
+ yield img, seed
57
+ except Exception as e:
58
+ print("Error during image generation:", e)
59
+ raise e
60
+
61
  examples = [
62
  "a tiny astronaut hatching from an egg on the moon",
63
  "a cat holding a sign that says hello world",