Update app.py

app.py

@@ -29,16 +29,16 @@ def chat_fn(message, history):
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
 # ---------------- IMAGE ----------------
-img_pipe.to("cpu")
-img_pipe.safety_checker = None
+# ---------- IMAGE MODEL ----------
+from diffusers import StableDiffusionPipeline
+
+img_pipe = StableDiffusionPipeline.from_pretrained(
+    "stabilityai/sdxl-turbo",
+    torch_dtype=torch.float32
+)
+
+img_pipe = img_pipe.to("cpu")
 
-def generate_image(prompt):
-    image = img_pipe(
-        prompt,
-        num_inference_steps=4,
-        guidance_scale=0.0
-    ).images[0]
-    return image
 
 # ---------------- UI ----------------
 with gr.Blocks(