from typing import Dict, List, Optional, Any import uuid from pydantic import BaseModel, Field, HttpUrl, root_validator, validator, SecretStr from enum import Enum from datetime import datetime class ImageRequest(BaseModel): prompt: str model: str = "black-forest-labs/flux-dev" enhance_prompt: bool = False provider: Optional[str] = "replicate" class Config: json_schema_extra = { "example": { "prompt": "Create a futuristic cityscape at sunset.", "model": "black-forest-labs/flux-dev", "enhance_prompt": True, } } # Models class TrainingStatus(str, Enum): STARTING = "starting" PROCESSING = "processing" SUCCEEDED = "succeeded" FAILED = "failed" CANCELED = "canceled" class WebhookEvent(str, Enum): START = "start" OUTPUT = "output" LOGS = "logs" COMPLETED = "completed" class AutocaptionOptions(BaseModel): enabled: bool = Field(False, description="Automatically caption images") prefix: Optional[str] = Field(None, description="Text to prepend to captions") suffix: Optional[str] = Field(None, description="Text to routerend to captions") class WandbConfig(BaseModel): api_key: SecretStr = Field(..., description="Weights & Biases API key") project: str = Field("flux_train_replicate", description="W&B project name") entity: Optional[str] = Field(None, description="W&B entity name") run: Optional[str] = Field(None, description="W&B run name") sample_prompts: Optional[str] = Field( None, description="Newline-separated prompts for W&B samples" ) sample_interval: int = Field( 100, description="Step interval for sampling output images" ) save_interval: int = Field( 100, description="Step interval for saving intermediate weights" ) class TrainingInput(BaseModel): destination: str = Field( ..., description="Model destination in format owner/name", pattern=r"^[A-Za-z0-9_\-]+/[A-Za-z0-9_\-]+$", ) model_owner: str = Field(..., description="Owner of the model to train") model_name: str = Field(..., description="Name of the model to train") version_id: str = Field( ..., description="Version ID of the model to train" ) # Optional field with default None # Required training parameters input_images: str = Field( ..., description="URL to zip file containing training data" ) trigger_word: str = Field("TOK", description="Trigger word for training") # Optional training parameters steps: int = Field(1000, ge=3, le=6000, description="Number of training steps") lora_rank: int = Field(16, ge=1, le=128, description="LoRA rank value") hf_repo_id: Optional[str] = Field(None, description="HuggingFace repository ID") hf_token: Optional[SecretStr] = Field(None, description="HuggingFace token") # W&B integration wandb: Optional[WandbConfig] = Field( None, description="Weights & Biases configuration" ) # Autocaption settings autocaption: AutocaptionOptions = Field( default_factory=AutocaptionOptions, description="Auto-captioning configuration" ) # Advanced training parameters learning_rate: float = Field(0.0004, description="Learning rate") batch_size: int = Field(1, description="Batch size") resolution: str = Field( "512,768,1024", description="Image resolutions for training" ) caption_dropout_rate: float = Field( 0.05, ge=0, le=1, description="Caption dropout rate" ) optimizer: str = Field("adamw8bit", description="Optimizer to use") cache_latents_to_disk: bool = Field(False, description="Cache latents to disk") layers_to_optimize_regex: Optional[str] = Field( None, description="Regex for layers to optimize" ) gradient_checkpointing: bool = Field( False, description="Enable gradient checkpointing" ) # Skip training option skip_training_and_use_pretrained_hf_lora_url: Optional[HttpUrl] = Field( None, description="URL to pretrained HF LoRA to use instead of training" ) # Webhook configuration webhook: Optional[HttpUrl] = Field( None, description="Webhook URL for notifications" ) webhook_events_filter: Optional[List[WebhookEvent]] = Field( None, description="Events that trigger webhooks" ) def to_replicate_payload(self) -> Dict[str, Any]: """Convert to Replicate API payload format""" # Build inputs dictionary with all the training parameters input_dict = { "input_images": str(self.input_images), "steps": self.steps, "lora_rank": self.lora_rank, "learning_rate": self.learning_rate, "batch_size": self.batch_size, "resolution": self.resolution, "caption_dropout_rate": self.caption_dropout_rate, "optimizer": self.optimizer, "cache_latents_to_disk": self.cache_latents_to_disk, "gradient_checkpointing": self.gradient_checkpointing, "autocaption": self.autocaption.enabled, } # Add optional parameters if they exist if self.autocaption.prefix: input_dict["autocaption_prefix"] = self.autocaption.prefix if self.autocaption.suffix: input_dict["autocaption_suffix"] = self.autocaption.suffix if self.layers_to_optimize_regex: input_dict["layers_to_optimize_regex"] = self.layers_to_optimize_regex if self.hf_repo_id: input_dict["hf_repo_id"] = self.hf_repo_id if self.hf_token: input_dict["hf_token"] = self.hf_token.get_secret_value() if self.skip_training_and_use_pretrained_hf_lora_url: input_dict["skip_training_and_use_pretrained_hf_lora_url"] = str( self.skip_training_and_use_pretrained_hf_lora_url ) # Handle W&B configuration if self.wandb: input_dict["wandb_api_key"] = self.wandb.api_key.get_secret_value() input_dict["wandb_project"] = self.wandb.project if self.wandb.entity: input_dict["wandb_entity"] = self.wandb.entity if self.wandb.run: input_dict["wandb_run"] = self.wandb.run if self.wandb.sample_prompts: input_dict["wandb_sample_prompts"] = self.wandb.sample_prompts input_dict["wandb_sample_interval"] = self.wandb.sample_interval input_dict["wandb_save_interval"] = self.wandb.save_interval # Build the final payload payload = {"destination": self.destination, "input": input_dict} # Add webhook configuration if present if self.webhook: payload["webhook"] = str(self.webhook) if self.webhook_events_filter: payload["webhook_events_filter"] = [ event.value for event in self.webhook_events_filter ] return payload class TrainingInputExtended(TrainingInput): image_urls: Optional[List[HttpUrl]] = Field( None, description="List of image URLs to be zipped for training" ) # Override the validator to allow either input_images or image_urls @root_validator(pre=True) def validate_input_source(cls, values): # Check if at least one of input_images or image_urls is provided if not values.get("input_images") and not values.get("image_urls"): raise ValueError("Either input_images or image_urls must be provided") return values class TrainingResponse(BaseModel): id: str model: str version: str input: Dict[str, Any] logs: Optional[str] = None error: Optional[str] = None status: TrainingStatus created_at: datetime started_at: Optional[datetime] = None completed_at: Optional[datetime] = None output: Optional[Dict[str, Any]] = None metrics: Optional[Dict[str, float]] = None urls: Dict[str, str] class PaginatedTrainings(BaseModel): next: Optional[str] = None previous: Optional[str] = None results: List[TrainingResponse] class HardwareInfo(BaseModel): name: str sku: str class ZipRequest(BaseModel): image_urls: List[str] = Field( ..., description="List of image URLs to include in the zip" ) folder_id: Optional[str] = Field( default_factory=lambda: str(uuid.uuid4()), description="Unique folder ID" ) caption: str = Field( "Describe this image", description="Caption instruction for image descriptions" ) focus_entity: Optional[str] = Field( None, description="Focus instruction for image descriptions" ) trigger_word: Optional[str] = Field(None, description="Trigger word for LoRA") image_urls_captions: Optional[List[str]] = None