Create app.py

app.py

@@ -0,0 +1,155 @@
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import Response
+from pydantic import BaseModel, Field
+import numpy as np
+import random
+import torch
+from diffusers import DiffusionPipeline
+import io
+import base64
+from PIL import Image
+import uvicorn
+import os
+
+# Initialize FastAPI app
+app = FastAPI(title="FLUX.1 Image Generation API", version="1.0.0")
+
+# Configuration
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
+
+# Load the model
+print("Loading FLUX.1 model...")
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+print("Model loaded successfully!")
+
+# Request models
+class ImageGenerationRequest(BaseModel):
+    prompt: str = Field(..., description="Text prompt for image generation")
+    seed: int = Field(default=42, ge=0, le=MAX_SEED, description="Random seed for generation")
+    randomize_seed: bool = Field(default=False, description="Whether to randomize the seed")
+    width: int = Field(default=1024, ge=256, le=MAX_IMAGE_SIZE, description="Image width")
+    height: int = Field(default=1024, ge=256, le=MAX_IMAGE_SIZE, description="Image height")
+    num_inference_steps: int = Field(default=4, ge=1, le=50, description="Number of inference steps")
+    return_format: str = Field(default="base64", description="Return format: 'base64' or 'bytes'")
+
+class ImageGenerationResponse(BaseModel):
+    image: str = Field(..., description="Generated image in base64 format")
+    seed: int = Field(..., description="Seed used for generation")
+    success: bool = Field(default=True, description="Whether generation was successful")
+
+# Helper functions
+def pil_to_base64(image: Image.Image) -> str:
+    """Convert PIL Image to base64 string"""
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    return img_str
+
+def generate_image(
+    prompt: str,
+    seed: int = 42,
+    randomize_seed: bool = False,
+    width: int = 1024,
+    height: int = 1024,
+    num_inference_steps: int = 4
+):
+    """Generate image using FLUX.1 model"""
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    # Ensure width and height are multiples of 32
+    width = (width // 32) * 32
+    height = (height // 32) * 32
+    
+    generator = torch.Generator().manual_seed(seed)
+    
+    try:
+        image = pipe(
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            guidance_scale=0.0
+        ).images[0]
+        
+        return image, seed
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Image generation failed: {str(e)}")
+
+# API endpoints
+@app.get("/")
+async def root():
+    """Root endpoint with API information"""
+    return {
+        "message": "FLUX.1 Image Generation API",
+        "version": "1.0.0",
+        "model": "black-forest-labs/FLUX.1-schnell",
+        "endpoints": {
+            "generate": "/generate",
+            "generate_image": "/generate/image",
+            "health": "/health"
+        }
+    }
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {"status": "healthy", "device": device, "model_loaded": True}
+
+@app.post("/generate", response_model=ImageGenerationResponse)
+async def generate_image_endpoint(request: ImageGenerationRequest):
+    """Generate image and return as base64"""
+    try:
+        image, used_seed = generate_image(
+            prompt=request.prompt,
+            seed=request.seed,
+            randomize_seed=request.randomize_seed,
+            width=request.width,
+            height=request.height,
+            num_inference_steps=request.num_inference_steps
+        )
+        
+        # Convert to base64
+        image_base64 = pil_to_base64(image)
+        
+        return ImageGenerationResponse(
+            image=image_base64,
+            seed=used_seed,
+            success=True
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/generate/image")
+async def generate_image_bytes(request: ImageGenerationRequest):
+    """Generate image and return as bytes"""
+    try:
+        image, used_seed = generate_image(
+            prompt=request.prompt,
+            seed=request.seed,
+            randomize_seed=request.randomize_seed,
+            width=request.width,
+            height=request.height,
+            num_inference_steps=request.num_inference_steps
+        )
+        
+        # Convert to bytes
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+        
+        return Response(
+            content=img_byte_arr.getvalue(),
+            media_type="image/png",
+            headers={"X-Generated-Seed": str(used_seed)}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)