Update app.py

app.py

@@ -7,8 +7,9 @@ from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, Autoen
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
 from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 
-dtype = torch.bfloat16
+# Use float16 on CUDA to avoid potential dtype issues
 device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
 
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 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
 # Bind the live preview helper function to the pipe
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
-@spaces.GPU(duration=75)
+# Increased GPU duration to 120 seconds
+@spaces.GPU(duration=120)
 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)):
-    # --- Fix: Truncate prompt to allowed token length ---
-    # Use the pipeline's tokenizer to tokenize the prompt.
-    tokenizer = pipe.tokenizer
-    max_length = tokenizer.model_max_length  # Typically 77 for CLIP
-    tokens = tokenizer.tokenize(prompt)
-    if len(tokens) > max_length:
-        prompt = tokenizer.decode(tokenizer.encode(prompt, truncation=True, max_length=max_length), skip_special_tokens=True)
-    # -------------------------------------------------------
-    
+    # --- Fix: Truncate the prompt to maximum token length ---
+    try:
+        tokenizer = pipe.tokenizer
+        max_length = tokenizer.model_max_length  # Typically 77 tokens for CLIP
+        # Tokenize and encode with truncation enabled
+        encoded = tokenizer(prompt, truncation=True, max_length=max_length, return_tensors="pt")
+        # Decode back to string (skipping special tokens)
+        prompt = tokenizer.decode(encoded.input_ids[0], skip_special_tokens=True)
+    except Exception as e:
+        print("Error during prompt truncation:", e)
+        raise e
+    # -----------------------------------------------------------
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
     
-    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator,
-            output_type="pil",
-            good_vae=good_vae,
-        ):
-        yield img, seed
-    
+    try:
+        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+                prompt=prompt,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                width=width,
+                height=height,
+                generator=generator,
+                output_type="pil",
+                good_vae=good_vae,
+            ):
+            yield img, seed
+    except Exception as e:
+        print("Error during image generation:", e)
+        raise e
+
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
     "a cat holding a sign that says hello world",