Update app.py

app.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI, Request, UploadFile, File, Query
 from PIL import Image, ImageDraw, ImageFont
 import torch
-from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionImg2ImgPipeline, DiffusionPipeline
 from starlette.responses import StreamingResponse
 from fastapi.responses import FileResponse
 import io
@@ -12,11 +12,20 @@ device = "cuda"
 # , torch_dtype=torch.float16
 device1 = "cpu"
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device1)
+pipe1 = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
 
 @app.get("/")
 async def root():
     return {"message": "Welcome to the Creating Ad Template With Stable Diffusion API!"}
-
+    
+@app.get("/generate", response_class=StreamingResponse)
+async def generate_image(prompt: str = Query(..., description="Text prompt for image generation")):
+    image = pipe1(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.post("/generate_ad")
 async def generate_ad(prompt: str,hex_code:str, button_color:str, punchline_color:str, image_file: UploadFile, logo: UploadFile):