Update app.py

app.py

@@ -3,18 +3,39 @@ import numpy as np
 import random
 import spaces
 import torch
+import os
+import time
 from diffusers import DiffusionPipeline
+from huggingface_hub import login
 
+# Ensure sentencepiece is installed in your environment
+try:
+    import sentencepiece
+except ImportError:
+    raise ImportError("The 'sentencepiece' library is required but not installed. Please add it to your environment.")
+
+# Access the API token securely from Hugging Face Secrets
+hf_api_token = os.getenv("HF_API_TOKEN")
+
+if hf_api_token:
+    login(token=hf_api_token)
+else:
+    raise ValueError("Hugging Face API token not found in secrets.")
+
+# Set the device and dtype
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+# Load the diffusion pipeline from the gated repository
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
 @spaces.GPU()
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidance_scale=7.5, progress=gr.Progress(track_tqdm=True)):
+    start_time = time.time()
+    
     if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
         raise ValueError("Image size exceeds the maximum allowed dimensions.")
     
@@ -24,23 +45,27 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
     
     try:
         image = pipe(
-                prompt=prompt, 
-                width=width,
-                height=height,
-                num_inference_steps=num_inference_steps, 
-                generator=generator,
-                guidance_scale=guidance_scale
-        ).images[0] 
+            prompt=prompt, 
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps, 
+            generator=generator,
+            guidance_scale=guidance_scale
+        ).images[0]
     except Exception as e:
+        print(f"Error generating image: {e}")
         return None, seed, f"Error: {str(e)}"
     
+    # Check if it took too long
+    if time.time() - start_time > 60:  # 60 seconds timeout
+        return None, seed, "Image generation took too long and was cancelled."
+
     return image, seed, None
  
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
     "a cat holding a sign that says hello world",
     "an anime illustration of a wiener schnitzel",
-    # Add more diverse examples
 ]
 
 css = """
@@ -55,7 +80,7 @@ with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""# Custom Image Creator
 12B param rectified flow transformer distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/) for 4 step generation
-[[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-schnell)]
+[[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1)]
         """)
         
         with gr.Row():