Deploy Multi-GPU Training

class FluxPreprocessorApp(fal.App):
    """
    Standalone Flux preprocessing app.
    
    Uses all GPUs in parallel to preprocess images:
    - Each GPU processes a subset of images
    - Results are gathered and saved
    - Returns URL to preprocessed data
    """
    
    machine_type = "GPU-H100"
    num_gpus = 2
    keep_alive = 300
    min_concurrency = 0
    max_concurrency = 2
    
    requirements = [
        "torch==2.4.0",
        "torchvision",  # Required by moondream
        "diffusers==0.30.3",
        "transformers==4.46.0",
        "tokenizers==0.20.1",
        "sentencepiece",
        "accelerate==1.4.0",
        "pyzmq==26.0.0",
        "huggingface_hub==0.26.5",
        "moondream==0.0.5",
        "einops",  # Required by moondream
        "pillow>=10.0.0",
        "timm",  # Required by moondream vision encoder
    ]

    async def setup(self) -> None:
        """
        Initialize the preprocessing runner.
        
        Downloads Flux weights and starts GPU workers for preprocessing.
        """
        import os
        from huggingface_hub import snapshot_download
        
        os.environ["HF_HOME"] = "/data/models"
        
        # Download Flux weights
        print("Downloading Flux model weights...")
        model_path = snapshot_download(
            repo_id="black-forest-labs/FLUX.1-dev",
            local_dir="/data/flux_weights",
        )
        print(f"Model downloaded to {model_path}")
        
        # Create preprocessing runner
        self.runner = DistributedRunner(
            worker_cls=FluxPreprocessorWorker,
            world_size=self.num_gpus,
        )
        
        # Start workers
        print(f"Starting {self.num_gpus} preprocessing workers...")
        await self.runner.start(model_path=model_path)
        
        print("Preprocessing workers ready!")

    @fal.endpoint("/")
    async def preprocess(
        self,
        request: PreprocessRequest,
        response: Response,
    ) -> PreprocessResponse:
        """
        Preprocess images for Flux LoRA training.
        
        This endpoint:
        1. Downloads ZIP of images
        2. Generates/loads captions
        3. Injects trigger word
        4. Encodes images with VAE (parallel across GPUs)
        5. Encodes captions with T5/CLIP (parallel across GPUs)
        6. Returns preprocessed data file
        """
        import tempfile
        
        with tempfile.TemporaryDirectory() as temp_dir:
            # Download images ZIP
            print(f"Downloading images from {request.images_data_url}")
            images_zip_path = str(download_file(
                request.images_data_url,
                target_dir=temp_dir
            ))
            
            # Generate unique request ID for this preprocessing job
            import time
            request_id = str(int(time.time() * 1000000))  # Microsecond timestamp
            
            # Preprocess in parallel across all GPUs
            print(f"Preprocessing with {self.num_gpus} GPUs (request_id: {request_id})...")
            result = await self.runner.invoke({
                "images_zip_url": images_zip_path,
                "request_id": request_id,
                "trigger_word": request.trigger_word,
                "resolution": request.resolution,
            })
            # ↑ runner.invoke() sends this dict to all workers' __call__() method
            # Workers process data in parallel, then rank 0 returns the result
            # See: /serverless/distributed/api-reference#invoke
            
            # Upload preprocessed data
            preprocessed_file = File.from_path(result["preprocessed_data_path"])
            num_images = result["num_images"]
            
            print(f"Preprocessing complete! Processed {num_images} images")
            
            return PreprocessResponse(
                preprocessed_data=preprocessed_file,
                num_images=num_images,
                message=f"Preprocessed {num_images} images with trigger word '{request.trigger_word}'"
            )

import fal_client

# Test preprocessing with your images
result = fal_client.subscribe(
    "username/your-preprocessor-app-id",
    arguments={
        "images_data_url": "https://example.com/training-images.zip",
        "trigger_word": "ohwx",
        "resolution": 512
    },
    with_logs=True,
)

print(f"Preprocessed {result['num_images']} images")
print(f"Data file: {result['preprocessed_data']['url']}")

