coder160/gpu_machine_learner

1from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, DiffusionPipeline as Pipe2import torch3 4class Generador:5    def img_to_bytes(image) -> bytes:6        import io7        _imgByteArr = io.BytesIO()8        image.save(_imgByteArr, format="png")9        return _imgByteArr.getvalue()10    def using_runway_sd_15(prompt:str)->bytes:11        try:            12            _generador = Pipe.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)13            _generador.to("cuda")14            _imagen = _generador(prompt).images[0]15            _response = bytes(Generador.img_to_bytes(image=_imagen))16        except Exception as e:17            _response = bytes(str(e), 'utf-8')18        finally:19            return _response20    def using_stability_sd_21(prompt:str)->bytes:21        try:22            _generador = Pipe.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)23            _generador.to("cuda")24            _imagen = _generador(prompt).images[0]25            _response = bytes(Generador.img_to_bytes(image=_imagen))26        except Exception as e:27            _response = bytes(str(e), 'utf-8')28        finally:29            return _response30    def using_realistic_v14(prompt:str)->bytes:31        try:32            _generador = Pipe.from_pretrained("SG161222/Realistic_Vision_V1.4", torch_dtype=torch.float16)33            _generador.to("cuda")34            _imagen = _generador(prompt).images[0]35            _response = bytes(Generador.img_to_bytes(image=_imagen))36        except Exception as e:37            _response = bytes(str(e), 'utf-8')38        finally:39            return _response40    def using_prompthero_openjourney(prompt:str)->bytes:41        try:42            _generador = Pipe.from_pretrained("prompthero/openjourney", torch_dtype=torch.float16)43            _generador.to("cuda")44            _imagen = _generador(prompt).images[0]45            _response = bytes(Generador.img_to_bytes(image=_imagen))46        except Exception as e:47            print(e)48            _response = bytes(str(e), 'utf-8')49        finally:50            return _response51class Difusor:52    def using_runway_sd_15(prompt:str)->bytes:53        try:54            controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-mlsd")55            _generador = StableDiffusionControlNetPipeline.from_pretrained(56                "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16)57            _generador.to("cuda")58            _imagen = _generador(prompt).images[0]59            _response = bytes(Generador.img_to_bytes(image=_imagen))60        except Exception as e:61            print(e)62            _response = bytes(str(e), 'utf-8')63        finally:64            return _response