Update app.py

app.py

@@ -13,36 +13,37 @@ pipe = StableDiffusionInpaintPipeline.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float32
 )
-# CPU-Modus (gratis, aber langsam); wenn GPU verfügbar, kannst du pipe.to("cuda") nutzen
+# CPU-Modus (gratis, aber langsam); wenn GPU verfügbar, pipe.to("cuda") für schnellere Ergebnisse
 pipe.to("cpu")
 
-# Reduziere die Anzahl Inference-Schritte (Default wäre ~50)
+# Wir reduzieren z. B. die Inference-Steps, damit es schneller geht
 DEFAULT_STEPS = 20
+GUIDANCE_SCALE = 7.5
+IMG_SIZE = 512  # Wir verkleinern Bilder auf 512x512, um CPU-Zeit zu sparen
 
 # ------------------------------------------------------------------------------
-# 2) INPAINT-FUNKTION
+# 2) INPAINT-FUNKTION mit manuellem Masking
 # ------------------------------------------------------------------------------
-def inpaint_room(original_image, prompt):
+def inpaint_masked(img, mask, prompt):
     """
-    Führt ein Inpainting auf dem gesamten Bild durch:
-    - 'mask_image' = original_image => Alles darf überschrieben werden.
-    - Wir verkleinern das Bild auf 512x512 und reduzieren num_inference_steps,
-      um Zeit zu sparen.
+    Nimmt ein Originalbild (img) + eine vom User gezeichnete Maske (mask).
+    Nur in maskierten Bereichen wird etwas verändert, z. B. Möbel eingefügt.
+    Der Rest bleibt unverändert.
     """
-
-    if original_image is None:
+    if img is None or mask is None:
         return None
 
-    # 1) Bild auf 512x512 verkleinern
-    original_image = original_image.resize((512, 512), resample=Image.LANCZOS)
+    # 1) Auf 512x512 verkleinern (Bild & Maske gleich groß halten!)
+    img = img.resize((IMG_SIZE, IMG_SIZE), resample=Image.LANCZOS)
+    mask = mask.resize((IMG_SIZE, IMG_SIZE), resample=Image.LANCZOS)
 
-    # 2) Inference-Aufruf mit weniger Schritten
+    # 2) Pipeline-Aufruf
     result = pipe(
         prompt=prompt,
-        image=original_image,
-        mask_image=original_image,
+        image=img,
+        mask_image=mask,
         num_inference_steps=DEFAULT_STEPS,
-        guidance_scale=7.5  # optionaler Wert, typisch ~7.5
+        guidance_scale=GUIDANCE_SCALE
     ).images[0]
 
     return result
@@ -52,40 +53,48 @@ def inpaint_room(original_image, prompt):
 # ------------------------------------------------------------------------------
 def build_app():
     with gr.Blocks() as demo:
-        gr.Markdown("## 🏠 KI-basiertes Virtual Home Staging (Demo, schnellere Variante)")
+        gr.Markdown("# 🏠 Virtuelles Home Staging mit manueller Maske")
         gr.Markdown(
-            "Lade ein Foto hoch (leer oder wenige Möbel) "
-            "und gib einen Stil an (z.B. 'modern furniture, bright lighting').<br>"
-            "**Achtung:** Auf CPU dauert es evtl. immer noch einige Sekunden."
+            "1. Lade dein Zimmerfoto hoch.\n"
+            "2. **Male in der Maske** die Bereiche aus, wo Möbel erscheinen sollen.\n"
+            "   - Weiß = Hier darf die KI etwas ändern.\n"
+            "   - Schwarz = Bleibt unverändert.\n"
+            "3. Prompt eingeben, auf **Inpaint** klicken."
         )
 
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(
-                    label="Schritt 1: Dein Zimmerfoto",
+                    label="(1) Zimmerfoto hochladen",
+                    type="pil"
+                )
+                mask_image = gr.Image(
+                    label="(2) Maske hier zeichnen (weiß = veränderbar)",
+                    tool="sketch",  # erlaubt freies Zeichnen
                     type="pil"
                 )
                 prompt_text = gr.Textbox(
-                    label="Schritt 2: Was soll eingefügt werden?",
-                    value="modern furniture, bright lighting, minimal decor",
+                    label="(3) Prompt: Was soll eingefügt werden?",
+                    value="Add modern furniture, keep walls the same",
                     lines=2
                 )
-                run_button = gr.Button("Starte Staging")
+                run_button = gr.Button("Inpaint")
 
             with gr.Column():
                 result_image = gr.Image(label="Ergebnis")
 
+        # Button-Event: Klick -> inpaint_masked()
         run_button.click(
-            fn=inpaint_room,
-            inputs=[input_image, prompt_text],
+            fn=inpaint_masked,
+            inputs=[input_image, mask_image, prompt_text],
             outputs=[result_image]
         )
 
         gr.Markdown(
-            "**Hinweis:**\n"
-            "- Dies ist eine schnelle Demo: Das Bild wird verkleinert (512×512) und nur 20 Schritte berechnet.\n"
-            "- Für bessere Qualität ggf. mehr Inference-Steps und höhere Auflösungen, dann dauert es länger.\n"
-            "- Wenn eine Community-GPU verfügbar ist, geht es wesentlich schneller."
+            "**Tipps:**\n"
+            "- Standard-Inpainting erfordert weiße Stellen für Bereiche, die verändert werden dürfen.\n"
+            "- Du kannst z. B. nur den Fußboden weiß maskieren, damit dort Möbel entstehen.\n"
+            "- Für bessere Qualität könntest du `DEFAULT_STEPS` oder `IMG_SIZE` anpassen."
         )
 
     return demo