class FluxPreprocessorWorker(DistributedWorker):
    """Worker for preprocessing images in parallel across GPUs"""
    
    def setup(self, model_path: str, **kwargs):
        """
        Load VAE, text encoders, and captioning model on each GPU.
        Called once per worker during runner.start().
        """
        from diffusers import AutoencoderKL
        from transformers import CLIPTextModel, T5EncoderModel
        
        self.rank_print(f"Loading preprocessing models on {self.device}")
        
        # Load VAE for image encoding
        self.vae = AutoencoderKL.from_pretrained(
            model_path,
            subfolder="vae",
            torch_dtype=torch.float16,
        ).to(self.device)
        
        # Load text encoders (T5 and CLIP)
        self.text_encoder = CLIPTextModel.from_pretrained(
            model_path, subfolder="text_encoder"
        ).to(self.device)
        
        self.text_encoder_2 = T5EncoderModel.from_pretrained(
            model_path, subfolder="text_encoder_2"
        ).to(self.device)
        
        # Load captioning model (Moondream)
        from moondream import Moondream
        self.caption_model = Moondream(device=self.device)
        
        self.rank_print("Models loaded successfully")
    
    def __call__(
        self,
        images_zip_url: str,
        request_id: str,
        trigger_word: str,
        resolution: int = 512,
        **kwargs
    ):
        """
        Process a subset of images on this GPU.
        Each worker processes different images in parallel.
        """
        import torch.distributed as dist
        
        # Step 1: Load and distribute images
        if self.rank == 0:
            # Only rank 0 loads the full dataset
            images = load_images_from_zip(images_zip_url)
            num_images = len(images)
        else:
            images = []
            num_images = 0
        
        # Broadcast number of images to all workers
        num_images = dist.broadcast_object_list([num_images], src=0)[0]
        
        # Step 2: Each worker processes a subset
        # Worker 0 processes images [0, 4, 8, ...]
        # Worker 1 processes images [1, 5, 9, ...]
        my_images = images[self.rank::self.world_size]
        
        self.rank_print(f"Processing {len(my_images)} images")
        
        # Step 3: Generate captions for my subset
        captions = []
        for img in my_images:
            caption = self.caption_model.caption(img)
            # Inject trigger word
            caption = f"{trigger_word} {caption}"
            captions.append(caption)
        
        # Step 4: Encode images with VAE
        latents_list = []
        for img in my_images:
            img_tensor = preprocess_image(img, resolution)
            with torch.no_grad():
                latent = self.vae.encode(img_tensor.to(self.device)).latent_dist.sample()
            latents_list.append(latent)
        
        # Step 5: Encode captions with text encoders
        text_embeddings_list = []
        for caption in captions:
            with torch.no_grad():
                text_emb = encode_text(caption, self.text_encoder, self.text_encoder_2)
            text_embeddings_list.append(text_emb)
        
        # Step 6: Gather results from all workers to rank 0
        if self.rank == 0:
            # Rank 0 collects from all workers
            all_latents = gather_from_all_workers(latents_list)
            all_text_embeddings = gather_from_all_workers(text_embeddings_list)
            
            # Save preprocessed data
            output_path = f"/tmp/preprocessed_{request_id}.pt"
            torch.save({
                "latents": torch.cat(all_latents),
                "text_embeddings": torch.cat(all_text_embeddings),
                "num_images": num_images,
            }, output_path)
            
            return {
                "preprocessed_data_path": output_path,
                "num_images": num_images,
            }
        
        return {}  # Other ranks return empty

