Replicate Fine-Tuning Tutorial

import os
from dotenv import load_dotenv
import replicate

# Load API key
load_dotenv()
REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
assert REPLICATE_API_TOKEN, "Set REPLICATE_API_TOKEN in .env file"
replicate_client = replicate.Client(api_token=REPLICATE_API_TOKEN)

# Example: Ensure training_images.zip exists
ZIP_PATH = "training_images.zip"
assert os.path.isfile(ZIP_PATH), f"Place your zip at {ZIP_PATH}"

try:
    # Try to get current user info through endpoint
    import requests
    headers = {"Authorization": f"Token {REPLICATE_API_TOKEN}"}
    response = requests.get("https://api.replicate.com/v1/account", headers=headers)
    if response.status_code == 200:
        user_data = response.json()
        username = user_data.get("username", user_data.get("name"))
    else:
        print(f"\n API request failed with status: {response.status_code}")
        username = "your_username"  # Set a fallback username
except Exception as e:
    print(f"Error: {e}")
    username = "your_username"  # Set a fallback username

# Set up destination name with your username
flux_destination = f"{username}/flux-my-finetune-v2"

print(f"\n📋 FLUX destination will be: {flux_destination}")

# Create the destination model first if it doesn't exist
try:
    print(f"Creating model: {flux_destination}")
    model = replicate_client.models.create(
        owner=username,
        name="flux-my-finetune-v2",
        visibility="private",
        description="Fine-tuned FLUX model",
        hardware="gpu-t4" 
    )
    print(f"✅ Model created: {model.name}")
except Exception as e:
    print(f"Error: {e}")

try:
    training = replicate_client.trainings.create(
        version="replicate/fast-flux-trainer:f463fbfc97389e10a2f443a8a84b6953b1058eafbf0c9af4d84457ff07cb04db",
        input={
            "input_images": open(ZIP_PATH, "rb"), 
            "trigger_word": "MYTOK",
            "steps": 1000
        },
        destination=flux_destination
    )
    print("✅ Training started successfully!")
    print("Training ID:", training.id)
    print("You can monitor training at:", training.urls.get)
    
except Exception as e:
    print(f"Error {e}")

# Polling loop (example)
import time

# Store training job
training_jobs = []
if 'training' in locals():
    training_jobs.append(training)

for job in training_jobs:
    print(f"Polling {job.id}...")
    while True:
        job = replicate_client.trainings.get(job.id)
        print(f"Status: {job.status}")
        if job.status in ["succeeded", "failed", "canceled"]:
            break
        time.sleep(30)  # Check every 30 seconds
    print(f"Job {job.id} finished with status {job.status}")
    
    if job.status == "succeeded":
        print(f"Training completed! View results at: {job.urls.get}")
        print("Note: To use the trained model, you'll need to:")
        print("1. Go to the training URL above")
        print("2. Create a model from the training results")
        print("3. Use the new model ID for generation")
    else:
        print(f"Training failed or was canceled")

# Generate an image of yourself wearing a space suit using your fine-tuned model
prompt = "A realistic photo of Devon wearing a space suit on Mars, cinematic lighting, astronaut helmet, NASA patch, high detail"
output = replicate.run(
    "devonsunml/flux-my-finetune-v2:1c68fbecae650b8a0dd9e4ea7f38f9c18cba1b9fc4af1077d0a0c9fe403db98a",
    input={
        "prompt": prompt,
        "model": "dev",
        "go_fast": False,
        "lora_scale": 1,
        "megapixels": "1",
        "num_outputs": 1,
        "aspect_ratio": "1:1",
        "output_format": "webp",
        "guidance_scale": 3,
        "output_quality": 80,
        "prompt_strength": 0.8,
        "extra_lora_scale": 1,
        "num_inference_steps": 28
    }
)

# To access the file URL:
img_url = output[0]
print("Image URL:", img_url)

import requests
from IPython.display import Image, display
img_data = requests.get(img_url).content
display(Image(data=img_data, width=512, height=512))