Update app.py

app.py

@@ -1,8 +1,10 @@
-#pip install fastapi uvicorn transformers
-from fastapi import FastAPI, Request
+#pip install fastapi uvicorn transformers diffusers
+from fastapi import FastAPI, Request, UploadFile, File
 from PIL import Image, ImageDraw, ImageFont
 import torch
+from io import BytesIO
 from diffusers import StableDiffusionImg2ImgPipeline
+
 app = FastAPI()
 
 device = "cuda"
@@ -42,20 +44,14 @@ async def generate_ad(request: Request):
     return {"message": "Welcome to the Ad Template API"}
 
 @app.get("/generate_ad")
-async def generate_ad(request: Request):
-    # Get the input data from the request
-    data = await request.json()
-
-    # Extract the inputs from the data
-    prompt = data["prompt"]
-    hex_code = data["hex_code"]
-    text = data["text"]
-    button_color = data["button_color"]
-    punchline_color = data["punchline_color"]
-    init_image = "flat-white-3402c4f.jpg"
-    # Generate the image using the text-to-image model
+async def generate_ad(prompt: str,hex_code:str, button_color:str, punchline_color:str, image_file: UploadFile):
+    image_bytes = await image_file.read()
+    init_image = Image.open(BytesIO(image_bytes)).convert("RGB")
+    init_image.thumbnail((768, 768))
+    # Generate the image using the img-to-image model
     generator = torch.Generator(device=device).manual_seed(1024)
-    image = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5, generator=generator).images[0]
+    ad_prompt = "From init image create me a new image that will attract customers to put in a (ad template)"
+    image = pipe(prompt=ad_prompt, image=init_image, strength=0.75, guidance_scale=7.5, generator=generator).images[0]
 
     # Create the ad template
     ad_template = create_ad_template(image, hex_code, text, button_color, punchline_color)