class FluxLoRATrainingApp(fal.App):
    """
    Flux LoRA Training App using DistributedRunner.
    
    This app focuses on TRAINING only and uses GPUs with DDP.
    For preprocessing, it calls the flux-preprocessor app.
    
    Benefits:
    - Training runner stays warm (no reload)
    - Preprocessor runs on separate GPUs (no conflict)
    - Clean separation of concerns
    - Can scale independently
    """
    
    machine_type = "GPU-H100"
    num_gpus = 2
    keep_alive = 3000
    min_concurrency = 1
    max_concurrency = 1
    
    requirements = [
        "torch==2.4.0",
        "diffusers==0.30.3",
        "transformers==4.46.0",
        "tokenizers==0.20.1",
        "sentencepiece",
        "peft==0.12.0",
        "safetensors==0.4.4",
        "accelerate==1.4.0",
        "pyzmq==26.0.0",
        "huggingface_hub==0.26.5",
        "fal-client",  # For calling preprocessor app
    ]
    
    async def setup(self) -> None:
        """
        Initialize the training runner.
        
        Downloads Flux weights and starts GPU workers for training.
        Preprocessing is handled by calling a separate app.
        """
        import os
        from huggingface_hub import snapshot_download
        
        os.environ["HF_HOME"] = "/data/models"
        
        # Download Flux weights
        print("Downloading Flux model weights...")
        model_path = snapshot_download(
            repo_id="black-forest-labs/FLUX.1-dev",
            local_dir="/data/flux_weights",
        )
        print(f"Model downloaded to {model_path}")
        
        # Create training runner (uses all GPUs)
        self.runner = DistributedRunner(
            worker_cls=FluxLoRATrainingWorker,
            world_size=self.num_gpus,
        )
        
        # Start workers
        print(f"Starting {self.num_gpus} training workers...")
        await self.runner.start(model_path=model_path)
        
        print("Training workers ready!")

    @fal.endpoint("/train")
    async def train(
        self,
        request: CompleteTrainingRequest,
    ) -> TrainingResponse:
        """
        Complete training pipeline: preprocess + train.
        
        This demonstrates microservices architecture where the training app
        calls a separate preprocessor app via API.
        """
        import fal_client
        import tempfile
        
        # Step 1: Call preprocessor app (runs on separate instance with GPUs)
        print(f"Calling preprocessor app: {request.preprocessor_app}")
        print(f"Preprocessing {request.images_data_url}...")
        
        try:
            preprocess_result = fal_client.subscribe(
                request.preprocessor_app,
                arguments={
                    "images_data_url": request.images_data_url,
                    "trigger_word": request.trigger_word,
                    "resolution": request.resolution,
                },
                with_logs=True,
            )
        except Exception as e:
            print(f"Preprocessing failed: {e}")
            print(f"Make sure {request.preprocessor_app} is deployed!")
            raise
        
        # Get preprocessed data URL
        preprocessed_data_url = preprocess_result["preprocessed_data"]["url"]
        num_images = preprocess_result["num_images"]
        
        print(f"✓ Preprocessed {num_images} images")
        
        # Step 2: Download preprocessed data and train
        with tempfile.TemporaryDirectory() as temp_dir:
            print("Downloading preprocessed data...")
            preprocessed_path = str(download_file(
                preprocessed_data_url,
                target_dir=temp_dir
            ))
            
            # Step 3: Train LoRA (on our GPUs)
            print(f"Training with {self.num_gpus} GPUs...")
            result = await self.runner.invoke({
                "training_data_path": preprocessed_path,
                "learning_rate": request.learning_rate,
                "b_up_factor": request.b_up_factor,
                "max_train_steps": request.steps,
                "batch_size": request.batch_size,
                "gradient_accumulation_steps": request.gradient_accumulation_steps,
                "guidance_scale": 1.0,
                "use_masks": False,
                "lr_scheduler": "constant",
                "lr_warmup_steps": 0,
                "seed": request.seed,
                "streaming": False,
            })
            # ↑ This payload dict is passed to each worker's __call__() method
            # All workers train in parallel with DDP, then rank 0 returns the checkpoint
            # See: /serverless/distributed/api-reference#invoke
            
            # Check for errors
            if "error" in result:
                raise RuntimeError(f"Training failed: {result['error']}")
            
            # Upload checkpoint
            checkpoint_file = File.from_path(result["checkpoint_path"])
            
            return TrainingResponse(
                checkpoint=checkpoint_file,
                final_loss=result["final_loss"],
                num_steps=result["num_steps"],
                message=f"Training complete! {request.steps} steps, final loss: {result['final_loss']:.6f}"
            )

