Spaces:
Runtime error
Runtime error
File size: 3,916 Bytes
c3a809c f673b98 3517ef3 40b3e9c 3517ef3 f673b98 d41934b f673b98 bc2361e f673b98 3517ef3 f673b98 bc2361e f673b98 bc2361e f673b98 3517ef3 f673b98 3517ef3 f673b98 3517ef3 f673b98 3517ef3 f673b98 3517ef3 f673b98 3517ef3 f673b98 3517ef3 f673b98 |
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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
import os
import io
import json
import base64
import time
import torch
from PIL import Image
from typing import Optional
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware
from safetensors.torch import save_file
from src.pipeline import FluxPipeline
from src.transformer_flux import FluxTransformer2DModel
from src.lora_helper import set_single_lora, set_multi_lora, unset_lora
# Define paths
base_path = "black-forest-labs/FLUX.1-dev"
lora_base_path = "./models"
# Initialize the model
print("Loading model...")
pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16)
transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.bfloat16)
pipe.transformer = transformer
pipe.to("cuda")
print("Model loaded successfully!")
# Function to clear cache
def clear_cache(transformer):
for name, attn_processor in transformer.attn_processors.items():
attn_processor.bank_kv.clear()
# Create FastAPI app
app = FastAPI(title="Ghibli Image Generator API",
description="Convert images to Ghibli Studio style using EasyControl")
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/")
def myfunc():
return {"status":"running"}
# Health check endpoint
@app.get("/health")
async def health_check():
return {"status": "healthy", "model": "loaded"}
# Main image conversion endpoint
@app.post("/generate-ghibli")
async def generate_ghibli(
file: UploadFile = File(...),
prompt: str = Form("Ghibli Studio style, Charming hand-drawn anime-style illustration"),
height: int = Form(768),
width: int = Form(768),
seed: int = Form(42)
):
try:
# Validate input image
if not file.content_type.startswith("image/"):
raise HTTPException(status_code=400, detail="File must be an image")
# Read and validate image
image_data = await file.read()
try:
spatial_img = Image.open(io.BytesIO(image_data))
except Exception as e:
raise HTTPException(status_code=400, detail=f"Invalid image: {str(e)}")
# Validate dimensions
if height < 256 or height > 1024 or width < 256 or width > 1024:
raise HTTPException(status_code=400, detail="Dimensions must be between 256 and 1024")
# Configure LoRA
lora_path = os.path.join(lora_base_path, "Ghibli.safetensors")
set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
# Generate image
with torch.cuda.amp.autocast():
output = pipe(
prompt,
height=height,
width=width,
guidance_scale=3.5,
num_inference_steps=25,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(seed),
subject_images=[],
spatial_images=[spatial_img],
cond_size=512,
).images[0]
# Clear cache
clear_cache(pipe.transformer)
# Convert output to bytes
img_byte_arr = io.BytesIO()
output.save(img_byte_arr, format='PNG')
img_byte_arr.seek(0)
# Return the image directly
return Response(
content=img_byte_arr.getvalue(),
media_type="image/png"
)
except HTTPException as e:
raise e
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
# Run the API with uvicorn
if __name__ == "__main__":
import uvicorn
uvicorn.run("app:app", host="0.0.0.0", port=7860) |