Update app.py

app.py

@@ -12,53 +12,14 @@ device = "cuda"
 device1 = "cpu"
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device1)
 
-def create_ad_template(image, hex_code, text, button_color, punchline_color):
-    # Load the logo image
-    logo_image = Image.open("flat-white-3402c4f.jpg")
-
-    # Create a blank canvas for the ad template
-    ad_template = Image.new("RGB", (800, 600), hex_code)
-
-    # Paste the logo image at the top of the ad template
-    ad_template.paste(logo_image, (0, 0))
-
-    # Paste the generated image in the middle of the ad template
-    ad_template.paste(image, (200, 100))
-
-    # Add the text at the bottom of the ad template
-    draw = ImageDraw.Draw(ad_template)
-    font = ImageFont.truetype("arial.ttf", 24)
-    text_width, text_height = draw.textsize(text, font=font)
-    text_position = (400 - text_width / 2, 500 - text_height / 2)
-    draw.text(text_position, text, font=font, fill=punchline_color)
-
-    # Add the button at the bottom of the ad template
-    button_width, button_height = 200, 50
-    button_position = (400 - button_width / 2, 550 - button_height / 2)
-    draw.rectangle(button_position, fill=button_color)
-
-    return ad_template
-
 @app.get("/")
-async def generate_ad(request: Request):
-    return {"message": "Welcome to the Ad Template API"}
-
-@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 img-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, text, button_color, punchline_color)
-
-    # Return the ad template as a response
-    return ad_template
-
-"""
-The application is now deployed and running. You can access it in your browser at http://localhost:8000/generate_ad. You can send a POST request to this endpoint with the necessary inputs to generate the ad template.
-"""
+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")