feat: FLUX.1-schnell image generation after correct answers

app.py

@@ -14,6 +14,7 @@ try:
     from core.chunker import chunk_text
     from core.questioner import generate_question, generate_mcq
     from core.evaluator import evaluate_answer
+    from core.image_gen import generate_concept_image
     print("✅ All imports successful")
 except Exception as e:
     import traceback
@@ -170,6 +171,13 @@ body { background: #0A0F1E !important; margin: 0 !important; }
     pointer-events: none !important;
 }
 
+/* ── Hidden concept image component ─────────────────────────────────────────── */
+#hidden-concept-image {
+    position: fixed !important;
+    top: -9999px !important;
+    opacity: 0 !important;
+}
+
 /* ── Load PDF button — defeat Gradio's orange primary override ──────────────── */
 .gr-button-primary,
 button.primary,
@@ -337,6 +345,29 @@ def submit_answer(question, chunk, student_answer, correct, total, language):
         return f"Erreur : {e}", correct, total, _score_label(correct, total)
 
 
+def _is_positive_verdict(feedback: str) -> bool:
+    """True for Correct or Partially Correct verdicts (not Incorrect / errors)."""
+    for line in feedback.splitlines():
+        low = line.lower()
+        if "verdict" in low:
+            return "correct" in low and "incorrect" not in low
+    return False
+
+
+@spaces.GPU(duration=60)
+def generate_image_fn(question, feedback):
+    if not question or not _is_positive_verdict(feedback):
+        return None
+    try:
+        concept = question.strip().replace("\n", " ")[:200]
+        print(f"[generate_image_fn] generating illustration for: {concept[:80]!r}")
+        return generate_concept_image(concept)
+    except Exception as e:
+        import traceback; traceback.print_exc()
+        print(f"[generate_image_fn] failed: {e}")
+        return None
+
+
 # ---------------------------------------------------------------------------
 # UI
 # ---------------------------------------------------------------------------
@@ -353,6 +384,7 @@ BRIDGE_JS = """() => {
     mode: 'open',
     mcqData: null,
     waitingMCQ: false, prevMCQ: '',
+    lastConceptImg: '',
   };
   const $ = id => document.getElementById(id);
 
@@ -456,6 +488,21 @@ BRIDGE_JS = """() => {
     return verdict;
   }
 
