intruct.py

from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import StreamingResponse
import io
import math
from PIL import Image, ImageOps, ImageDraw
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline, StableDiffusionInpaintPipeline
from fastapi import FastAPI, Response
from fastapi.responses import FileResponse
import torch
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
from huggingface_hub import hf_hub_download, login
from safetensors.torch import load_file
from io import BytesIO
import os
import base64
from typing import List

# Initialize FastAPI app
app = FastAPI()

# Load the pre-trained InstructPix2Pix model for editing
model_id = "timbrooks/instruct-pix2pix"
pipe_edit = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    model_id, torch_dtype=torch.float16, safety_checker=None
).to("cuda")

# Load the pre-trained Inpainting model
inpaint_model_id = "stabilityai/stable-diffusion-2-inpainting"
pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
    inpaint_model_id, torch_dtype=torch.float16, safety_checker=None
).to("cuda")

# Default configuration values
DEFAULT_STEPS = 50
DEFAULT_TEXT_CFG = 7.5
DEFAULT_IMAGE_CFG = 1.5
DEFAULT_SEED = 1371

HF_TOKEN = os.getenv("HF_TOKEN")

def load_model():
    try:
        # Login to Hugging Face if token is provided
        if HF_TOKEN:
            login(token=HF_TOKEN)
            
        base = "stabilityai/stable-diffusion-xl-base-1.0"
        repo = "ByteDance/SDXL-Lightning"
        ckpt = "sdxl_lightning_4step_unet.safetensors"

        # Load model with explicit error handling
        unet = UNet2DConditionModel.from_config(
            base, 
            subfolder="unet"
        ).to("cuda", torch.float16)
        
        unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
        pipe = StableDiffusionXLPipeline.from_pretrained(
            base, 
            unet=unet, 
            torch_dtype=torch.float16, 
            variant="fp16"
        ).to("cuda")
        
        # Configure scheduler
        pipe.scheduler = EulerDiscreteScheduler.from_config(
            pipe.scheduler.config, 
            timestep_spacing="trailing"
        )
        
        return pipe
    
    except Exception as e:
        raise Exception(f"Failed to load model: {str(e)}")

# Load model at startup with error handling
try:
    pipe_generate = load_model()
except Exception as e:
    print(f"Model initialization failed: {str(e)}")
    raise

@app.get("/generate")
async def generate_image(prompt: str):
    try:
        # Generate image
        image = pipe_generate(
            prompt,
            num_inference_steps=4,
            guidance_scale=0
        ).images[0]
        
        # Save image to buffer
        buffer = BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)
        
        return Response(content=buffer.getvalue(), media_type="image/png")
    
    except Exception as e:
        return {"error": str(e)}

@app.get("/generate_multiple")
async def generate_multiple_images(prompts: List[str]):
    try:
        # List to store base64-encoded images
        generated_images = []
        
        # Generate an image for each prompt
        for prompt in prompts:
            image = pipe_generate(
                prompt,
                num_inference_steps=4,
                guidance_scale=0
            ).images[0]
            
            # Save image to buffer
            buffer = BytesIO()
            image.save(buffer, format="PNG")
            buffer.seek(0)
            
            # Encode the image as base64
            image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
            generated_images.append({
                "prompt": prompt,
                "image_base64": image_base64
            })
        
        return {"images": generated_images}
    
    except Exception as e:
        return {"error": str(e)}

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

