predict.py

from cog import BasePredictor, Input, Path
import os
import time
import subprocess
import torch
from PIL import Image
from diffusers import (
    FluxControlNetPipeline,
    FluxControlNetModel,
    FluxMultiControlNetModel,
    FluxControlNetInpaintPipeline
)

# from huggingface_hub import hf_hub_download
# from torchao.quantization import quantize_, int8_weight_only
from sd_embed.embedding_funcs import get_weighted_text_embeddings_flux1
import numpy as np

MODEL_CACHE = "FLUX.1-dev"
CONTROLNET_CACHE = "FLUX.1-dev-ControlNet-Union-Pro"

class Predictor(BasePredictor):
    def setup(self) -> None:
        """Initialize all models during setup"""
        controlnet_union = FluxControlNetModel.from_pretrained(
            CONTROLNET_CACHE,
            torch_dtype=torch.bfloat16
        )

        controlnet = FluxMultiControlNetModel([controlnet_union])

        print("Loading main pipeline...")
        self.pipe = FluxControlNetPipeline.from_pretrained(
            MODEL_CACHE,
            controlnet=controlnet,
            torch_dtype=torch.bfloat16
        ).to("cuda")

        # Quantize transformer
        # quantize_(self.pipe.transformer, int8_weight_only())

        # Create inpaint pipe and copy components
        print("Creating inpaint pipeline...")
        try:
            self.inpaint_pipe = FluxControlNetInpaintPipeline(
                vae=self.pipe.vae,
                text_encoder=self.pipe.text_encoder,
                tokenizer=self.pipe.tokenizer, 
                text_encoder_2=self.pipe.text_encoder_2,
                tokenizer_2=self.pipe.tokenizer_2,
                transformer=self.pipe.transformer,
                controlnet=controlnet,
                scheduler=self.pipe.scheduler,
            ).to("cuda")
            print("inpaint pipeline loaded")
        except Exception as e:
            print('error:', e)
            for component in [
                self.inpaint_pipe.vae,
                self.inpaint_pipe.text_encoder,
                self.inpaint_pipe.text_encoder_2,
                self.inpaint_pipe.transformer
            ]:
                if hasattr(component, "to"):
                    component.to(dtype=torch.bfloat16)

    def predict(
        self,
        prompt: str = Input(
            description="The text prompt that guides image generation",
            default=""
        ),
        canny_image: Path = Input(
            description="Preprocessed canny image. Required.",
            default=None
        ),
        depth_image: Path = Input(
            description="Preprocessed depth image. Required.",
            default=None
        ),
        canny_strength: float = Input(
            description="Strength for canny conditionning.",
            default=0.6,
            ge=0,
            le=2
        ),
        depth_strength: float = Input(
            description="Strength for depth conditionning.",
            default=0.6,
            ge=0,
            le=2
        ),
        guidance_scale: float = Input(
            description="Controls how closely the image follows the prompt.",
            default=3.5,
            ge=0,
            le=50
        ),
        steps: int = Input(
            description="Number of denoising steps.",
            default=8,
            ge=1,
            le=50
        ),
        seed: int = Input(
            description="Random seed for reproducibility.",
            default=None
        ),
        masked_image: Path = Input(
            description="Input image for inpainting. Optionnal.",
            default=None
        ),
        mask_image: Path = Input(
            description="Mask image for inpainting (white pixels will be repainted). Optionnal.",
            default=None
        ),
        fill_strength: float = Input(
            description="Strength of inpainting (how much to repaint masked area).",
            default=0.6,
            ge=0,
            le=1
        ),
        control_guidance_start: float = Input(
            description="Percentage of steps at which conditionning starts applying.",
            default=0.0,
            ge=0.0,
            le=1.0
        ),
        control_guidance_end: float = Input(
            description="Percentage of steps at which conditionning stops applying.",
            default=1.0,
            ge=0.0,
            le=1.0
        ),
    ) -> Path:
        if seed is None:
            seed = int.from_bytes(os.urandom(2), "big")
        print(f"Using seed: {seed}")
        generator = torch.Generator("cuda").manual_seed(seed)

        # Process control images
        control_images = []
        control_strengths = []
        reference_size = None

        image_configs = [
            (canny_image, canny_strength),
            (depth_image, depth_strength),
        ]

        # First, determine reference size from the first available image
        for img_path, _ in image_configs:
            if img_path:
                img = Image.open(img_path)
                width = (img.width // 16) * 16
                height = (img.height // 16) * 16
                reference_size = (width, height)
                print(f"Reference size set to: {reference_size}")
                break

        if reference_size is None:
            raise ValueError("At least one image must be provided")

        # Process images with chosen detector types
        for img_path, strength in image_configs:
            if img_path:
                img = Image.open(img_path)
                img = img.resize(reference_size, Image.LANCZOS)
                
                control_images.append(img)
                control_strengths.append(strength)

        if not canny_image:
            raise ValueError("Preprocessed Canny image is required.")

        if not depth_image:
            raise ValueError("Preprocessed Depth image is required.")

        # Handle custom dimensions
        final_width = reference_size[0]
        final_height = reference_size[1]

        # Prepare common generation parameters
        generation_params = {
            'control_image': control_images,
            'controlnet_conditioning_scale': control_strengths,
            'control_mode':[0, 2],
            'width': final_width,
            'height': final_height,
            'control_guidance_start': control_guidance_start,
            'control_guidance_end': control_guidance_end,
            'num_inference_steps': steps,
            'guidance_scale': guidance_scale,
            'generator': generator,
            'prompt': prompt
        }

        # Add inpainting-specific parameters if needed
        if mask_image is not None:
            if masked_image is None:
                raise ValueError("Inpainting requires both masked_image and mask_image")
            init_image = Image.open(masked_image)
            mask = Image.open(mask_image)
            
            generation_params.update({
                'image': init_image,
                'mask_image': mask,
                'strength': fill_strength
            })
            
            selected_pipe = self.inpaint_pipe
        else:
            selected_pipe = self.pipe

        generated_image = selected_pipe(**generation_params).images[0]

        output_path = f"/tmp/output.png"
        generated_image.save(output_path)
        print(f"Generation complete. Output saved to: {output_path}")
        return Path(output_path)