Update app.py

app.py

@@ -1,23 +1,34 @@
 # Import necessary libraries
 import torch
 from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
-import gradio as gr
-from PIL import Image
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 import io
-import tempfile
-import os
+import base64
+from PIL import Image
 import time
 
+# Initialize FastAPI
+app = FastAPI(title="Children's Book Illustrator API")
+
+# Add CORS middleware to allow requests from n8n
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 # Force CPU usage
 device = "cpu"
 print(f"Using device: {device}")
 
-# Load a HIGHER QUALITY model
+# Load model
 model_id = "stabilityai/stable-diffusion-2-1"
-
 print("Loading pipeline... This may take a few minutes.")
+
 try:
-    # Use torch.float32 for CPU compatibility
     pipe = StableDiffusionPipeline.from_pretrained(
         model_id,
         torch_dtype=torch.float32,
@@ -25,17 +36,12 @@ try:
         safety_checker=None,
         requires_safety_checker=False
     )
-    
-    # Use a faster scheduler for quicker generation
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    
-    # Move the pipeline to the CPU
     pipe = pipe.to(device)
     print("Model loaded successfully on CPU!")
-
 except Exception as e:
     print(f"Error loading model: {e}")
-    # Fallback to original model if needed
+    # Fallback
     model_id = "dreamlike-art/dreamlike-diffusion-1.0"
     pipe = StableDiffusionPipeline.from_pretrained(
         model_id,
@@ -47,59 +53,59 @@ except Exception as e:
     pipe = pipe.to(device)
     print(f"Fell back to {model_id}")
 
-# Define the image generation function
-def generate_image(prompt):
-    """
-    This function takes a text prompt and returns a generated image.
-    """
-    # PROFESSIONAL-QUALITY PROMPT ENHANCEMENT
-    enhanced_prompt = f"masterpiece, best quality, 4K, ultra detailed, photorealistic, sharp focus, studio lighting, professional photography, {prompt}"
-    
-    # Remove any negative aspects (optional but improves quality)
-    negative_prompt = "blurry, low quality, low resolution, watermark, signature, text, ugly, deformed"
-    
-    print(f"Generating image for prompt: {enhanced_prompt}")
-    
-    # Generate the image with better settings
-    image = pipe(
-        prompt=enhanced_prompt,
-        negative_prompt=negative_prompt,
-        width=512,
-        height=512,
-        guidance_scale=9.0,
-        num_inference_steps=25,
-        generator=torch.Generator(device=device)
-    ).images[0]
-    
-    # Convert to RGB to ensure proper color format
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-    
-    print("Image generated successfully!")
-    
-    # Create a temporary file and save the image
-    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp_file:
-        image.save(tmp_file, format='PNG')
-        tmp_file_path = tmp_file.name
-    
-    return tmp_file_path
+# Request model
+class GenerateRequest(BaseModel):
+    prompt: str
+    width: int = 512
+    height: int = 512
+    steps: int = 25
 
-# Create the Gradio Interface
-demo = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Textbox(
-        label="Enter your scene description",
-        lines=2,
-        placeholder="A dragon reading a book under a magical tree"
-    ),
-    outputs=gr.File(label="Download Generated Illustration"),
-    title="Premium Children's Book Illustrator 🤖🎨",
-    description="Generating high-quality, sharp, detailed images for your stories. Enter a scene description."
-)
+# Health check endpoint
+@app.get("/")
+async def health_check():
+    return {"status": "healthy", "model": model_id}
+
+# Main API endpoint
+@app.post("/generate")
+async def generate_image(request: GenerateRequest):
+    try:
+        # Enhanced prompt
+        enhanced_prompt = f"masterpiece, best quality, 4K, ultra detailed, photorealistic, sharp focus, studio lighting, professional photography, {request.prompt}"
+        negative_prompt = "blurry, low quality, low resolution, watermark, signature, text, ugly, deformed"
+        
+        print(f"Generating image for prompt: {enhanced_prompt}")
+        
+        # Generate image
+        image = pipe(
+            prompt=enhanced_prompt,
+            negative_prompt=negative_prompt,
+            width=request.width,
+            height=request.height,
+            guidance_scale=9.0,
+            num_inference_steps=request.steps,
+            generator=torch.Generator(device=device)
+        ).images[0]
+        
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        print("Image generated successfully!")
+        
+        # Convert to base64 for API response
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        img_base64 = base64.b64encode(buffered.getvalue()).decode()
+        
+        return {
+            "status": "success",
+            "image": f"data:image/png;base64,{img_base64}",
+            "prompt": request.prompt
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
 
-# Launch the app
-demo.launch(
-    debug=True,
-    server_name="0.0.0.0",
-    share=False
-)
+# Run the app
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)