Working on it

routers/training.py

@@ -6,12 +6,14 @@ from routers.donut_evaluate import run_evaluate_donut
 from routers.donut_training import run_training_donut
 import utils
 import torch
-from diffusers import StableDiffusionPipeline
+import requests
+from PIL import Image
+from io import BytesIO
 
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionImg2ImgPipeline
 
-pipe = DiffusionPipeline.from_pretrained("radames/stable-diffusion-v1-5-img2img")
-# pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)
+model_id_or_path = "runwayml/stable-diffusion-v1-5"
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
 pipe = pipe.to("cuda")
 
 
@@ -19,7 +21,22 @@ router = APIRouter()
 
 @router.get("/hi")
 async def hifunction():
-    prompt = "a photograph of an astronaut riding a horse"
-    image = pipe(prompt).images[0]
-    print(image)
-    return ["done"]
+    
+    url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+    response = requests.get(url)
+    init_image = Image.open(BytesIO(response.content)).convert("RGB")
+    init_image = init_image.resize((768, 512))
+    prompt = "A fantasy landscape, trending on artstation"
+    images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
+    print(images)
+    print(images[0])
+
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    img_str = base64.b64encode(buffered.getvalue())
+
+    # images[0].save("fantasy_landscape.png")
+
+    return {
+        "image": img_str
+    }