File size: 2,278 Bytes
803f9b1 |
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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from fastapi import FastAPI, Response
from fastapi.middleware.cors import CORSMiddleware
from diffusers import StableDiffusionPipeline
import torch
from torch import autocast
from io import BytesIO
import base64
import datetime # Import the datetime module
import uuid # Import the uuid module
import logging
import os
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_credentials=True,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"]
)
device = "cuda"
model_id = "rowenac/vintagechinesecomics"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to(device)
# Define a folder path for saving images
image_folder = "image_gallery"
# Create the directory if it doesn't exist
os.makedirs(image_folder, exist_ok=True)
@app.get("/")
def generate(prompt: str):
try:
with autocast(device):
image = pipe(prompt, guidance_scale=8.5).images[0]
# Generate a unique filename for each image using a timestamp and UUID
timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S%f") # Format the timestamp
image_uuid = uuid.uuid4()
image_filename = f"{image_folder}/image_{timestamp}_{image_uuid}.png"
# Save the image to the specified folder
image.save(image_filename)
logging.info(f"Image saved to {image_filename}")
buffer = BytesIO()
image.save(buffer, format="PNG")
imgstr = base64.b64encode(buffer.getvalue()).decode('utf-8')
return Response(content=imgstr, media_type="image/png")
except Exception as e:
logging.error(f"Error while generating image: {str(e)}")
return Response(content=str(e), status_code=500)
# New API endpoint to fetch image filenames
@app.get("/api/get_images")
def get_images():
# List all files in the image folder
images = os.listdir(image_folder)
return images
# New API endpoint to serve individual images
@app.get("/api/get_image/{image_name}")
def get_image(image_name: str):
# Define the path to the image
image_path = os.path.join(image_folder, image_name)
# Open and serve the image
with open(image_path, "rb") as image_file:
return Response(content=image_file.read(), media_type="image/png")
|