def process_image(input_image: Image.Image, instruction: str, steps: int, text_cfg_scale: float, image_cfg_scale: float, seed: int):
    """
    Process the input image with the given instruction using InstructPix2Pix.
    """
    # Resize image to fit model requirements
    width, height = input_image.size
    factor = 512 / max(width, height)
    factor = math.ceil(min(width, height) * factor / 64) * 64 / min(width, height)
    width = int((width * factor) // 64) * 64
    height = int((height * factor) // 64) * 64
    input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)

    if not instruction:
        return input_image

    # Set the random seed for reproducibility
    generator = torch.manual_seed(seed)

    # Generate the edited image
    edited_image = pipe_edit(
        instruction,
        image=input_image,
        guidance_scale=text_cfg_scale,
        image_guidance_scale=image_cfg_scale,
        num_inference_steps=steps,
        generator=generator,
    ).images[0]

    return edited_image

@app.post("/edit-image/")
async def edit_image(
    file: UploadFile = File(...),
    instruction: str = Form(...),
    steps: int = Form(default=DEFAULT_STEPS),
    text_cfg_scale: float = Form(default=DEFAULT_TEXT_CFG),
    image_cfg_scale: float = Form(default=DEFAULT_IMAGE_CFG),
    seed: int = Form(default=DEFAULT_SEED)
):
    """
    Endpoint to edit an image based on a text instruction.
    """
    # Read and convert the uploaded image
    image_data = await file.read()
    input_image = Image.open(io.BytesIO(image_data)).convert("RGB")

    # Process the image
    edited_image = process_image(input_image, instruction, steps, text_cfg_scale, image_cfg_scale, seed)

    # Convert the edited image to bytes
    img_byte_arr = io.BytesIO()
    edited_image.save(img_byte_arr, format="PNG")
    img_byte_arr.seek(0)

    # Return the image as a streaming response
    return StreamingResponse(img_byte_arr, media_type="image/png")

# New endpoint for inpainting
@app.post("/inpaint/")
async def inpaint_image(
    file: UploadFile = File(...),
    prompt: str = Form(...),
    mask_coordinates: str = Form(...),  # Format: "x1,y1,x2,y2" (top-left and bottom-right of the rectangle to inpaint)
    steps: int = Form(default=DEFAULT_STEPS),
    guidance_scale: float = Form(default=7.5),
    seed: int = Form(default=DEFAULT_SEED)
):
    """
    Endpoint to perform inpainting on an image.
    - file: The input image to inpaint.
    - prompt: The text prompt describing what to generate in the inpainted area.
    - mask_coordinates: Coordinates of the rectangular area to inpaint (format: "x1,y1,x2,y2").
    - steps: Number of inference steps.
    - guidance_scale: Guidance scale for the inpainting process.
    - seed: Random seed for reproducibility.
    """
    try:
        # Read and convert the uploaded image
        image_data = await file.read()
        input_image = Image.open(io.BytesIO(image_data)).convert("RGB")

        # Resize image to fit model requirements (must be divisible by 8 for inpainting)
        width, height = input_image.size
        factor = 512 / max(width, height)
        factor = math.ceil(min(width, height) * factor / 8) * 8 / min(width, height)
        width = int((width * factor) // 8) * 8
        height = int((height * factor) // 8) * 8
        input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)

        # Create a mask for inpainting
        mask = Image.new("L", (width, height), 0)  # Black image (0 = no inpainting)
        draw = ImageDraw.Draw(mask)

        # Parse the mask coordinates
        try:
            x1, y1, x2, y2 = map(int, mask_coordinates.split(","))
            # Adjust coordinates based on resized image
            x1 = int(x1 * factor)
            y1 = int(y1 * factor)
            x2 = int(x2 * factor)
            y2 = int(y2 * factor)
        except ValueError:
            return {"error": "Invalid mask coordinates format. Use 'x1,y1,x2,y2'."}

        # Draw a white rectangle on the mask (255 = area to inpaint)
        draw.rectangle([x1, y1, x2, y2], fill=255)

        # Set the random seed for reproducibility
        generator = torch.manual_seed(seed)

        # Perform inpainting
        inpainted_image = pipe_inpaint(
            prompt=prompt,
            image=input_image,
            mask_image=mask,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            generator=generator,
        ).images[0]

        # Convert the inpainted image to bytes
        img_byte_arr = io.BytesIO()
        inpainted_image.save(img_byte_arr, format="PNG")
        img_byte_arr.seek(0)

        # Return the image as a streaming response
        return StreamingResponse(img_byte_arr, media_type="image/png")

    except Exception as e:
        return {"error": str(e)}

@app.get("/")
async def root():
    """
    Root endpoint for basic health check.
    """
    return {"message": "InstructPix2Pix API is running. Use POST /edit-image/ or /inpaint/ to edit images."}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)