Create app.py

app.py

@@ -0,0 +1,196 @@
+import torch
+from diffusers import Flux2KleinPipeline
+import gradio as gr
+import os
+from datetime import datetime
+from ddgs import DDGS
+import requests
+from PIL import Image
+from io import BytesIO
+import re
+
+torch.cuda.empty_cache()
+
+# =========================
+# ✅ CSS
+# =========================
+custom_css = """
+* { font-family: Arial, sans-serif !important; }
+html, body, .gradio-container { background-color: #0f172a !important; }
+h1, h2, h3, label, span, p { color: #22c55e !important; }
+input, textarea, select {
+    background-color: #1f2937 !important;
+    color: white !important;
+    border-radius: 8px !important;
+    border: 1px solid #22c55e !important;
+}
+button { background: black !important; color: #22c55e !important; }
+"""
+
+# =========================
+# PIPELINE
+# =========================
+pipe = Flux2KleinPipeline.from_pretrained(
+    "lea97338/KTXFlux-2.0",
+    token="hf_xxx",
+    cache_dir="./model",
+    torch_dtype=torch.bfloat16,
+)
+pipe.enable_model_cpu_offload()
+
+pipe.enable_attention_slicing()
+
+pipe.transformer.to(memory_format=torch.channels_last)
+pipe.transformer.eval()
+torch.compile(pipe.transformer)
+pipe.text_encoder.eval()
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.set_float32_matmul_precision("medium")
+
+OUTPUT_DIR = "./outputs"
+
+# =========================
+# RESOLUTION ✅ presets retour
+# =========================
+def get_dimensions(res):
+    return {
+        "256x256": (256,256),
+        "512x512": (512,512),
+        "768x768": (768,768),
+        "Portrait": (512,768),
+        "Landscape": (768,512),
+        "Panorama": (1024,512),
+    }[res]
+
+# =========================
+# CLEAN TEXT
+# =========================
+def clean_text(text):
+    text = re.sub(r"#\w+|http\S+|\d+", "", text.lower())
+    return " ".join([w for w in text.split() if len(w) > 3][:10])
+
+# =========================
+# DDGS ✅ FAST
+# =========================
+def ddgs_search(prompt, width, height):
+
+    texts = []
+    images = []
+
+    with DDGS() as ddgs:
+
+        # TEXTE (2 rapides)
+        for r in ddgs.text(prompt, max_results=3):
+            t = clean_text(r.get("body", ""))
+            if t:
+                texts.append(t)
+            if len(texts) >= 2:
+                break
+
+        # IMAGE (3 rapides)
+        for r in ddgs.images(prompt, max_results=4):
+            try:
+                img = Image.open(
+                    BytesIO(requests.get(r["image"], timeout=3).content)
+                ).convert("RGB")
+
+                # ✅ resize direct → BOOST PERF
+                img = img.resize((width, height), Image.LANCZOS)
+
+                images.append(img)
+
+                if len(images) >= 3:
+                    break
+
+            except:
+                continue
+
+    return " ".join(texts), images
+
+# =========================
+# GENERATE ✅ OPTIMISÉ
+# =========================
+def generate_image(prompt, steps, guidance, resolution):
+
+    torch.cuda.empty_cache()
+
+    width, height = get_dimensions(resolution)
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    # 🔍 DDGS
+    extra_text, images = ddgs_search(prompt, width, height)
+    final_prompt = f"{prompt}, {extra_text}"
+
+    print("🧠 Prompt:", final_prompt)
+    print("🖼 Images:", len(images))
+
+    ref_image = images[0] if images else None
+
+    # ✅ SANS strength (fix crash)
+    if ref_image:
+        result = pipe(
+            prompt=final_prompt,
+            image=ref_image,
+            height=height,
+            width=width,
+            num_inference_steps=int(steps),
+            guidance_scale=float(guidance),
+        )
+    else:
+        result = pipe(
+            prompt=final_prompt,
+            height=height,
+            width=width,
+            num_inference_steps=int(steps),
+            guidance_scale=float(guidance),
+        )
+
+    image = result.images[0]
+
+    path = os.path.join(
+        OUTPUT_DIR,
+        datetime.now().strftime("img_%Y%m%d_%H%M%S.png")
+    )
+    image.save(path)
+
+    return image, path
+
+# =========================
+# UI
+# =========================
+with gr.Blocks(css=custom_css) as demo:
+
+    gr.Markdown("# ⚡ Générateur rapide (DDGS optimisé)")
+
+    with gr.Row():
+        with gr.Column():
+
+            prompt = gr.Textbox(value="minecraft village", label="Prompt")
+
+            steps = gr.Slider(1, 30, 5, step=1, label="Steps")
+            guidance = gr.Slider(1.0, 10.0, 3.0, step=0.5, label="Guidance")
+
+            resolution = gr.Dropdown(
+                ["256x256","512x512","768x768","Portrait","Landscape","Panorama"],
+                value="512x512",
+                label="Résolution"
+            )
+
+            btn = gr.Button("⚡ Générer")
+            path_out = gr.Textbox(label="Fichier")
+
+        with gr.Column():
+            image_out = gr.Image(label="Résultat")
+
+    btn.click(
+        generate_image,
+        inputs=[prompt, steps, guidance, resolution],
+        outputs=[image_out, path_out]
+    )
+
+# =========================
+# RUN
+# =========================
+if __name__ == "__main__":
+    demo.launch(share=True)