Upload folder using huggingface_hub

803f9b1 about 3 years ago

2.28 kB

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")

	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")