Update app.py

app.py

@@ -11,6 +11,7 @@ from huggingface_hub import InferenceClient
 import tempfile
 from PIL import Image
 import concurrent.futures
+from fastapi import FastAPI
 
 def run_with_timeout(func, timeout=300, *args, **kwargs):
     """Run function with timeout"""
@@ -28,7 +29,28 @@ def pil_to_tempfile(image):
     image.save(tmp_path, format="PNG")
     return tmp_path
 
+def resize_and_crop(image, target_res=(832, 480)):
+    tw, th = target_res
+    iw, ih = image.size
+    scale = max(tw / iw, th / ih)
+    nw, nh = int(iw * scale), int(ih * scale)
+    image = image.resize((nw, nh), Image.LANCZOS)
+    left = (nw - tw) // 2
+    if ih > iw:
+        top = int((nh - th) * 0.25)
+    else:
+        top = (nh - th) // 2
+    right = left + tw
+    bottom = top + th
+    return image.crop((left, top, right, bottom))
+
+def aligned_num_frames(duration, fps=16):
+    n = int(duration * fps)
+    return ((n - 1) // 4) * 4 + 1
+
 token = os.getenv("HF_TOKEN")
+if not token:
+    raise RuntimeError("Please set HF_TOKEN environment variable")
 
 client = InferenceClient(token=token)
 
@@ -48,25 +70,6 @@ text_to_image_client = InferenceClient(
     api_key=token
 )
 
-def resize_and_crop(image, target_res=(832, 480)):
-    tw, th = target_res
-    iw, ih = image.size
-    scale = max(tw / iw, th / ih)
-    nw, nh = int(iw * scale), int(ih * scale)
-    image = image.resize((nw, nh), Image.LANCZOS)
-    left = (nw - tw) // 2
-    if ih > iw:
-        top = int((nh - th) * 0.25)
-    else:
-        top = (nh - th) // 2
-    right = left + tw
-    bottom = top + th
-    return image.crop((left, top, right, bottom))
-
-def aligned_num_frames(duration, fps=16):
-    n = int(duration * fps)
-    return ((n - 1) // 4) * 4 + 1
-
 image_output = None    
 video_output = None    
 
@@ -81,7 +84,8 @@ def video_tool(
     prompt: str = "high quality, detailed, sharp, cinematic", 
     duration: float = 4, 
     steps: int = 20, 
-    guidance: float = 3.0
+    guidance: float = 3.0,
+    progress: gr.Progress = gr.Progress(),
 ) -> str:
     """
     Generates a video from a starting image using Wan 2.1.
@@ -97,6 +101,8 @@ def video_tool(
     global video_output
     
     try:
+        progress(0.125, desc="Performing diffusion… this might take awhile..")
+
         FPS = 12  
         num_frames = aligned_num_frames(duration, FPS) 
         
@@ -113,10 +119,10 @@ def video_tool(
         
         video_bytes = run_with_timeout(generate_video, timeout=300)
         
+        progress(0.90, desc="Exporting..")
         out = tempfile.mktemp(suffix=".mp4")
         with open(out, "wb") as f:
-            f.write(video_bytes)
-        
+            f.write(video_bytes)      
         video_output = out
         return "Video successfully generated and stored for Gradio UI."
         
@@ -125,7 +131,7 @@ def video_tool(
         return f"Video generation failed: {e}"
       
 @tool
-def nsfw_detection_tool(nsfw_detection_input: Image.Image) -> str:
+def nsfw_detection_tool(nsfw_detection_input: Image.Image, progress: gr.Progress = gr.Progress(),) -> str:
     """
     Suitable for filtering through score explicit or inappropriate content in images.
     Args:
@@ -134,6 +140,8 @@ def nsfw_detection_tool(nsfw_detection_input: Image.Image) -> str:
         str: Highest score result.
     """
     try:
+        progress(0.125, desc="Checking image..")
+
         tmp_path = pil_to_tempfile(nsfw_detection_input)
         
         outputs = client.image_classification(
@@ -149,13 +157,14 @@ def nsfw_detection_tool(nsfw_detection_input: Image.Image) -> str:
             f"Confidence: {top_result.score:.2%}"
         )
 
+        progress(0.90, desc="Returning verdict..")
         return verdict
 
     except Exception as e:
         return f"NSFW detection failed: {e}"
 
 @tool
-def image_tool(prompt: str) -> str:
+def image_tool(prompt: str, progress: gr.Progress = gr.Progress(),) -> str:
     """
     Generate an image from text using SD3-Medium.
     Args:
@@ -166,6 +175,8 @@ def image_tool(prompt: str) -> str:
     global image_output  
     
     try:
+        progress(0.125, desc="Generating image.. this might take awhile...")
+
         def generate_image():
             return text_to_image_client.text_to_image(
                 prompt=prompt,
@@ -176,6 +187,7 @@ def image_tool(prompt: str) -> str:
                 height=992
             )
         
+        progress(0.90, desc="Exporting..")
         image = run_with_timeout(generate_image, timeout=300)
         image_output = image
         return "Image successfully generated and stored for Gradio UI."
@@ -186,7 +198,7 @@ def image_tool(prompt: str) -> str:
 
 
 @tool
-def search_tool(query: str) -> str:
+def search_tool(query: str, progress: gr.Progress = gr.Progress(),) -> str:
     """
     Search the web and return the most relevant results.
     Args:
@@ -195,7 +207,11 @@ def search_tool(query: str) -> str:
         str: The search results.
     """
     try:
+        progress(0.125, desc="Searchign the web...")
+
         web_search_tool = DuckDuckGoSearchTool(max_results=5, rate_limit=2.0)
+
+        progress(0.90, desc="Returning Results..")
         results = web_search_tool(query)
         return results
         
@@ -252,13 +268,14 @@ def run_agent(
     video_prompt_param="", 
     video_duration_param=4.0, 
     video_steps_param=20, 
-    video_guidance_param=3.0, 
+    video_guidance_param=3.0,
+    progress: gr.Progress = gr.Progress(),
 ):
     global image_output, video_output
     image_output = None
     video_output = None
 
-    yield None, None, "⏳ Jerry is thinking... please wait"
+    progress(0.05, desc="Jerry is thinking …")
 
     try:
 
@@ -285,6 +302,8 @@ def run_agent(
             }
         )
 
+        progress(1, desc="Done…")
+
         yield image_output, video_output, str(response)
 
     except Exception as e:
@@ -429,7 +448,7 @@ with gr.Blocks(title="Jerry AI Assistant") as demo:
             ],
             outputs=[image_output_display, gr.Video(visible=False), agent_response]
         )
-
+        
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",