OpenGVLab/InternVL

1# --------------------------------------------------------2# InternVL3# Copyright (c) 2024 OpenGVLab4# Licensed under The MIT License [see LICENSE for details]5# --------------------------------------------------------6 7from io import BytesIO8 9import torch10from diffusers import StableDiffusion3Pipeline11from fastapi import FastAPI12from fastapi.responses import Response13from pydantic import BaseModel14 15# Initialize pipeline16pipe = StableDiffusion3Pipeline.from_pretrained('stabilityai/stable-diffusion-3-medium-diffusers',17                                                torch_dtype=torch.float16)18pipe = pipe.to('cuda')19 20# Create a FastAPI application21app = FastAPI()22 23 24# Define the input data model25class CaptionRequest(BaseModel):26    caption: str27 28 29# Defining API endpoints30@app.post('/generate_image/')31async def generate_image(request: CaptionRequest):32    caption = request.caption33    negative_prompt = 'blurry, low resolution, artifacts, unnatural, poorly drawn, bad anatomy, out of focus'34    image = pipe(35        caption,36        negative_prompt=negative_prompt,37        num_inference_steps=20,38        guidance_scale=7.039    ).images[0]40 41    # Converts an image to a byte stream42    img_byte_arr = BytesIO()43    image.save(img_byte_arr, format='PNG')44    img_byte_arr = img_byte_arr.getvalue()45 46    return Response(content=img_byte_arr, media_type='image/png')47 48 49# Run the Uvicorn server50if __name__ == '__main__':51    import argparse52 53    import uvicorn54    parser = argparse.ArgumentParser()55    parser.add_argument('--port', default=11005, type=int)56    args = parser.parse_args()57 58    uvicorn.run(app, host='0.0.0.0', port=args.port)59