mcp is alive and well!

.gitignore

@@ -1,2 +1,3 @@
 __pycache__
-.venv
+.venv
+.gradio

README.md

@@ -0,0 +1,8 @@
+
+
+## [Modal Flux Fintune Tutorial](https://modal.com/docs/examples/diffusers_lora_finetune)
+
+## Setup 
+```
+uv pip install modal
+```

app.py

@@ -0,0 +1,303 @@
+## Gradio MCP server that launches modal finetune
+
+import gradio as gr
+import requests
+import json
+import time
+from typing import Optional, Dict, Any
+
+# Configuration - Update these URLs to match your deployed Modal app
+MODAL_BASE_URL = "https://stillerman--jason-lora-flux"  # Update with your actual Modal app URL
+START_TRAINING_URL = f"{MODAL_BASE_URL}-api-start-training.modal.run"
+JOB_STATUS_URL = f"{MODAL_BASE_URL}-api-job-status.modal.run"
+
+def start_training(
+    dataset_id: str,
+    hf_token: str,
+    output_repo: str,
+    instance_name: Optional[str] = None,
+    class_name: Optional[str] = None,
+    max_train_steps: int = 500
+) -> tuple[str, str]:
+    """
+    Start a LoRA training job for Flux image generation model.
+    
+    This function initiates a LoRA (Low-Rank Adaptation) training job on a dataset of images.
+    It sends a request to a Modal API endpoint to start the training process.
+    
+    Parameters:
+    - dataset_id (str, required): The HuggingFace dataset ID containing training images, format: "username/dataset-name"
+    - hf_token (str, required): HuggingFace access token with read permissions, format: "hf_xxxxxxxxxxxx"
+    - output_repo (str, required): HuggingFace repository where trained LoRA will be uploaded, format: "username/repo-name"
+    - instance_name (str, optional): Name of the subject being trained (e.g., 'Fluffy', 'MyDog', 'John')
+    - class_name (str, optional): Class category of the subject (e.g., 'person', 'dog', 'cat', 'building')
+    - max_train_steps (int, optional): Number of training steps, range 100-2000, default 500
+    
+    Returns:
+    - tuple[str, str]: (status_message, job_id)
+      - status_message: Human-readable status with training details or error message
+      - job_id: Unique identifier for the training job, empty string if failed
+    
+    Example usage:
+    status, job_id = start_training(
+        dataset_id="myuser/dog-photos",
+        hf_token="hf_abcdef123456",
+        output_repo="myuser/my-dog-lora",
+        instance_name="Fluffy",
+        class_name="dog", 
+        max_train_steps=500
+    )
+    """
+    
+    if not dataset_id or not hf_token or not output_repo:
+        return "❌ Error: Dataset ID, HuggingFace token, and output repo are required", ""
+    
+    payload = {
+        "dataset_id": dataset_id,
+        "hf_token": hf_token,
+        "output_repo": output_repo,
+        "max_train_steps": max_train_steps
+    }
+    
+    # Add optional parameters if provided
+    if instance_name and instance_name.strip():
+        payload["instance_name"] = instance_name.strip()
+    if class_name and class_name.strip():
+        payload["class_name"] = class_name.strip()
+    
+    try:
+        response = requests.post(
+            START_TRAINING_URL,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=30
+        )
+        
+        if response.status_code == 200:
+            result = response.json()
+            if result.get("status") == "started":
+                job_id = result.get("job_id", "")
+                message = f"✅ Training started successfully!\n\n"
+                message += f"**Job ID:** `{job_id}`\n"
+                message += f"**Dataset:** {dataset_id}\n"
+                message += f"**Output Repo:** {output_repo}\n"
+                message += f"**Training Steps:** {max_train_steps}\n\n"
+                message += "Copy the Job ID to check status below."
+                return message, job_id
+            else:
+                return f"❌ Error: {result.get('message', 'Unknown error')}", ""
+        else:
+            return f"❌ HTTP Error {response.status_code}: {response.text}", ""
+            
+    except requests.exceptions.Timeout:
+        return "❌ Error: Request timed out. The service might be starting up.", ""
+    except requests.exceptions.RequestException as e:
+        return f"❌ Error: Failed to connect to training service: {str(e)}", ""
+    except json.JSONDecodeError:
+        return "❌ Error: Invalid response from server", ""
+
+def check_job_status(job_id: str) -> str:
+    """
+    Check the current status of a LoRA training job.
+    
+    This function queries the Modal API to get the current status of a training job
+    using its unique job ID. It returns detailed information about the job progress.
+    
+    Parameters:
+    - job_id (str, required): The unique job identifier returned from start_training function
+    
+    Returns:
+    - str: Detailed status message containing:
+      - Job status (completed, running, failed, error)
+      - Training results if completed (dataset used, steps completed, training prompt)
+      - Error messages if failed
+      - Progress information if still running
+    
+    Possible status values:
+    - "completed": Training finished successfully, LoRA model is ready
+    - "running": Training is still in progress
+    - "failed": Training failed due to an error
+    - "error": System error occurred
+    
+    Example usage:
+    status_info = check_job_status("job_12345abcdef")
+    """
+    
+    if not job_id or not job_id.strip():
+        return "❌ Error: Job ID is required"
+    
+    try:
+        response = requests.get(
+            JOB_STATUS_URL,
+            params={"job_id": job_id.strip()},
+            timeout=10
+        )
+        
+        if response.status_code == 200:
+            result = response.json()
+            status = result.get("status", "unknown")
+            
+            if status == "completed":
+                message = "🎉 **Training Completed!**\n\n"
+                training_result = result.get("result", {})
+                if isinstance(training_result, dict):
+                    message += f"**Status:** {training_result.get('status', 'completed')}\n"
+                    message += f"**Message:** {training_result.get('message', 'Training finished')}\n"
+                    if training_result.get('dataset_used'):
+                        message += f"**Dataset Used:** {training_result['dataset_used']}\n"
+                    if training_result.get('training_steps'):
+                        message += f"**Training Steps:** {training_result['training_steps']}\n"
+                    if training_result.get('training_prompt'):
+                        message += f"**Training Prompt:** {training_result['training_prompt']}\n"
+                else:
+                    message += f"**Result:** {training_result}"
+                return message
+                
+            elif status == "running":
+                return f"🔄 **Training in Progress**\n\nThe training job is still running. Check back in a few minutes."
+                
+            elif status == "failed":
+                error_msg = result.get("message", "Training failed with unknown error")
+                return f"❌ **Training Failed**\n\n**Error:** {error_msg}"
+                
+            elif status == "error":
+                error_msg = result.get("message", "Unknown error occurred")
+                return f"❌ **Error**\n\n**Message:** {error_msg}"
+                
+            else:
+                return f"❓ **Unknown Status**\n\n**Status:** {status}\n**Response:** {json.dumps(result, indent=2)}"
+                
+        else:
+            return f"❌ HTTP Error {response.status_code}: {response.text}"
+            
+    except requests.exceptions.Timeout:
+        return "❌ Error: Request timed out"
+    except requests.exceptions.RequestException as e:
+        return f"❌ Error: Failed to connect to status service: {str(e)}"
+    except json.JSONDecodeError:
+        return "❌ Error: Invalid response from server"
+
+def check_and_update_status(job_id: str) -> str:
+    """
+    Wrapper function to check job status for Gradio interface.
+    
+    This is a simple wrapper around check_job_status that provides the same functionality
+    but is specifically designed for use with Gradio button callbacks.
+    
+    Parameters:
+    - job_id (str, required): The unique job identifier from training
+    
+    Returns:
+    - str: Status message from check_job_status function
+    
+    Example usage:
+    status = check_and_update_status("job_12345abcdef")
+    """
+    return check_job_status(job_id)
+
+# Create simplified single-page Gradio interface
+with gr.Blocks(title="FluxFoundry LoRA Training", theme=gr.themes.Soft()) as app:
+    gr.Markdown("""
+    # 🎨 FluxFoundry LoRA Training
+    
+    Train custom LoRA models for Flux image generation and check training status.
+    """)
+    
+    # Training Section
+    gr.Markdown("## 🚀 Start Training")
+    
+    with gr.Row():
+        with gr.Column():
+            dataset_id = gr.Textbox(
+                label="HuggingFace Dataset ID",
+                placeholder="username/dataset-name",
+                info="The HuggingFace dataset containing your training images"
+            )
+            hf_token = gr.Textbox(
+                label="HuggingFace Token",
+                placeholder="hf_...",
+                type="password",
+                info="Your HuggingFace access token with read permissions"
+            )
+            output_repo = gr.Textbox(
+                label="Output Repository",
+                placeholder="username/my-lora-model",
+                info="HuggingFace repository where the trained LoRA will be uploaded"
+            )
+        
+        with gr.Column():
+            instance_name = gr.Textbox(
+                label="Instance Name (Optional)",
+                placeholder="subject",
+                info="Name of the subject being trained (e.g., 'Fluffy', 'MyDog')"
+            )
+            class_name = gr.Textbox(
+                label="Class Name (Optional)",
+                placeholder="person",
+                info="Class of the subject (e.g., 'person', 'dog', 'cat')"
+            )
+            max_train_steps = gr.Slider(
+                minimum=100,
+                maximum=2000,
+                value=500,
+                step=50,
+                label="Max Training Steps",
+                info="Number of training steps (more steps = longer training)"
+            )
+    
+    start_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
+    
+    with gr.Row():
+        training_output = gr.Markdown(label="Training Status")
+        job_id_output = gr.Textbox(
+            label="Job ID",
+            placeholder="Copy this ID to check status",
+            interactive=False
+        )
+    
+    start_btn.click(
+        fn=start_training,
+        inputs=[dataset_id, hf_token, output_repo, instance_name, class_name, max_train_steps],
+        outputs=[training_output, job_id_output]
+    )
+    
+    # Status Section
+    gr.Markdown("## 📊 Check Status")
+    
+    job_id_input = gr.Textbox(
+        label="Job ID",
+        placeholder="Paste your job ID here",
+        info="The Job ID returned when you started training"
+    )
+    
+    with gr.Row():
+        status_btn = gr.Button("📊 Check Status", variant="secondary")
+        refresh_btn = gr.Button("🔄 Refresh", variant="secondary")
+    
+    status_output = gr.Markdown(label="Job Status")
+    
+    status_btn.click(
+        fn=check_and_update_status,
+        inputs=[job_id_input],
+        outputs=[status_output]
+    )
+    
+    refresh_btn.click(
+        fn=check_and_update_status,
+        inputs=[job_id_input],
+        outputs=[status_output]
+    )
+
+if __name__ == "__main__":
+    print("🎨 Starting FluxFoundry Training Interface...")
+    print(f"📡 Modal API Base URL: {MODAL_BASE_URL}")
+    print("⚠️  Make sure to update the MODAL_BASE_URL in the code with your actual Modal deployment URL")
+    
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True,
+        show_error=True,
+        mcp_server=True
+    )
+

