Create models/local_image_model.py

models/local_image_model.py

@@ -0,0 +1,27 @@
+try:
+    from diffusers import StableDiffusionPipeline
+    import torch
+    from io import BytesIO
+except ImportError:
+    StableDiffusionPipeline = None
+    torch = None
+    BytesIO = None
+
+class ImageLocal:
+    def __init__(self):
+        self.pipe = None
+        self.load_model()
+
+    def load_model(self):
+        if StableDiffusionPipeline is None:
+            raise RuntimeError("Diffusers no disponible")
+        self.pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+        if torch.cuda.is_available():
+            self.pipe.to("cuda")
+
+    def generate(self, prompt: str) -> bytes:
+        image = self.pipe(prompt).images[0]
+        buf = BytesIO()
+        image.save(buf, format="PNG")
+        return buf.getvalue()
+