class FluxLoRATrainingWorker(DistributedWorker):
    """
    Production-ready distributed worker for Flux LoRA training.
    """

    def setup(self, model_path: str = "/data/flux_weights", **kwargs: Any) -> None:
        """
        Initialize the model on each GPU worker with proper LoRA configuration.
        """
        from diffusers import FluxTransformer2DModel
        from peft import LoraConfig

        self.rank_print(f"Loading Flux model on {self.device}")
        
        # Load the transformer model
        self.transformer = FluxTransformer2DModel.from_pretrained(
            model_path,
            subfolder="transformer",
            torch_dtype=torch.bfloat16,
        ).to(self.device)
        
        # Configure LoRA targeting all blocks (both double and single stream)
        target_modules = []
        
        # Double stream blocks (19 blocks) - these handle text-image interaction
        for block_num in range(19):
            target_modules.extend([
                f"transformer_blocks.{block_num}.attn.to_q",
                f"transformer_blocks.{block_num}.attn.to_k",
                f"transformer_blocks.{block_num}.attn.to_v",
                f"transformer_blocks.{block_num}.attn.to_out.0",
                # ... more attention and FF layers
            ])
        
        # Single stream blocks (38 blocks) - CRITICAL for image generation
        for block_num in range(38):
            target_modules.extend([
                f"single_transformer_blocks.{block_num}.attn.to_q",
                f"single_transformer_blocks.{block_num}.attn.to_k",
                f"single_transformer_blocks.{block_num}.attn.to_v",
                # ... more attention layers
            ])
        
        lora_config = LoraConfig(
            r=16,  # rank
            lora_alpha=16,
            target_modules=target_modules,
            lora_dropout=0.0,
            bias="none",
            init_lora_weights="gaussian",
        )
        
        # Add LoRA adapters
        self.transformer.add_adapter(lora_config)
        
        # Freeze base model, only train LoRA
        self.transformer.requires_grad_(False)
        for name, param in self.transformer.named_parameters():
            if "lora" in name:
                param.requires_grad = True
        
        # Wrap with DDP for synchronized training
        self.transformer = DDP(
            self.transformer,
            device_ids=[self.rank],
            output_device=self.rank,
            find_unused_parameters=False,
        )
        
        self.rank_print("Model loaded and wrapped with DDP")

    def __call__(
        self,
        streaming: bool = False,
        training_data_path: str = None,
        learning_rate: float = 4e-4,
        max_train_steps: int = 100,
        batch_size: int = 1,
        seed: int = 42,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """
        Production-ready training function with proper Flux training logic.
        """
        self.rank_print(f"Starting training: lr={learning_rate}, steps={max_train_steps}")
        
        # Step 1: Load training data (only rank 0)
    if self.rank == 0:
            self.rank_print(f"Loading training data from {training_data_path}")
            data = torch.load(training_data_path, map_location="cpu")
        latents = data["latents"]
        text_embeddings = data["text_embeddings"]
            pooled_embeddings = data["pooled_embeddings"]
            text_ids = data["text_ids"]
    else:
        latents = torch.empty(0)
        text_embeddings = torch.empty(0)
            pooled_embeddings = torch.empty(0)
            text_ids = torch.empty(0, dtype=torch.long)
    
        # Step 2: Broadcast data to all ranks
        torch.cuda.set_device(self.device)
        objects = [latents, text_embeddings, pooled_embeddings, text_ids]
    dist.broadcast_object_list(objects, src=0)
        latents, text_embeddings, pooled_embeddings, text_ids = objects
        
        # Move to GPU
        latents = latents.to(self.device, dtype=torch.bfloat16)
        text_embeddings = text_embeddings.to(self.device, dtype=torch.bfloat16)
        pooled_embeddings = pooled_embeddings.to(self.device, dtype=torch.bfloat16)
        text_ids = text_ids.to(self.device, dtype=torch.long)
        
        num_samples = latents.shape[0]
        self.rank_print(f"Loaded {num_samples} training samples")
        
        # Step 3: Set up optimizer with DUAL learning rates for lora_A and lora_B
        params_A = []
        params_B = []
        
        model = self.transformer.module
        
        for name, param in model.named_parameters():
            if param.requires_grad:
                if "lora_A" in name:
                    params_A.append(param)
                elif "lora_B" in name:
                    params_B.append(param)
        
        optimizer = torch.optim.AdamW([
            {"params": params_A, "lr": learning_rate, "weight_decay": 0.1},
            {"params": params_B, "lr": learning_rate * b_up_factor, "weight_decay": 0.1},
        ], betas=(0.9, 0.999), eps=1e-8)
        
        # Step 4: Training loop
        self.transformer.train()
        total_loss = 0.0
        
    for step in range(max_train_steps):
            # Each GPU gets different batch indices
            batch_indices = torch.randperm(num_samples, device="cpu")[:batch_size * self.world_size]
        local_indices = batch_indices[self.rank * batch_size:(self.rank + 1) * batch_size]
        
        # Get local batch
        batch_latents = latents[local_indices]
        batch_text_emb = text_embeddings[local_indices]
            batch_pooled_emb = pooled_embeddings[local_indices]
        
        # Compute loss
            loss = self.compute_flux_loss(
                latents=batch_latents,
                text_embeddings=batch_text_emb,
                pooled_embeddings=batch_pooled_emb,
                # ... other args
            )
        
        # Backward pass (DDP automatically syncs gradients)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
            total_loss += loss.item()
            
            # Stream progress (only rank 0)
            if streaming and self.rank == 0 and step % 10 == 0:
                avg_loss = total_loss / (step + 1)
                self.add_streaming_result({
                    "step": step,
                    "loss": loss.item(),
                    "avg_loss": avg_loss,
                }, as_text_event=True)
                # ↑ Sends intermediate results to the client during streaming
                # Only call from rank 0 to avoid duplicate messages
                # See: /serverless/distributed/api-reference#add_streaming_result
    
        # Step 5: Save checkpoint (only rank 0)
    if self.rank == 0:
            self.rank_print("Saving checkpoint...")
        checkpoint_path = self.save_lora_checkpoint()
            
            return {
                "checkpoint_path": checkpoint_path,
                "final_loss": total_loss / max_train_steps,
                "num_steps": max_train_steps,
            }
    
    return {}  # Other ranks return empty dict

# Wait for all GPUs to finish setup
dist.barrier()

if self.rank == 0:
    # Save checkpoint
    # Log metrics
    # Upload results