intruct.py

from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import StreamingResponse
import io
import math
from PIL import Image, ImageOps
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline
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  # Added for encoding images as base64
from typing import List  # Added for type hinting the list of prompts



# Initialize FastAPI app
app = FastAPI()

# Load the pre-trained model once at startup
model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    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 = 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(
            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)}

# New endpoint to handle a list of prompts
@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(
                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(
        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.
    - file: The input image to edit.
    - instruction: The text instruction for editing the image.
    - steps: Number of inference steps.
    - text_cfg_scale: Text CFG weight.
    - image_cfg_scale: Image CFG weight.
    - seed: Random seed for reproducibility.
    """
    # 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")

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

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