""" title: WAN Image Generator description: Generate images with WAN2.1 model (prunaai/wan-image) hosted on Replicate author: Dean Tjodimedjo version: 1.0.0 license: MIT """ import requests # Used for making HTTP requests to the Replicate API import os # Used to access environment variables for API tokens from pydantic import ( BaseModel, Field, conint, constr, # For string constraints like enum ) import time # Import for time.sleep import random # Imported, though not currently used for explicit random seed generation class Tools: class Valves(BaseModel): REPLICATE_API_TOKEN: str = Field( default="", description="Get your token from https://replicate.com/account", title="Replicate API Token", ) # WAN Image specific parameters, excluding 'prompt' from valves SEED: int = Field( default=-1, # Default to -1 as per WAN Image's API for random seed. description="Random seed. Leave -1 for a random seed.", title="Seed", json_schema_extra={ "minimum": -1, }, ) JUCED: bool = Field( default=False, description="If true, apply a 'juiced' effect.", title="Juiced Effect", ) MEGAPIXELS: conint(ge=1, le=2) = Field( # Constrain to 1 or 2 as per WAN Image default=2, description="Approximate number of megapixels for generated image.", title="Megapixels", json_schema_extra={ "enum": [1, 2], }, ) ASPECT_RATIO: constr(pattern=r"^(1:1|16:9|9:16|4:3|3:4|21:9)$") = Field( default="16:9", description="Aspect ratio for the generated image.", title="Aspect Ratio", json_schema_extra={ "enum": ["1:1", "16:9", "9:16", "4:3", "3:4", "21:9"], }, ) OUTPUT_FORMAT: constr(pattern=r"^(png|jpg|webp)$") = Field( default="jpg", description="Format of the output images.", title="Output Format", json_schema_extra={ "enum": ["png", "jpg", "webp"], }, ) OUTPUT_QUALITY: conint(ge=1, le=100) = Field( default=80, description="Quality when saving the output images, from 1 to 100. Not relevant for .png outputs.", title="Output Quality", json_schema_extra={ "minimum": 1, "maximum": 100, }, ) # Add a timeout valve for polling, default to 300 seconds (5 minutes) # This allows the user to configure how long to wait POLLING_TIMEOUT: conint(ge=30, le=600) = Field( default=300, # 5 minutes description="Maximum time in seconds to wait for image generation.", title="Polling Timeout (seconds)", json_schema_extra={ "minimum": 30, "maximum": 600, }, ) POLLING_INTERVAL: conint(ge=1, le=10) = Field( default=5, # Check every 5 seconds description="Interval in seconds between checking image generation status.", title="Polling Interval (seconds)", json_schema_extra={ "minimum": 1, "maximum": 10, }, ) def __init__(self): self.valves = self.Valves( REPLICATE_API_TOKEN=os.getenv("REPLICATE_API_TOKEN", ""), SEED=int(os.getenv("WAN_IMAGE_SEED", "-1")), JUCED=os.getenv("WAN_IMAGE_JUCED", "False").lower() == "true", MEGAPIXELS=int(os.getenv("WAN_IMAGE_MEGAPIXELS", "2")), ASPECT_RATIO=os.getenv("WAN_IMAGE_ASPECT_RATIO", "16:9"), OUTPUT_FORMAT=os.getenv("WAN_IMAGE_OUTPUT_FORMAT", "jpg"), OUTPUT_QUALITY=int(os.getenv("WAN_IMAGE_OUTPUT_QUALITY", "80")), POLLING_TIMEOUT=int(os.getenv("WAN_IMAGE_POLLING_TIMEOUT", "300")), POLLING_INTERVAL=int(os.getenv("WAN_IMAGE_POLLING_INTERVAL", "5")), ) async def generate_image( self, image_prompt: str, __event_emitter__=None, ) -> str: try: await __event_emitter__( { "type": "status", "data": { "description": "Initiating image generation...", "done": False, }, } ) if not self.valves.REPLICATE_API_TOKEN: raise ValueError( "API token missing. Please set it in the tool's Valves." ) model_identifier = "prunaai/wan-image" initial_prediction_url = "https://api.replicate.com/v1/predictions" headers = { "Authorization": f"Bearer {self.valves.REPLICATE_API_TOKEN}", "Content-Type": "application/json", # REMOVED: "Prefer": "wait" because we will now poll explicitly } input_data = { "prompt": image_prompt, "juiced": self.valves.JUCED, "megapixels": self.valves.MEGAPIXELS, "aspect_ratio": self.valves.ASPECT_RATIO, "output_format": self.valves.OUTPUT_FORMAT, "output_quality": self.valves.OUTPUT_QUALITY, } if self.valves.SEED != -1: input_data["seed"] = self.valves.SEED payload = { "version": model_identifier, "input": input_data, } # 1. Send the initial request to create the prediction response = requests.post( initial_prediction_url, headers=headers, json=payload ) response.raise_for_status() initial_result = response.json() prediction_id = initial_result.get("id") if not prediction_id: raise ValueError("Failed to get prediction ID from Replicate.") prediction_get_url = initial_result.get("urls", {}).get("get") if not prediction_get_url: raise ValueError("Failed to get prediction URL for polling.") start_time = time.time() # 2. Poll for the prediction status while True: if time.time() - start_time > self.valves.POLLING_TIMEOUT: raise TimeoutError( f"Image generation timed out after {self.valves.POLLING_TIMEOUT} seconds." ) # Update status message during polling await __event_emitter__( { "type": "status", "data": { "description": f"Generating image (status: {initial_result.get('status', 'unknown')})...", "done": False, }, } ) poll_response = requests.get(prediction_get_url, headers=headers) poll_response.raise_for_status() result = poll_response.json() # print(f"Replicate Poll Response: {result}") # Debugging line - keep this if you want to see progress in logs status = result.get("status") if status == "succeeded": break # Prediction complete, break the loop elif status == "failed": error_detail = result.get( "error", "No specific error provided by Replicate." ) raise ValueError(f"Replicate prediction failed: {error_detail}") elif status == "canceled": raise ValueError("Replicate prediction was canceled.") elif status in ["starting", "processing", "queued"]: time.sleep( self.valves.POLLING_INTERVAL ) # Wait before polling again else: raise ValueError(f"Unknown Replicate prediction status: {status}") # Prediction succeeded, now process the output if "output" not in result or not result["output"]: raise ValueError( "Failed to generate image. No output URL found after successful prediction." ) if isinstance(result["output"], list): image_urls = result["output"] else: image_urls = [result["output"]] generated_images_markdown = "" for i, image_url in enumerate(image_urls): generated_images_markdown += ( f"\n\n" ) await __event_emitter__( {"type": "status", "data": {"description": "Done!", "done": True}} ) await __event_emitter__( { "type": "message", "data": {"content": generated_images_markdown}, } ) return "✅ Image(s) successfully generated and displayed above!" except requests.exceptions.HTTPError as e: error_msg = f"API Error: {str(e)}" if e.response.status_code == 401: error_msg += "\n🔑 **Fix:** Check your Replicate API token at https://replicate.com/account" elif e.response.status_code == 422: try: replicate_error_details = e.response.json() if "detail" in replicate_error_details: error_msg += ( f"\nReplicate details: {replicate_error_details['detail']}" ) except: pass return f"Tell the user: {error_msg}" except TimeoutError as e: return f"Tell the user: {str(e)}" except Exception as e: return f"Tell the user: An unexpected error occurred: {str(e)}"