Update app.py

app.py

@@ -4,7 +4,7 @@ from PIL import Image, ImageDraw, ImageFont
 import torch
 from io import BytesIO
 from diffusers import StableDiffusionImg2ImgPipeline
-from starlette.responses import StreamingResponse
+#from starlette.responses import StreamingResponse
 
 app = FastAPI()
 
@@ -17,10 +17,19 @@ pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion
 async def root():
     return {"message": "Welcome to the Text-to-Image API!"}
 
-@app.get("/generate", response_class=StreamingResponse)
-async def generate_image(prompt: str = Query(..., description="Text prompt for image generation")):
-    image = pipe(prompt).images[0]
-    image_data = io.BytesIO()
-    image.save(image_data, format="PNG")
-    image_data.seek(0)
-    return StreamingResponse(image_data, media_type="image/png")
+@app.get("/generate_ad")
+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 text-to-image model
+    generator = torch.Generator(device=device).manual_seed(1024)
+    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, prompt, button_color, punchline_color)
+
+    # Return the ad template as a response
+    return ad_template