Update app.py

app.py

@@ -1,113 +1,86 @@
 import gradio as gr
-import sqlite3
-import threading
-import time
 import os
 import shutil
 from gradio_client import Client, handle_file
-# import spaces
+from smolagents import Tool, CodeAgent, HfApiModel
 
-# Database setup
-conn = sqlite3.connect('/tmp/jobs.db', check_same_thread=False)
-c = conn.cursor()
-c.execute('''CREATE TABLE IF NOT EXISTS jobs
-             (id INTEGER PRIMARY KEY, image_path TEXT, job_id TEXT, status TEXT, output_path TEXT)''')
-conn.commit()
+# import spaces - if using ZeroGPU
 
-# TRELLIS API client
-trellis_client = Client("jkorstad/TRELLIS")
+# Define tools from Spaces
+spaces = [
+    {"repo_id": "black-forest-labs/FLUX.1-schnell", 
+     "name": "image_generator", 
+     "description": "Generate an image from a prompt"},
+    
+    {"repo_id": "jamesliu1217/EasyControl_Ghibli", 
+     "name": "Ghibli_style_Image_control", 
+     "description": "Create Ghibli style image"},
+]
 
+tools = []
+for space in spaces:
+    # Access repo_id, name, and description
+    repo_id = space['repo_id']
+    name = space.get('name', repo_id)  # Use repo_id as name if not specified
+    description = space.get('description', '')  # Use empty string if not specified
 
-# @spaces.GPU
-# Processing logic with three-step TRELLIS workflow
-def process_job(job_id):
+    # Create Tool instance
+    tool = Tool.from_space(repo_id, name=name, description=description)
+    tools.append(tool)
+
+# Define a custom tool
+class CustomTool(Tool):
+    name = "custom_tool"
+    description = "A custom tool that processes input text"
+    inputs = {"input": {"type": "string", "description": "Some input text to process"}}
+    output_type = "string"
+    def forward(self, input: str):
+        return f"Processed: {input}"
+
+tools.append(CustomTool())
+
+
+# Initialize the model
+model = HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct")
+
+# Create the agent
+agent = CodeAgent(tools=tools, model=model)
+
+# Function to run the agent and return the image path
+def generate_and_transform(prompt):
+    result = agent.run(prompt)
+    
+    if isinstance(result, str):  # Assuming result is a file path
+        # Copy the temporary file to a permanent location
+        permanent_path = "ghibli_output.webp"
+        shutil.copy(result, permanent_path)
+        return permanent_path
+    else:
+        raise ValueError("Unexpected result type from agent")
+
+# Gradio interface function
+def gradio_interface(prompt):
     try:
-        # Get the uploaded image path
-        c.execute("SELECT image_path FROM jobs WHERE id=?", (job_id,))
-        image_path = c.fetchone()[0]
-        
-        # Step 1: Preprocess the image
-        c.execute("UPDATE jobs SET status='preprocessing' WHERE id=?", (job_id,))
-        conn.commit()
-        preprocessed_image = trellis_client.predict(
-            image=handle_file(image_path),
-            api_name="/preprocess_image"
-        )
-        
-        # Step 2: Generate 3D asset
-        c.execute("UPDATE jobs SET status='generating' WHERE id=?", (job_id,))
-        conn.commit()
-        time.sleep(30)  # Wait between steps; adjust based on observed timing
-        result_3d = trellis_client.predict(
-            image=handle_file(preprocessed_image),  # Use preprocessed image
-            multiimages=[],
-            seed=0,  # Default; could make configurable
-            ss_guidance_strength=7.5,
-            ss_sampling_steps=12,
-            slat_guidance_strength=3,
-            slat_sampling_steps=12,
-            multiimage_algo="stochastic",
-            api_name="/image_to_3d"
-        )
-        video_path = result_3d['video']  # Extract video filepath from dict
-        
-        # Step 3: Extract GLB
-        c.execute("UPDATE jobs SET status='extracting' WHERE id=?", (job_id,))
-        conn.commit()
-        time.sleep(65)  # Wait for 3D processing; adjust as needed
-        glb_result = trellis_client.predict(
-            mesh_simplify=0.95,
-            texture_size=1024,
-            api_name="/extract_glb"
-        )
-        glb_path = glb_result[0]  # First element is the GLB filepath
-        
-        # Move GLB to persistent storage
-        output_path = f'/tmp/outputs/result_{job_id}.glb'
-        os.makedirs('/tmp/outputs', exist_ok=True)
-        shutil.move(glb_path, output_path)
-        
-        # Update job status
-        c.execute("UPDATE jobs SET status='completed', output_path=? WHERE id=?", (output_path, job_id))
-        conn.commit()
-        
+        image_path = generate_and_transform(prompt)
+        return image_path
     except Exception as e:
-        c.execute("UPDATE jobs SET status='failed' WHERE id=?", (job_id,))
-        conn.commit()
-        print(f"Error processing job {job_id}: {e}")
+        return str(e)
 
-# Gradio interface
-def submit_images(files):
-    if not files:
-        return "No files uploaded."
-    for file in files:
-        c.execute("INSERT INTO jobs (status) VALUES ('submitted')")
-        job_id = c.lastrowid
-        conn.commit()
-        image_path = f'/tmp/inputs/input_{job_id}.jpg'
-        os.makedirs('/tmp/inputs', exist_ok=True)
-        shutil.copy(file.name, image_path)
-        c.execute("UPDATE jobs SET image_path=? WHERE id=?", (image_path, job_id))
-        conn.commit()
-        threading.Thread(target=process_job, args=(job_id,), daemon=True).start()
-    return "Jobs submitted. Check the status tab."
+# Create the Gradio app
+with gr.Blocks() as app:
+    gr.Markdown("### Smolagent Image Generator with Ghibli Style")
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="e.g., Generate an image of a dog and then make a Ghibli style version of that image")
+        submit_button = gr.Button("Generate")
+    output_image = gr.Image(label="Generated Image")
+    download_button = gr.File(label="Download Image")
 
-def get_status():
-    c.execute("SELECT id, image_path, status, output_path FROM jobs")
-    return c.fetchall()
+    # Connect the button to the function
+    def on_submit(prompt):
+        image_path = gradio_interface(prompt)
+        return image_path, image_path
 
-with gr.Blocks(title="TRELLIS 3D Generator") as demo:
-    with gr.Tab("Upload"):
-        files_input = gr.File(file_count="multiple", label="Upload Images (JPG/PNG)")
-        submit_btn = gr.Button("Submit")
-        output_msg = gr.Textbox(label="Message")
-        submit_btn.click(fn=submit_images, inputs=files_input, outputs=output_msg)
-    with gr.Tab("Status"):
-        status_table = gr.DataFrame(
-            headers=["ID", "Image Path", "Status", "Output Path"],
-            label="Job Status"
-        )
-        refresh_btn = gr.Button("Refresh")
-        refresh_btn.click(fn=get_status, inputs=None, outputs=status_table)
+    submit_button.click(on_submit, inputs=prompt_input, outputs=[output_image, download_button])
 
-demo.launch()
+# Launch the app
+app.launch()