diffusers_lora_finetune.py

@@ -34,6 +34,7 @@
 
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Optional
 
 import modal
 
@@ -52,11 +53,12 @@ app = modal.App(name="jason-lora-flux")
 
 image = modal.Image.debian_slim(python_version="3.10").pip_install(
     "accelerate==0.31.0",
-    "datasets~=2.13.0",
+    "datasets==3.6.0",
+    "pillow",
     "fastapi[standard]==0.115.4",
     "ftfy~=6.1.0",
     "gradio~=5.5.0",
-    "huggingface-hub==0.26.2",
+    "huggingface-hub==0.32.4",
     "hf_transfer==0.1.8",
     "numpy<2",
     "peft==0.11.1",
@@ -184,6 +186,325 @@ def load_images(image_urls: list[str]) -> Path:
     return img_path
 
 
+def load_images_from_hf_dataset(dataset_id: str, hf_token: str) -> Path:
+    """Load images from a HuggingFace dataset."""
+    import PIL.Image
+    from datasets import load_dataset
+
+    img_path = Path("/img")
+    img_path.mkdir(parents=True, exist_ok=True)
+    
+    # Load dataset from HuggingFace
+    dataset = load_dataset(dataset_id, token=hf_token, split="train")
+    
+    for ii, example in enumerate(dataset):
+        # Assume the dataset has an 'image' column
+        if 'image' in example:
+            image = example['image']
+            if isinstance(image, PIL.Image.Image):
+                image.save(img_path / f"{ii}.png")
+            else:
+                # Handle other image formats
+                pil_image = PIL.Image.open(image)
+                pil_image.save(img_path / f"{ii}.png")
+        else:
+            print(f"Warning: No 'image' field found in dataset example {ii}")
+    
+    print(f"{len(dataset)} images loaded from HuggingFace dataset")
+    return img_path
+
+
+# ## Stateless API Training Function
+
+@dataclass
+class APITrainConfig:
+    """Configuration for the API training function."""
+    
+    # Basic model info
+    model_name: str = "black-forest-labs/FLUX.1-dev"
+    
+    # Training prompt components
+    instance_name: str = "subject"
+    class_name: str = "person"
+    prefix: str = "a photo of"
+    postfix: str = ""
+    
+    # Training hyperparameters
+    resolution: int = 512
+    train_batch_size: int = 3
+    rank: int = 16  # lora rank
+    gradient_accumulation_steps: int = 1
+    learning_rate: float = 4e-4
+    lr_scheduler: str = "constant"
+    lr_warmup_steps: int = 0
+    max_train_steps: int = 500
+    checkpointing_steps: int = 1000
+    seed: int = 117
+
+
+@app.function(
+    image=image,
+    gpu="A100-80GB",  # fine-tuning is VRAM-heavy and requires a high-VRAM GPU
+    timeout=3600,  # 60 minutes
+)
+def train_lora_stateless(
+    dataset_id: str,
+    hf_token: str,
+    output_repo: str,
+    instance_name: Optional[str] = None,
+    class_name: Optional[str] = None,
+    max_train_steps: int = 500,
+):
+    """
+    Stateless LoRA training function that reads from HF dataset and uploads to HF repo.
+    
+    Args:
+        dataset_id: HuggingFace dataset ID (e.g., "username/dataset-name")
+        hf_token: HuggingFace API token
+        output_repo: HuggingFace repository to upload the trained LoRA to
+        instance_name: Name of the subject (optional, defaults to "subject")
+        class_name: Class of the subject (optional, defaults to "person")
+        max_train_steps: Number of training steps
+    """
+    import subprocess
+    import tempfile
+    from pathlib import Path
+    
+    import torch
+    from accelerate.utils import write_basic_config
+    from diffusers import DiffusionPipeline
+    from huggingface_hub import snapshot_download, upload_folder, login, create_repo
+    
+    # Login to HuggingFace
+    login(token=hf_token)
+    
+    # Create temporary directories
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_path = Path(temp_dir)
+        model_dir = temp_path / "model"
+        output_dir = temp_path / "output"
+        
+        # Download base model
+        print("📥 Downloading base model...")
+        snapshot_download(
+            "black-forest-labs/FLUX.1-dev",
+            local_dir=str(model_dir),
+            ignore_patterns=["*.pt", "*.bin"],  # using safetensors
+            token=hf_token
+        )
+        
+        # Load and validate model
+        DiffusionPipeline.from_pretrained(str(model_dir), torch_dtype=torch.bfloat16)
+        print("✅ Base model loaded successfully")
+        
+        # Load training images from HF dataset
+        print(f"📥 Loading images from dataset: {dataset_id}")
+        img_path = load_images_from_hf_dataset(dataset_id, hf_token)
+        
+        # Set up training configuration
+        config = APITrainConfig(
+            instance_name=instance_name or "subject",
+            class_name=class_name or "person", 
+            max_train_steps=max_train_steps
+        )
+        
+        # Set up hugging face accelerate library for fast training
+        write_basic_config(mixed_precision="bf16")
+        
+        # Define the training prompt
+        instance_phrase = f"{config.instance_name} the {config.class_name}"
+        prompt = f"{config.prefix} {instance_phrase} {config.postfix}".strip()
+        
+        print(f"🎯 Training prompt: {prompt}")
+        print(f"🚀 Starting training for {max_train_steps} steps...")
+        
+        # Execute training subprocess
+        def _exec_subprocess(cmd: list[str]):
+            """Executes subprocess and prints log to terminal while subprocess is running."""
+            process = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+            )
+            with process.stdout as pipe:
+                for line in iter(pipe.readline, b""):
+                    line_str = line.decode()
+                    print(f"{line_str}", end="")
+
+            if exitcode := process.wait() != 0:
+                raise subprocess.CalledProcessError(exitcode, "\n".join(cmd))
+        
+        # Run training
+        _exec_subprocess([
+            "accelerate",
+            "launch",
+            "examples/dreambooth/train_dreambooth_lora_flux.py",
+            "--mixed_precision=bf16",
+            f"--pretrained_model_name_or_path={model_dir}",
+            f"--instance_data_dir={img_path}",
+            f"--output_dir={output_dir}",
+            f"--instance_prompt={prompt}",
+            f"--resolution={config.resolution}",
+            f"--train_batch_size={config.train_batch_size}",
+            f"--gradient_accumulation_steps={config.gradient_accumulation_steps}",
+            f"--learning_rate={config.learning_rate}",
+            f"--lr_scheduler={config.lr_scheduler}",
+            f"--lr_warmup_steps={config.lr_warmup_steps}",
+            f"--max_train_steps={config.max_train_steps}",
+            f"--checkpointing_steps={config.checkpointing_steps}",
+            f"--seed={config.seed}",
+        ])
+        
+        print("✅ Training completed!")
+        
+        # Upload trained LoRA to HuggingFace repository
+        
+        print(f"📤 Uploading LoRA to repository: {output_repo}")
+
+        # Create repository if it doesn't exist
+        create_repo(
+            repo_id=output_repo,
+            repo_type="model",
+            token=hf_token
+        )
+
+        # print contents of output_dir
+        print(f"Contents of {output_dir}:")
+        for file in output_dir.iterdir():
+            print(file)
+
+        upload_folder(
+            folder_path=str(output_dir),
+            repo_id=output_repo,
+            repo_type="model",
+            token=hf_token,
+            commit_message=f"Add LoRA trained on {dataset_id}",
+        )
+        
+        print(f"🎉 Successfully uploaded LoRA to {output_repo}")
+        
+        return {
+            "status": "success",
+            "message": f"LoRA training completed and uploaded to {output_repo}",
+            "dataset_used": dataset_id,
+            "training_steps": max_train_steps,
+            "training_prompt": prompt
+        }
+
+
+# ## API Endpoints with Job ID System
+
+@app.function(
+    image=image,
+    keep_warm=1,  # Keep one container warm for faster response
+)
+@modal.fastapi_endpoint(method="POST")
+def api_start_training(item: dict):
+    """
+    Start LoRA training and return a job ID.
+    
+    Expected JSON payload:
+    {
+        "dataset_id": "username/dataset-name",
+        "hf_token": "hf_...",
+        "output_repo": "username/output-repo",
+        "instance_name": "optional_subject_name",
+        "class_name": "optional_class_name", 
+        "max_train_steps": 500
+    }
+    """
+    try:
+        # Extract required parameters
+        dataset_id = item["dataset_id"]
+        hf_token = item["hf_token"] 
+        output_repo = item["output_repo"]
+        
+        # Extract optional parameters
+        instance_name = item.get("instance_name")
+        class_name = item.get("class_name")
+        max_train_steps = item.get("max_train_steps", 500)
+        
+        # Start training (non-blocking)
+        call_handle = train_lora_stateless.spawn(
+            dataset_id=dataset_id,
+            hf_token=hf_token,
+            output_repo=output_repo,
+            instance_name=instance_name,
+            class_name=class_name,
+            max_train_steps=max_train_steps
+        )
+        
+        job_id = call_handle.object_id
+        
+        return {
+            "status": "started",
+            "job_id": job_id,
+            "message": "Training job started successfully",
+            "dataset_id": dataset_id,
+            "output_repo": output_repo,
+            "max_train_steps": max_train_steps
+        }
+        
+    except KeyError as e:
+        return {
+            "status": "error",
+            "message": f"Missing required parameter: {e}"
+        }
+    except Exception as e:
+        return {
+            "status": "error", 
+            "message": f"Failed to start training: {str(e)}"
+        }
+
+
+@app.function(
+    image=image,
+    keep_warm=1,
+)
+@modal.fastapi_endpoint(method="GET")
+def api_job_status(job_id: str):
+    """
+    Check the status of a training job.
+    Pass job_id as a query parameter: /job_status?job_id=xyz
+    """
+    try:
+        from modal.functions import FunctionCall
+        
+        # Get the function call handle
+        call_handle = FunctionCall.from_id(job_id)
+        
+        if call_handle is None:
+            return {
+                "status": "error",
+                "message": "Job not found"
+            }
+        
+        # Check if the job is finished
+        try:
+            result = call_handle.get(timeout=0)  # Non-blocking check
+            return {
+                "status": "completed",
+                "result": result
+            }
+        except TimeoutError:
+            return {
+                "status": "running",
+                "message": "Job is still running"
+            }
+        except Exception as e:
+            return {
+                "status": "failed",
+                "message": f"Job failed: {str(e)}"
+            }
+            
+    except Exception as e:
+        return {
+            "status": "error",
+            "message": f"Error checking job status: {str(e)}"
+        }
+
+
+
 # ## Low-Rank Adapation (LoRA) fine-tuning for a text-to-image model
 
 # The base model we start from is trained to do a sort of "reverse [ekphrasis](https://en.wikipedia.org/wiki/Ekphrasis)":

requirements.txt

@@ -0,0 +1,2 @@
+gradio>=4.0.0
+requests>=2.25.0