+  // ── Concept image ──────────────────────────────────────────────────────────
+  function hideConceptImage() {
+    const wrap = $('concept-image-display');
+    if (wrap) { wrap.classList.add('hidden'); wrap.classList.remove('fade-in-up'); }
+  }
+
+  function showConceptImage(src) {
+    const wrap = $('concept-image-display'), img = $('concept-image');
+    if (!wrap || !img) return;
+    img.src = src;
+    wrap.classList.remove('hidden');
+    void wrap.offsetWidth; // restart fade-in animation
+    wrap.classList.add('fade-in-up');
+  }
+
   // ── Question display ───────────────────────────────────────────────────────
   function showQuestion(text) {
     S.currentQuestion = text; S.waitingQ = false;
@@ -492,6 +539,7 @@ BRIDGE_JS = """() => {
     $('empty-state')?.classList.remove('hidden');
     $('question-area')?.classList.add('hidden');
     $('feedback-card')?.classList.add('hidden');
+    hideConceptImage();
   }
 
   // ── MCQ display ────────────────────────────────────────────────────────────
@@ -572,6 +620,7 @@ BRIDGE_JS = """() => {
     if (es) es.classList.add('hidden');
     if (qa) qa.classList.remove('hidden');
     if (eb) eb.classList.remove('hidden');
+    hideConceptImage();
     if (S.mode === 'mcq') {
       S.waitingMCQ = true;
       S.prevMCQ = document.querySelector('#hidden-mcq-output textarea')?.value || '';
@@ -642,6 +691,7 @@ BRIDGE_JS = """() => {
         if(ai){ai.value='';} if(cc) cc.textContent='0 chars'; if(sb) sb.disabled=true;
         if(fc) fc.classList.add('hidden'); if(es) es.classList.remove('hidden');
         if(qa) qa.classList.add('hidden'); if(eb) eb.classList.add('hidden');
+        hideConceptImage();
         if(qc) qc.scrollIntoView({behavior:'smooth',block:'start'});
       }],
     ];
@@ -747,6 +797,16 @@ BRIDGE_JS = """() => {
         updateScore();
       }
     }
+
+    // Concept image: Gradio writes the generated image into the hidden
+    // gr.Image — surface it in the custom UI when a new src appears.
+    const conceptEl = document.querySelector('#hidden-concept-image img');
+    if (conceptEl && conceptEl.src && conceptEl.src !== S.lastConceptImg) {
+      S.lastConceptImg = conceptEl.src;
+      // Only show if feedback is still on screen (user hasn't moved on)
+      const fc = $('feedback-card');
+      if (fc && !fc.classList.contains('hidden')) showConceptImage(conceptEl.src);
+    }
   }, 300);
 }"""
 
@@ -777,6 +837,7 @@ with gr.Blocks(title="PaperProf", css=CSS) as demo:
     answer_box       = gr.Textbox(label="✏️ Your Answer", lines=4, placeholder="Write your answer here…", elem_id="hidden-answer-input")
     submit_btn       = gr.Button("Submit Answer", variant="primary", elem_id="hidden-submit-btn")
     feedback_box     = gr.Textbox(label="💬 Feedback", interactive=False, lines=7, elem_id="hidden-feedback-output")
+    concept_image    = gr.Image(label="concept", type="pil", interactive=False, visible=True, elem_id="hidden-concept-image")
     question_trigger = gr.Textbox(value="0",       label="qtrig",   elem_id="hidden-question-trigger")
     answer_trigger   = gr.Textbox(value="0",       label="atrig",   elem_id="hidden-answer-trigger")
     chunk_display    = gr.Textbox(value="",        label="chunk",   elem_id="hidden-chunk-output")
@@ -800,6 +861,10 @@ with gr.Blocks(title="PaperProf", css=CSS) as demo:
         submit_answer,
         inputs=[question_box, chunk_state, answer_box, correct_state, total_state, language_box],
         outputs=[feedback_box, correct_state, total_state, score_box],
+    ).then(
+        generate_image_fn,
+        inputs=[question_box, feedback_box],
+        outputs=[concept_image],
     )
 
     mcq_trigger.change(

core/image_gen.py

@@ -0,0 +1,40 @@
+"""
+Concept image generation via FLUX.1-schnell.
+"""
+
+from functools import lru_cache
+
+import torch
+from PIL import Image
+
+
+@lru_cache(maxsize=1)
+def _load_pipeline():
+    # Imported lazily so the app can start (and CUDA stays untouched)
+    # outside the @spaces.GPU context that calls generate_concept_image.
+    from diffusers import FluxPipeline
+
+    pipe = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-schnell",
+        torch_dtype=torch.float16,
+    )
+    pipe.to("cuda")
+    return pipe
+
+
+def generate_concept_image(concept: str) -> Image.Image:
+    pipe = _load_pipeline()
+    prompt = (
+        f"Educational illustration of: {concept}, "
+        "minimalist flat design, dark background, purple and cyan colors, "
+        "scientific diagram style"
+    )
+    result = pipe(
+        prompt,
+        height=512,
+        width=512,
+        num_inference_steps=4,
+        guidance_scale=0.0,
+        max_sequence_length=256,
+    )
+    return result.images[0]

requirements.txt

@@ -10,6 +10,9 @@ spaces>=0.50.0
 # PDF extraction
 PyMuPDF>=1.24.0
 
+# Image generation
+diffusers>=0.30.0
+
 # LLM inference
 torch>=2.3.0
 transformers==4.57.1

ui/index.html

@@ -281,6 +281,22 @@
     .feedback-body .section-content { color: var(--text); }
     .feedback-raw { white-space: pre-wrap; }
 
+    /* ── Concept image card ───────────────────────────────────────────────── */
+    #concept-image-display { text-align: center; }
+    .concept-image-title {
+      font-size: .72rem; font-weight: 700; letter-spacing: .1em;
+      text-transform: uppercase; color: var(--cyan-l);
+      margin-bottom: 14px; text-align: left;
+    }
+    #concept-image {
+      width: 100%; max-width: 512px; display: inline-block;
+      border-radius: var(--r-sm); border: 1px solid var(--border);
+      box-shadow: 0 8px 24px rgba(0,0,0,.4);
+    }
+    .concept-image-caption {
+      margin-top: 10px; font-size: .8rem; color: var(--muted);
+    }
+
     /* ── Collapsible source ───────────────────────────────────────────────── */
     details.source-details { margin-top: 18px; }
     details.source-details summary {
@@ -534,6 +550,13 @@
       <div class="feedback-body" id="feedback-body"></div>
       <div class="mcq-explanations hidden" id="mcq-explanations"></div>
     </div>
+
+    <!-- ── Concept image card ────────────────────────────────────────────── -->
+    <div class="card hidden" id="concept-image-display">
+      <div class="concept-image-title">Visual Summary</div>
+      <img id="concept-image" alt="Visual summary of this concept" />
+      <p class="concept-image-caption">Visual summary of this concept</p>
+    </div>
   </div>
 
   <!-- ── Session modal ───────────────────────────────────────────────────── -->