import gradio as gr
import torch
from diffusers import AutoPipelineForImage2Image, LCMScheduler
from PIL import Image

# Lade das Basismodell und den LCM-LoRA Adapter
# Verwende ein unzensiertes Basismodell wie Lykon/dreamshaper-8 oder SG_161222/Realistic_Vision_V5.1_noVAE
# Hier wird 'Lykon/dreamshaper-8' als Beispiel verwendet.
# Für CPU-Nutzung ist torch.float32 besser geeignet.

pipeline = AutoPipelineForImage2Image.from_pretrained(
    "Lykon/dreamshaper-8",
    torch_dtype=torch.float32,
    variant="fp16" # variant="fp16" ist hier nur ein Platzhalter, da es für CPU nicht relevant ist, aber für das Laden des Modells benötigt wird.
)
pipeline.to("cpu")

# Setze den Scheduler
pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)

# Lade den LCM-LoRA Adapter
pipeline.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")

def image_to_image_editing(input_image: Image.Image, prompt: str) -> Image.Image:
    # Skaliere das Bild auf eine passende Größe, falls nötig
    # Stable Diffusion Modelle arbeiten oft am besten mit 512x512 oder 768x768
    input_image = input_image.resize((512, 512))

    # Führe die Inferenz durch
    edited_image = pipeline(
        prompt=prompt,
        image=input_image,
        num_inference_steps=4,  # Wenige Schritte für schnelle CPU-Inferenz
        guidance_scale=1.0,     # Niedriger guidance_scale für LCM-LoRA
        strength=0.6            # Stärke der Bildbearbeitung
    ).images[0]
    return edited_image

# Erstelle die Gradio-Oberfläche
iface = gr.Interface(
    fn=image_to_image_editing,
    inputs=[
        gr.Image(type="pil", label="Eingabebild"),
        gr.Textbox(label="Prompt (Beschreibung der gewünschten Änderung)")
    ],
    outputs=gr.Image(type="pil", label="Bearbeitetes Bild"),
    title="CPU-optimierter unzensierter Bildeditor mit LCM-LoRA",
    description="Lade ein Bild hoch und beschreibe die gewünschte Änderung. Das Modell ist für CPU optimiert und verwendet ein unzensiertes Basismodell mit LCM-LoRA für schnelle Ergebnisse."
)

if __name__ == "__main__":
    iface.launch()