Create main.py

main.py

@@ -0,0 +1,156 @@
+from fastapi import FastAPI, Form, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, FileResponse
+from diffusers import StableDiffusionPipeline
+import torch
+import uuid
+import base64
+import io
+from PIL import Image
+import os
+
+# Initialize FastAPI app
+app = FastAPI(title="PromptAgro Image Generator API")
+
+# Add CORS middleware to allow frontend connections
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # In production, specify your frontend domains
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load Stable Diffusion LCM model (your original approach)
+print("🚀 Loading Stable Diffusion Model...")
+model_id = "rupeshs/LCM-runwayml-stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+print("✅ Model Loaded.")
+
+@app.get("/")
+async def root():
+    """Health check endpoint"""
+    return {
+        "status": "alive",
+        "service": "PromptAgro Image Generator",
+        "model_loaded": True,
+        "device": "cuda" if torch.cuda.is_available() else "cpu"
+    }
+
+@app.post("/generate/")
+async def generate_image(prompt: str = Form(...)):
+    """
+    Generate product packaging image from input prompt.
+    Returns image file directly (your original approach).
+    """
+    print(f"🖌️ Generating image for prompt: {prompt}")
+
+    # Generate image (your original approach)
+    image = pipe(prompt).images[0]
+
+    # Save image to temp file (your original approach)
+    filename = f"/tmp/{uuid.uuid4().hex}.png"
+    image.save(filename)
+
+    print(f"📦 Image saved to {filename}")
+
+    # Return image file as response (your original approach)
+    return FileResponse(filename, media_type="image/png")
+
+@app.post("/generate-json/")
+async def generate_image_json(
+    prompt: str = Form(...),
+    width: int = Form(512),
+    height: int = Form(512),
+    num_inference_steps: int = Form(4),  # LCM works well with few steps
+    guidance_scale: float = Form(1.0)    # LCM uses low guidance
+):
+    """
+    Generate image and return as JSON with base64 data (for frontend integration).
+    """
+    print(f"🖌️ Generating image for prompt: {prompt}")
+    
+    try:
+        # Generate image with parameters optimized for LCM
+        image = pipe(
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale
+        ).images[0]
+        
+        # Convert image to base64 for JSON response
+        buffer = io.BytesIO()
+        image.save(buffer, format='PNG')
+        img_str = base64.b64encode(buffer.getvalue()).decode()
+        
+        print("✅ Image generated successfully")
+        
+        return JSONResponse({
+            "success": True,
+            "image_data": f"data:image/png;base64,{img_str}",
+            "prompt_used": prompt,
+            "dimensions": {"width": width, "height": height},
+            "steps": num_inference_steps
+        })
+        
+    except Exception as e:
+        print(f"❌ Generation failed: {e}")
+        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
+
+@app.post("/generate-packaging/")
+async def generate_packaging_specific(
+    product_name: str = Form(...),
+    colors: str = Form("green,yellow"),
+    emotion: str = Form("trust"),
+    platform: str = Form("farmers-market")
+):
+    """
+    Generate packaging with PromptAgro-specific prompt engineering
+    """
+    # Create professional prompt for agricultural packaging
+    prompt = f"""Professional agricultural product packaging design for {product_name}, 
+    modern clean style, {colors.replace(',', ' and ')} color scheme, premium typography, 
+    conveying {emotion}, suitable for {platform}, product photography style, 
+    white background, high quality commercial design, realistic packaging mockup, 
+    professional studio lighting, eco-friendly agricultural branding"""
+    
+    prompt = prompt.strip().replace('\n', ' ').replace('  ', ' ')
+    
+    print(f"🎨 Generating packaging for: {product_name}")
+    print(f"📝 Using prompt: {prompt}")
+    
+    try:
+        # Generate with packaging-optimized settings
+        image = pipe(
+            prompt=prompt,
+            width=768,
+            height=768,
+            num_inference_steps=6,
+            guidance_scale=1.5
+        ).images[0]
+        
+        # Convert to base64
+        buffer = io.BytesIO()
+        image.save(buffer, format='PNG')
+        img_str = base64.b64encode(buffer.getvalue()).decode()
+        
+        return JSONResponse({
+            "success": True,
+            "image_data": f"data:image/png;base64,{img_str}",
+            "prompt_used": prompt,
+            "product_name": product_name,
+            "generator": "Stable Diffusion LCM",
+            "cost": "FREE",
+            "processing_time": "~3-5 seconds"
+        })
+        
+    except Exception as e:
+        print(f"❌ Packaging generation failed: {e}")
+        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)