bug fix: updating duration func

app.py

@@ -62,7 +62,7 @@ def remote_text_encoder(prompts):
     except Exception as e:
         raise Exception(f"Failed to encode prompt: {str(e)}")
 
-def get_duration(num_inference_steps: int, input_image: Image.Image = None):
+def get_duration(prompt: str, input_image: Image.Image = None, num_inference_steps: int = 28, guidance_scale: float = 4.0, seed: int = 42, progress=None):
     """Calculate dynamic GPU duration based on inference steps and input image."""
     num_images = 0 if input_image is None else 1
     step_duration = 1 + 0.7 * num_images
@@ -87,6 +87,10 @@ def generate_image(
         guidance_scale: How closely to follow the prompt (higher = more strict)
         seed: Random seed for reproducibility (-1 for random)
     """
+    print(f"=== Starting generation ===")
+    print(f"Prompt: {prompt[:100]}...")
+    print(f"CUDA available: {torch.cuda.is_available()}")
+    
     if not prompt or prompt.strip() == "":
         raise gr.Error("Please enter a prompt!")
     
@@ -94,21 +98,32 @@ def generate_image(
     
     try:
         # Load pipeline (lazy loading)
+        print("Loading pipeline...")
         pipeline = load_pipeline()
+        print("Pipeline loaded successfully")
         
         progress(0.1, desc="Encoding prompt...")
+        print("Encoding prompt...")
         
         # Get prompt embeddings from remote encoder
-        prompt_embeds = remote_text_encoder(prompt)
+        try:
+            prompt_embeds = remote_text_encoder(prompt)
+            print(f"Prompt embeds shape: {prompt_embeds.shape}")
+        except Exception as e:
+            print(f"Error encoding prompt: {str(e)}")
+            raise gr.Error(f"Failed to encode prompt. Please check your HuggingFace token. Error: {str(e)}")
         
         progress(0.3, desc="Generating image...")
         
         # Set up generator
         generator_device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Generator device: {generator_device}")
+        
         if seed == -1:
             import random
             seed = random.randint(0, 2**32 - 1)
         
+        print(f"Using seed: {seed}")
         generator = torch.Generator(device=generator_device).manual_seed(int(seed))
         
         # Prepare pipeline arguments
@@ -123,25 +138,43 @@ def generate_image(
         if input_image is not None:
             pipe_kwargs["image"] = input_image
             progress(0.4, desc="Processing input image...")
+            print("Processing with input image")
+        
+        print(f"Starting generation with {num_inference_steps} steps...")
         
         # Generate image
         with torch.inference_mode():
-            image = pipeline(**pipe_kwargs).images[0]
+            result = pipeline(**pipe_kwargs)
+            image = result.images[0]
         
+        print("Generation complete!")
         progress(1.0, desc="Done!")
         
         return image
     
+    except gr.Error:
+        # Re-raise Gradio errors as-is
+        raise
     except Exception as e:
         import traceback
         error_msg = f"Error generating image: {str(e)}\n{traceback.format_exc()}"
         print(error_msg)
-        raise gr.Error(f"Error generating image: {str(e)}")
+        
+        # Provide more helpful error messages
+        if "CUDA" in str(e):
+            raise gr.Error(f"GPU Error: {str(e)}. The model requires GPU to run.")
+        elif "token" in str(e).lower() or "401" in str(e):
+            raise gr.Error("Authentication failed. Please ensure your HuggingFace token is set correctly.")
+        elif "timeout" in str(e).lower():
+            raise gr.Error("Request timed out. Please try again.")
+        else:
+            raise gr.Error(f"Error: {str(e)}")
 
 
 # Create Gradio interface
 with gr.Blocks(
     title="Flux2 Image Generator",
+    theme=gr.themes.Soft(),
 ) as demo:
     gr.Markdown(
         """