Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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
|
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 |
-
|
|
|
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
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
|
|
|
|
|
|
|
|
35 |
if randomize_seed:
|
36 |
seed = random.randint(0, MAX_SEED)
|
37 |
generator = torch.Generator().manual_seed(seed)
|
38 |
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
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",
|