from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, DiffusionPipeline as Pipe
import torch

class Generador:
    def img_to_bytes(image) -> bytes:
        import io
        _imgByteArr = io.BytesIO()
        image.save(_imgByteArr, format="png")
        return _imgByteArr.getvalue()
    def using_runway_sd_15(prompt:str)->bytes:
        try:            
            _generador = Pipe.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
            _generador.to("cuda")
            _imagen = _generador(prompt).images[0]
            _response = bytes(Generador.img_to_bytes(image=_imagen))
        except Exception as e:
            _response = bytes(str(e), 'utf-8')
        finally:
            return _response
    def using_stability_sd_21(prompt:str)->bytes:
        try:
            _generador = Pipe.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
            _generador.to("cuda")
            _imagen = _generador(prompt).images[0]
            _response = bytes(Generador.img_to_bytes(image=_imagen))
        except Exception as e:
            _response = bytes(str(e), 'utf-8')
        finally:
            return _response
    def using_realistic_v14(prompt:str)->bytes:
        try:
            _generador = Pipe.from_pretrained("SG161222/Realistic_Vision_V1.4", torch_dtype=torch.float16)
            _generador.to("cuda")
            _imagen = _generador(prompt).images[0]
            _response = bytes(Generador.img_to_bytes(image=_imagen))
        except Exception as e:
            _response = bytes(str(e), 'utf-8')
        finally:
            return _response
    def using_prompthero_openjourney(prompt:str)->bytes:
        try:
            _generador = Pipe.from_pretrained("prompthero/openjourney", torch_dtype=torch.float16)
            _generador.to("cuda")
            _imagen = _generador(prompt).images[0]
            _response = bytes(Generador.img_to_bytes(image=_imagen))
        except Exception as e:
            print(e)
            _response = bytes(str(e), 'utf-8')
        finally:
            return _response
class Difusor:
    def using_runway_sd_15(prompt:str)->bytes:
        try:
            controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-mlsd")
            _generador = StableDiffusionControlNetPipeline.from_pretrained(
                "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16)
            _generador.to("cuda")
            _imagen = _generador(prompt).images[0]
            _response = bytes(Generador.img_to_bytes(image=_imagen))
        except Exception as e:
            print(e)
            _response = bytes(str(e), 'utf-8')
        finally:
            return _response