Spaces:
Sleeping
Sleeping
File size: 2,130 Bytes
4e216fb b50100a 4e216fb ad21ed7 4e216fb b50100a 4e216fb 8184fca 4e216fb | 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 | from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import FileResponse
from diffusers import AutoPipelineForImage2Image
from PIL import Image
from huggingface_hub import login
import torch
import os
os.environ["HF_HOME"] = "/app/huggingface_cache"
# Ensure the cache directory exists
os.makedirs("/app/huggingface_cache", exist_ok=True)
# Initialize FastAPI app
app = FastAPI()
hf_token = os.environ.get("HUGGINGFACE_TOKEN")
login(token=hf_token)
# Load the model
model_id = "kandinsky-community/kandinsky-2-2-decoder"
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = AutoPipelineForImage2Image.from_pretrained(
model_id,
torch_dtype=torch.float32, use_safetensors=True
).to(device)
# Create an output directory
os.makedirs("output_images", exist_ok=True)
@app.get("/")
async def root():
return {"message": "Welcome to the Image-to-Image API!"}
@app.post("/image-to-image/")
async def image_to_image(
prompt: str = Form(...),
stren: float = Form(...),
negative_prompt: str = Form(...),
image: UploadFile = File(...)
):
"""
Perform image-to-image transformation using a given prompt and input image.
Args:
- prompt (str): Text prompt describing the desired transformation.
- image (UploadFile): Input image file.
Returns:
- FileResponse: The transformed image file.
"""
try:
# Open and preprocess the input image
input_image = Image.open(image.file).convert("RGB")
original_size = input_image.size # Save the original size
# Generate the output image using the pipeline
generated_image = pipe(prompt=prompt,negative_prompt=negative_prompt ,image=input_image,strength=stren).images[0]
resized_image = generated_image.resize(original_size, Image.LANCZOS)
# Save the generated image
output_path = f"output_images/generated_{image.filename}"
resized_image.save(output_path)
# Return the generated image as a response
return FileResponse(output_path, media_type="image/png")
except Exception as e:
return {"error": str(e)}
|