|
|
|
|
|
import os |
|
|
import torch |
|
|
from torchvision import transforms |
|
|
from PIL import Image |
|
|
from huggingface_hub import hf_hub_download |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
|
|
|
MODEL_REPO = "stabilityai/stable-diffusion-2" |
|
|
|
|
|
|
|
|
def load_model(): |
|
|
print(f"Descargando y cargando modelo desde {MODEL_REPO} ...") |
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, trust_remote_code=True) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
MODEL_REPO, |
|
|
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, |
|
|
device_map="auto", |
|
|
trust_remote_code=True |
|
|
) |
|
|
return model, tokenizer |
|
|
|
|
|
MODEL, TOKENIZER = load_model() |
|
|
|
|
|
@torch.inference_mode() |
|
|
def main(config=None, ckpt=None, prompt="A beautiful sci-fi landscape", steps=20, seed=42, out_dir="outputs/"): |
|
|
os.makedirs(out_dir, exist_ok=True) |
|
|
torch.manual_seed(seed) |
|
|
|
|
|
inputs = TOKENIZER(prompt, return_tensors="pt").to(MODEL.device) |
|
|
print(f"Prompt recibido: {prompt}") |
|
|
|
|
|
try: |
|
|
output = MODEL.generate(**inputs, max_new_tokens=steps) |
|
|
text_result = TOKENIZER.decode(output[0], skip_special_tokens=True) |
|
|
print("Resultado del modelo:", text_result) |
|
|
|
|
|
img_path = os.path.join(out_dir, "result_placeholder.png") |
|
|
Image.new("RGB", (720, 320), color=(20, 20, 20)).save(img_path) |
|
|
return img_path |
|
|
|
|
|
except Exception as e: |
|
|
print("Error durante inferencia:", e) |
|
|
return None |
|
|
|
