Update main.py

main.py

@@ -1,47 +1,59 @@
+import os
 import io
 import base64
+import torch
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
-from pydantic import BaseModel
 from diffusers import FluxImg2ImgPipeline
 from PIL import Image
-import torch
 
 app = FastAPI()
 
-# Menggunakan model Schnell karena paling efisien untuk CPU
-model_id = "black-forest-labs/FLUX.1-schnell"
+# 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"
 
-print("🚀 Memuat Mamboro Engine di CPU...")
-# FluxImg2ImgPipeline bisa menangani Text-to-Image maupun Image-to-Image
-pipe = FluxImg2ImgPipeline.from_pretrained(
-    model_id, 
-    torch_dtype=torch.float32  # CPU bekerja paling stabil dengan float32
-)
-pipe.to("cpu")
+print("🚀 Menghubungkan ke Hugging Face...")
 
-# Optimasi memori agar tidak crash di RAM terbatas
-pipe.enable_attention_slicing()
+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
+    )
+    # Pindahkan ke CPU
+    pipe.to("cpu")
+    
+    # Optimasi RAM agar tidak crash di Hugging Face Free Tier
+    pipe.enable_attention_slicing()
+    
+    print("✅ Mamboro AI Engine Berhasil Dimuat di CPU")
+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.get_value(buffered.getvalue()).decode('utf-8')}"
+    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 murni, kita beri gambar kosong/noise atau panggil method dasar
-        image = pipe(
+        # Untuk generate di pipeline Img2Img, kita tidak butuh image_file
+        # Schnell sangat cepat, hanya butuh 4 steps
+        result = pipe(
             prompt=prompt, 
             num_inference_steps=4, 
             guidance_scale=0.0
         ).images[0]
         
-        buffered = io.BytesIO()
-        image.save(buffered, format="JPEG")
-        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        return {"image": f"data:image/jpeg;base64,{img_str}"}
+        return {"image": image_to_base64(result)}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -50,28 +62,32 @@ async def generate(prompt: str = Form(...)):
 async def edit(
     prompt: str = Form(...), 
     image_file: UploadFile = File(...),
-    strength: float = Form(0.6)
+    strength: float = Form(0.6) # Nilai 0.1 s/d 0.9
 ):
     try:
-        # Membuka gambar yang dikirim dari aplikasi Android
-        init_image = Image.open(image_file.file).convert("RGB")
-        init_image = init_image.resize((512, 512)) # Standar agar CPU tidak terlalu berat
+        # 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))
 
-        # Proses mengubah gambar berdasarkan prompt
-        image = pipe(
+        # Proses mengubah gambar berdasarkan referensi
+        result = pipe(
             prompt=prompt, 
             image=init_image, 
-            strength=strength, # 0.1 (mirip asli) s/d 0.9 (berubah total)
+            strength=strength, 
             num_inference_steps=4
         ).images[0]
         
-        buffered = io.BytesIO()
-        image.save(buffered, format="JPEG")
-        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        return {"image": f"data:image/jpeg;base64,{img_str}"}
+        return {"image": image_to_base64(result)}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")
 def health_check():
-    return {"status": "Mamboro AI CPU Server is Online"}
+    return {
+        "status": "Online",
+        "engine": "Mamboro AI (Flux Schnell CPU)",
+        "device": "CPU"
+    }