Update main.py

main.py

@@ -19,6 +19,18 @@ async def generate(prompt):
     result = await bot.ask(prompt=prompt, conversation_style=ConversationStyle.precise)
     return result
 
+def dummy(images, **kwargs):
+    return images, False
+
+async def generate_image(prompt):
+    model_id = "SG161222/Realistic_Vision_V1.4"
+    scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
+    pipe = await StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16)
+    pipe = pipe.to("cuda")
+    pipe.safety_checker = dummy
+    image = await pipe(prompt).images[0]
+    return image
+
 @app.get("/")
 def read_root():
     return "Hello, I'm Linlada"
@@ -29,8 +41,11 @@ def hi(hello: str):
 
 @app.get("/image/{image}")
 def img(image: str):
-    img = Gimager(image)
-    return img
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    result = loop.run_until_complete(generate_image(image))
+    loop.close()
+    return result
     
 @app.get('/linlada/{prompt}')
 def generate_image_route(prompt: str):