| from diffusers import 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 |
| |
