Update main.py

main.py

@@ -1,96 +1,80 @@
 import os
-import warnings
-# Menghindari error protobuf di beberapa environment
-os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 import io
 import base64
 import torch
+import warnings
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
-from diffusers import FluxImg2ImgPipeline
+from diffusers import AutoPipelineForImage2Image
 from PIL import Image
 
+# Sembunyikan warning CUDA karena kita pakai CPU
+warnings.filterwarnings("ignore")
+os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
+
 app = FastAPI()
 
-# 1. Konfigurasi Token & Model
-# Pastikan Anda sudah membuat Secret bernama HF_TOKEN di Settings Space Anda
-HF_TOKEN = os.getenv("HF_TOKEN")
-MODEL_ID = "black-forest-labs/FLUX.1-schnell"
+# Menggunakan SDXL Turbo: Jauh lebih ringan dari Flux untuk CPU
+MODEL_ID = "stabilityai/sdxl-turbo"
 
-print("🚀 Menghubungkan ke Hugging Face...")
+print("🚀 Memuat Mamboro AI Engine (Mode Ringan untuk CPU)...")
 
 try:
-    # Menggunakan FluxImg2ImgPipeline karena bisa menangani Generate dan Edit sekaligus
-    pipe = FluxImg2ImgPipeline.from_pretrained(
-        MODEL_ID, 
-        torch_dtype=torch.float32, # CPU lebih stabil dengan float32
-        token=HF_TOKEN
+    # Memuat model dengan optimasi RAM
+    pipe = AutoPipelineForImage2Image.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float32,
+        variant="fp16" # Menggunakan varian yang lebih kecil
     )
-    # Pindahkan ke CPU
     pipe.to("cpu")
     
-    # Optimasi RAM agar tidak crash di Hugging Face Free Tier
+    # Fitur hemat memori
     pipe.enable_attention_slicing()
     
-    print("✅ Mamboro AI Engine Berhasil Dimuat di CPU")
+    print("✅ Mamboro AI Engine Online!")
 except Exception as e:
     print(f"❌ Gagal memuat model: {e}")
 
-# Fungsi pembantu untuk konversi gambar ke Base64
 def image_to_base64(image):
     buffered = io.BytesIO()
     image.save(buffered, format="JPEG")
     return f"data:image/jpeg;base64,{base64.b64encode(buffered.getvalue()).decode('utf-8')}"
 
-# --- ENDPOINT 1: GENERATE (Text to Image) ---
 @app.post("/generate")
 async def generate(prompt: str = Form(...)):
-    if not prompt:
-        raise HTTPException(status_code=400, detail="Prompt tidak boleh kosong")
-        
     try:
-        # Untuk generate di pipeline Img2Img, kita tidak butuh image_file
-        # Schnell sangat cepat, hanya butuh 4 steps
-        result = pipe(
+        # SDXL Turbo hanya butuh 1-2 langkah di CPU, sangat cepat!
+        image = pipe(
             prompt=prompt, 
-            num_inference_steps=4, 
+            num_inference_steps=2, 
             guidance_scale=0.0
         ).images[0]
-        
-        return {"image": image_to_base64(result)}
+        return {"image": image_to_base64(image)}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-# --- ENDPOINT 2: EDIT (Image to Image / Reference) ---
 @app.post("/edit")
 async def edit(
     prompt: str = Form(...), 
     image_file: UploadFile = File(...),
-    strength: float = Form(0.6) # Nilai 0.1 s/d 0.9
+    strength: float = Form(0.5)
 ):
     try:
-        # Membaca gambar referensi dari HP
         init_content = await image_file.read()
         init_image = Image.open(io.BytesIO(init_content)).convert("RGB")
-        
-        # Perkecil ukuran gambar agar CPU tidak meledak (PENTING!)
-        init_image = init_image.resize((512, 512))
+        init_image = init_image.resize((512, 512)) # Ukuran aman untuk CPU
 
-        # Proses mengubah gambar berdasarkan referensi
-        result = pipe(
+        # Edit gambar berdasarkan referensi
+        image = pipe(
             prompt=prompt, 
             image=init_image, 
             strength=strength, 
-            num_inference_steps=4
+            num_inference_steps=2
         ).images[0]
         
-        return {"image": image_to_base64(result)}
+        return {"image": image_to_base64(image)}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")
-def health_check():
-    return {
-        "status": "Online",
-        "engine": "Mamboro AI (Flux Schnell CPU)",
-        "device": "CPU"
-    }
+def home():
+    return {"status": "Mamboro AI Light-Mode is Ready", "model": "SDXL Turbo"}