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)}")