SkyOSFullCore / code_runner.py
sky-meilin's picture
Update code_runner.py
980095c verified
Raw
History Blame Contribute Delete
1.58 kB
import torch
from diffusers import AutoPipelineForText2Image, LCMScheduler
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
lcm_lora_id = "latent-consistency/lcm-lora-sdxl"
# Dynamische Geräte-Erkennung (CPU vs. CUDA)
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32
print(f"CodeRunner initialisiert. Nutze Gerät: {device} mit {torch_dtype}")
try:
pipe = AutoPipelineForText2Image.from_pretrained(
model_id,
torch_dtype=torch_dtype,
variant="fp16" if device == "cuda" else None # fp16-Variante läuft auf reinem CPU-Torch oft instabil
)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
# Sicherer Transfer auf das verfügbare Gerät
if device == "cuda":
pipe.to("cuda")
else:
# Überlebenswichtig für CPU Spaces: Reduziert den RAM-Verbrauch drastisch
print("Optimiere Speicher für CPU-Betrieb...")
pipe.enable_attention_slicing()
# LoRA Gewichte laden und fusionieren für maximalen Speedup
pipe.load_lora_weights(lcm_lora_id)
pipe.fuse_lora()
prompt = "Sky Meilin, consistent AI persona, creator system, cinematic portrait"
# Generierung starten
image = pipe(
prompt=prompt,
num_inference_steps=4,
guidance_scale=1.5
).images[0]
image.save("sky_output.png")
print("Bild erfolgreich generiert und als 'sky_output.png' gespeichert.")
except Exception as e:
print(f"Fehler bei der Bildgenerierung in code_runner: {str(e)}")