Update app.py

app.py

@@ -112,9 +112,48 @@ def convert_panel_to_visual_prompt(panel_text, character_mode, character_desc):
     return f"{base_style} Scene: {panel_text}."
 
 # ============================
-#  IMAGE GENERATOR (FLUX)
+#  IMAGE GENERATORS (3 MODEL)
 # ============================
+def generate_image_pastel_mix(prompt):
+    """
+    fofr/pastel-mix — paling cocok untuk komik pastel
+    """
+    try:
+        out = replicate.run(
+            "fofr/pastel-mix",
+            input={"prompt": prompt}
+        )
+        if isinstance(out, list) and out:
+            return out[0]
+    except:
+        return ""
+    return ""
+
+def generate_image_sdxl(prompt):
+    """
+    stability-ai/sdxl — detail tinggi
+    """
+    try:
+        out = replicate.run(
+            "stability-ai/sdxl",
+            input={
+                "prompt": prompt,
+                "cfg_scale": 7,
+                "steps": 30
+            }
+        )
+        if isinstance(out, dict) and "output" in out and out["output"]:
+            return out["output"][0]
+        if isinstance(out, list) and out:
+            return out[0]
+    except:
+        return ""
+    return ""
+
 def generate_image_flux(prompt):
+    """
+    black-forest-labs/flux-schnell-free — cepat, tapi tidak selalu stabil
+    """
     try:
         out = replicate.run(
             "black-forest-labs/flux-schnell-free",
@@ -126,6 +165,20 @@ def generate_image_flux(prompt):
         return ""
     return ""
 
+def generate_image_by_model(prompt, model_name):
+    """
+    Pilih model berdasarkan dropdown user.
+    """
+    if model_name == "Pastel-Mix":
+        return generate_image_pastel_mix(prompt)
+    elif model_name == "SDXL":
+        return generate_image_sdxl(prompt)
+    elif model_name == "Flux":
+        return generate_image_flux(prompt)
+    else:
+        # fallback ke Pastel-Mix
+        return generate_image_pastel_mix(prompt)
+
 def download_image(url, filename):
     if not url:
         return ""
@@ -157,7 +210,8 @@ class ComicPDF(FPDF):
         self.cell(0, 10, f"Page {self.page_no()}", 0, 0, "C")
 
 def build_comic_pdf(username, goal, genre, tone, style, cover_brief,
-                    chapters_count, character_mode, character_desc, comic_name):
+                    chapters_count, character_mode, character_desc,
+                    comic_name, image_model):
 
     ensure_outputs_dir()
     pdf = ComicPDF(format="A4")
@@ -168,7 +222,7 @@ def build_comic_pdf(username, goal, genre, tone, style, cover_brief,
         f"Cute pastel cartoon cover, soft colors, kid-friendly, Islamic-friendly, "
         f"{cover_brief}, theme: {goal}"
     )
-    cover_url = generate_image_flux(cover_prompt)
+    cover_url = generate_image_by_model(cover_prompt, image_model)
     cover_path = download_image(cover_url, f"cover_{uuid.uuid4().hex}.png")
 
     pdf.add_page()
@@ -224,8 +278,8 @@ def build_comic_pdf(username, goal, genre, tone, style, cover_brief,
                 character_desc
             )
 
-            # Generate image
-            illus_url = generate_image_flux(visual_prompt)
+            # Generate image by selected model
+            illus_url = generate_image_by_model(visual_prompt, image_model)
             illus_path = download_image(
                 illus_url,
                 f"chapter_{idx}_panel_{p_idx}_{uuid.uuid4().hex}.png"
@@ -284,7 +338,7 @@ with gr.Blocks() as demo:
     # MAIN UI
     with gr.Group(visible=False) as main_ui:
 
-        gr.Markdown("# 🌟 AIPromptLab — Comic Generator 2.0 (Cartoon Pastel Edition)")
+        gr.Markdown("# 🌟 AIPromptLab — Comic Generator 2.1 (Cartoon Pastel Edition)")
 
         with gr.Tabs():
 
@@ -362,6 +416,13 @@ with gr.Blocks() as demo:
                     [character_desc]
                 )
 
+                # IMAGE MODEL DROPDOWN
+                image_model = gr.Dropdown(
+                    label="Model Gambar",
+                    choices=["Pastel-Mix", "SDXL", "Flux"],
+                    value="Pastel-Mix"  # default sesuai pilihanmu
+                )
+
                 generate_pdf_btn = gr.Button("🚀 Generate Comic PDF")
                 pdf_output = gr.File(label="Download Comic PDF")
 
@@ -386,7 +447,8 @@ with gr.Blocks() as demo:
     # PDF GENERATOR HANDLER
     def handle_generate_pdf(login_status, login_user,
                             comic_name, goal, genre, tone, style,
-                            chapters, cover_brief, character_mode, character_desc):
+                            chapters, cover_brief, character_mode,
+                            character_desc, image_model):
 
         if not login_status:
             return None
@@ -401,7 +463,8 @@ with gr.Blocks() as demo:
             chapters,
             character_mode,
             character_desc,
-            comic_name
+            comic_name,
+            image_model
         )
         return pdf_path
 
@@ -411,7 +474,7 @@ with gr.Blocks() as demo:
             login_status, login_user,
             comic_name, comic_goal, comic_genre, comic_tone,
             comic_style, chapters_slider, cover_brief,
-            character_mode, character_desc
+            character_mode, character_desc, image_model
         ],
         pdf_output
     )