app.py

import spaces
import os
import gradio as gr
from smolagents import (
    tool,
    CodeAgent,
    DuckDuckGoSearchTool,
    InferenceClientModel,
    FinalAnswerTool,
    LocalPythonExecutor,
)
from huggingface_hub import InferenceClient
import tempfile
from PIL import Image

# ==========================================
# 🛠️ HELPER FUNCTIONS
# ==========================================

def pil_to_tempfile(image):
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    tmp_path = tmp.name
    tmp.close()
    image.save(tmp_path, format="PNG")
    return tmp_path

def aligned_num_frames(duration, fps=16):
    n = int(duration * fps)
    return ((n - 1) // 4) * 4 + 1

def align(x, base=16):
    return (x // base) * base

image_output = None    
video_output = None    

# ==========================================
# 🧰 AGENT TOOLS DEFINITION
# ==========================================

@tool
def video_tool(
    video_image_input: Image.Image, 
    prompt: str = "high quality, detailed, sharp, cinematic", 
    duration: float = 4, 
    steps: int = 20, 
    guidance: float = 3.0,
    hf_visitor_token: str = None
) -> str:
    """
    Generates a video from a starting image using Wan 2.1.
    Args:
        video_image_input (Image.Image): The source image to be animated.
        prompt (str): The prompt for video generation.
        duration (float): Duration in seconds.
        steps (int): Number of inference steps.
        guidance (float): Guidance scale.
        hf_visitor_token (str): The visitor's authenticated OAuth token.
    Returns:
        str: A confirmation message.
    """
    global video_output
    try:
        dynamic_video_client = InferenceClient(
            model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
            provider="fal-ai",
            token=hf_visitor_token,
        )

        MAX_RES = 640
        w, h = video_image_input.size
        scale = min(MAX_RES / w, MAX_RES / h, 1)
        new_w = align(int(w * scale))
        new_h = align(int(h * scale))
        
        image = video_image_input.resize((new_w, new_h), Image.LANCZOS)
        FPS = 16 
        num_frames = aligned_num_frames(duration, FPS)
        
        video_bytes = dynamic_video_client.image_to_video(
            image=image, 
            width=new_w,
            height=new_h,
            prompt=prompt,
            negative_prompt="low quality, deformed, grainy, blurry, pixelated",
            num_frames=num_frames,
            num_inference_steps=steps,  
            guidance_scale=guidance, 
            decode_chunk_size=8, 
        )        
        
        out = tempfile.mktemp(suffix=".mp4")
        with open(out, "wb") as f:
            f.write(video_bytes)      
        video_output = out
        return "Video successfully generated and stored for Gradio UI."
    except Exception as e:
        video_output = None
        return f"Video generation failed: {e}"
      
@tool
def nsfw_detection_tool(nsfw_detection_input: Image.Image, hf_visitor_token: str = None) -> str:
    """
    Suitable for filtering through score explicit or inappropriate content in images.
    Args:
        nsfw_detection_input (Image.Image): The image to check.
        hf_visitor_token (str): The visitor's authenticated OAuth token.
    Returns:
        str: Highest score result.
    """
    try:
        dynamic_nsfw_client = InferenceClient(token=hf_visitor_token)
        tmp_path = pil_to_tempfile(nsfw_detection_input)
        outputs = dynamic_nsfw_client.image_classification(
            tmp_path,
            model="Falconsai/nsfw_image_detection"
        )        
        os.remove(tmp_path)
        top_result = max(outputs, key=lambda x: x.score)
        return f"Verdict: {top_result.label.upper()}\nConfidence: {top_result.score:.2%}"
    except Exception as e:
        return f"NSFW detection failed: {e}"

@tool
def image_tool(image_prompt_param: str, hf_visitor_token: str = None) -> str:
    """
    Generate an image from text using SD3-Medium.
    Args:
        image_prompt_param (str): image description.
        hf_visitor_token (str): The visitor's authenticated OAuth token.
    Returns:
        str: A confirmation message.
    """    
    global image_output  
    try:
        dynamic_img_client = InferenceClient(
            model="stabilityai/stable-diffusion-3-medium",
            token=hf_visitor_token
        )
        image = dynamic_img_client.text_to_image(
            prompt=image_prompt_param,
            negative_prompt="low quality, deformed",
            guidance_scale=7.0,
            num_inference_steps=28,
            width=832,
            height=1280
        )
        image_output = image
        return "Image successfully generated and stored for Gradio UI."
    except Exception as e:
        image_output = None
        return f"Image generation failed: {e}"

@tool
def search_tool(query: str) -> str:
    """
    Search the web and return the most relevant results.
    Args:
        query (str): The search query.
    Returns:
        str: The search results.
    """
    try:
        web_search_tool = DuckDuckGoSearchTool(max_results=5, rate_limit=2.0)
        return web_search_tool(query)
    except Exception as e:
        return f"Search failed: {e}" 

@tool
def sentiment_analysis_tool(text: str, hf_visitor_token: str = None) -> str:
    """
    Analyzes the raw sentiment label and confidence score of a text string.
    Args:
        text (str): The text or prompt to evaluate.
        hf_visitor_token (str): The visitor's authenticated OAuth token.
    Returns:
        str: The classification result.
    """
    try:
        client = InferenceClient(token=hf_visitor_token)

        outputs = client.text_classification(
            text, 
            model="distilbert/distilbert-base-uncased-finetuned-sst-2-english"
        )
        
        top_result = max(outputs, key=lambda x: x.score)
        return f"Sentiment: {top_result.label.upper()}\nConfidence: {top_result.score:.2%}"
    except Exception as e:
        return f"Sentiment processing failed: {e}"   
        
# ==========================================
# 🤖 AGENT CONFIGURATION
# ==========================================

final_answer = FinalAnswerTool()

executor = LocalPythonExecutor(
    additional_authorized_imports=[],
    timeout_seconds=300
) 

agent = CodeAgent(
    model=None,
    tools=[video_tool, image_tool, nsfw_detection_tool, search_tool, sentiment_analysis_tool, final_answer],
    max_steps=6,
    verbosity_level=2,
    executor=executor,
)

agent.prompt_templates["system_prompt"] += """
    You are a tool calling agent.
    You have access to these tools: 
    - search_tool(query: str) -> str: Search the web and return the most relevant results.
    - video_tool(video_image_input: Image.Image, prompt: str, duration: float, steps: int, guidance: float) -> str: Generates a video from a starting image.
    - image_tool(image_prompt_param: str) -> str: Generate an image from a text prompt.
    - nsfw_detection_tool(nsfw_detection_input: Image.Image) -> str: Check an image file for explicit content score metrics.
    - sentiment_analysis_tool(text: str) -> str: Evaluates raw sentiment category and confidence percentage.

    CRITICAL INSTRUCTIONS:
    - When sentiment analysis is requested, or if the user prompt starts with "Analyze the sentiment:", you MUST execute the sentiment_analysis_tool on the text.
    - When generating a video, to save time the image must not use the nsfw_detection_tool first.
    - You must construct a well-formatted human-readable answer.
    - You must introduce yourself as Jerry and greet the user warmly in the final answer text.
    - You must try to include clear breaks like newlines, bullets, numbering, and proper punctuation.
    - You must use this answer in final_answer.
"""

# ==========================================
# 🚀 GRADIO APPLICATION RUNTIME
# ==========================================

@spaces.GPU(duration=500)
def run_agent(
    query,
    image_prompt_param,
    nsfw_detection_input, 
    video_image_input, 
    video_prompt_param, 
    video_duration_param, 
    video_steps_param, 
    video_guidance_param,
    progress=gr.Progress(),
    oauth_token: gr.OAuthToken | None = None,
):
    global image_output, video_output
    image_output = None
    video_output = None

    if oauth_token is None:
        yield None, None, "⚠️ Please log in using the Hugging Face button to use Jerry under your own quota!"
        return

    visitor_token = oauth_token.token
    progress(0, desc="Jerry is working...")

    try:
        if video_image_input is not None or (video_prompt_param and video_prompt_param.strip()):
            actual_query = "Generate a video"
            progress(0.05, desc="Generating video...")
        elif image_prompt_param and image_prompt_param.strip():
            actual_query = "Generate an image"
            progress(0.05, desc="Generating image...")
        elif nsfw_detection_input is not None:
            actual_query = "Check this image for NSFW content"
            progress(0.05, desc="Checking NSFW context...")
        elif query and query.strip():
            actual_query = query
            progress(0.05, desc="Thinking...")
        else:
            actual_query = "What can I help you with?"

        agent.model = InferenceClientModel(
            model_id="Qwen/Qwen2.5-72B-Instruct",
            token=visitor_token,
            max_tokens=2096,
            temperature=0.6,
        )

        response = agent.run(
            actual_query,
            additional_args={
                "image_prompt_param": image_prompt_param,
                "nsfw_detection_input": nsfw_detection_input, 
                "video_image_input": video_image_input,
                "prompt": video_prompt_param,
                "duration": video_duration_param,
                "steps": video_steps_param,
                "guidance": video_guidance_param,
                "hf_visitor_token": visitor_token,
            }
        )

        progress(1, desc="Done!")
        yield image_output, video_output, str(response)

    except Exception as e:
        yield None, None, f"❌ Agent Error: {str(e)}"

# ==========================================
# 🎨 GRADIO INTERFACE LAYOUT
# ==========================================

with gr.Blocks(title="Jerry AI Assistant") as demo:
    gr.Markdown("# 🤖 Jerry - Your AI Assistant")
    gr.LoginButton()

    agent_response = gr.Textbox(label="Response", lines=5, interactive=False)

    with gr.Tab("💬 Chat & Sentiment"):
        query_chat = gr.Textbox(lines=3, label="Ask me anything or paste text for sentiment analysis...")
        run_chat_btn = gr.Button("🚀 Run", variant="primary")

        run_chat_btn.click(
            fn=run_agent,
            inputs=[
                query_chat,
                gr.State(""),      
                gr.State(None),    
                gr.State(None),    
                gr.State(""),      
                gr.State(4.0),     
                gr.State(20),      
                gr.State(3.0),     
            ],
            outputs=[gr.Image(visible=False), gr.Video(visible=False), agent_response],
        )

    with gr.Tab("🎬 Video Tools"):
        with gr.Row():
            with gr.Column():
                video_image_input = gr.Image(type="pil", label="Input Image")
                prompt_txt = gr.Textbox(lines=3, label="Prompt")
                with gr.Accordion("Settings", open=False):
                    dur_slider = gr.Slider(1, 4, value=4, step=0.1, label="Duration")
                    step_slider = gr.Slider(4, 35, value=20, step=1, label="Steps")
                    guidance_slider = gr.Slider(1.0, 6.0, value=3.0, step=0.1, label="Guidance Strength")
                gen_btn = gr.Button("Generate Video", variant="primary")
            with gr.Column():
                output_vid = gr.Video(label="Generated Video")

        gen_btn.click(
            fn=run_agent,
            inputs=[
                gr.State(""),        
                gr.State(""),        
                gr.State(None),      
                video_image_input,   
                prompt_txt,                  
                dur_slider,                  
                step_slider,              
                guidance_slider,                                                                                
            ],
            outputs=[gr.Image(visible=False), output_vid, agent_response],
        )

    with gr.Tab("🎨 Image Tools"):
        with gr.Row():
            with gr.Column():
                nsfw_detection_input = gr.Image(type="pil", label="Upload for NSFW Check")
                check_nsfw_btn = gr.Button("🔍 Check NSFW")
                query_img = gr.Textbox(lines=2, label="Image generation prompt")
                run_img_btn = gr.Button("🎨 Generate Image", variant="primary")
            with gr.Column():
                image_output_display = gr.Image(label="Generated Image")

        check_nsfw_btn.click(
            fn=run_agent,
            inputs=[
                gr.State(""),        
                gr.State(""),        
                nsfw_detection_input,         
                gr.State(None),      
                gr.State(""),        
                gr.State(4.0),       
                gr.State(20),        
                gr.State(3.0),       
            ],
            outputs=[gr.Image(visible=False), gr.Video(visible=False), agent_response],
        )

        run_img_btn.click(
            fn=run_agent,
            inputs=[
                gr.State(""),        
                query_img,
                gr.State(None),      
                gr.State(None),      
                gr.State(""),        
                gr.State(4.0),       
                gr.State(20),        
                gr.State(3.0),       
            ],
            outputs=[image_output_display, gr.Video(visible=False), agent_response],
        )

    gr.Examples(
        examples=[
            ["A cyberpunk cat with neon glowing eyes"],
            ["A serene Japanese garden with cherry blossoms"],
        ],
        inputs=[query_img],
        label="💡 Image Generation Examples:"
    )

    gr.Examples(
        examples=[
            ["The people all raise a glass and cheer"],
            ["A beautiful cinematic timelapse of a sunrise over mountains"],
        ],
        inputs=[prompt_txt],
        label="💡 Video Prompts:"
    )

    gr.Examples(
        examples=[
            ["How do i cook a curry quickly"],
            ["Analyze the sentiment: This is terrible service"],
        ],
        inputs=[query_chat],
        label="💡 Chat & Sentiment Examples:"
    )
    
if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, theme=gr.themes.Soft())