File size: 1,652 Bytes
3f1b7f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
# --------------------------------------------------------
# InternVL
# Copyright (c) 2024 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
from io import BytesIO
import torch
from diffusers import StableDiffusion3Pipeline
from fastapi import FastAPI
from fastapi.responses import Response
from pydantic import BaseModel
# Initialize pipeline
pipe = StableDiffusion3Pipeline.from_pretrained('stabilityai/stable-diffusion-3-medium-diffusers',
torch_dtype=torch.float16)
pipe = pipe.to('cuda')
# Create a FastAPI application
app = FastAPI()
# Define the input data model
class CaptionRequest(BaseModel):
caption: str
# Defining API endpoints
@app.post('/generate_image/')
async def generate_image(request: CaptionRequest):
caption = request.caption
negative_prompt = 'blurry, low resolution, artifacts, unnatural, poorly drawn, bad anatomy, out of focus'
image = pipe(
caption,
negative_prompt=negative_prompt,
num_inference_steps=20,
guidance_scale=7.0
).images[0]
# Converts an image to a byte stream
img_byte_arr = BytesIO()
image.save(img_byte_arr, format='PNG')
img_byte_arr = img_byte_arr.getvalue()
return Response(content=img_byte_arr, media_type='image/png')
# Run the Uvicorn server
if __name__ == '__main__':
import argparse
import uvicorn
parser = argparse.ArgumentParser()
parser.add_argument('--port', default=11005, type=int)
args = parser.parse_args()
uvicorn.run(app, host='0.0.0.0', port=args.port)
|