Update app.py

app.py

@@ -1,126 +1,85 @@
 import gradio as gr
-import torch
+import requests
+import io
 import random
 import os
-import io
+import time
 from PIL import Image
 from deep_translator import GoogleTranslator
-from diffusers import DiffusionPipeline
-from huggingface_hub import hf_hub_download, login
-
-# Autentikasi Hugging Face
-HF_TOKEN = os.getenv("HF_READ_TOKEN")  # Ganti dengan token Anda atau set env variable
-login(token=HF_TOKEN)
-
-# Konfigurasi Model
-BASE_MODEL = "black-forest-labs/FLUX.1-dev"
-LORA_REPO = "burhansyam/uncen"
-LORA_WEIGHTS_NAME = "uncen.safetensors"  # Ganti jika nama file berbeda
-torch_dtype = torch.float16  # Gunakan float16 untuk kompatibilitas lebih luas
 
-# Inisialisasi Pipeline dengan LoRA
-def init_pipeline():
-    # Muat model dasar
-    pipe = DiffusionPipeline.from_pretrained(
-        BASE_MODEL,
-        torch_dtype=torch_dtype,
-        use_auth_token=HF_TOKEN
-    )
-    
-    # Muat weights LoRA
-    lora_path = hf_hub_download(
-        repo_id=LORA_REPO,
-        filename=LORA_WEIGHTS_NAME,
-        token=HF_TOKEN
-    )
-    pipe.load_lora_weights(lora_path, adapter_name="uncen")
-    
-    # Optimasi GPU jika tersedia
-    if torch.cuda.is_available():
-        pipe.to("cuda")
-        try:
-            pipe.enable_xformers_memory_efficient_attention()
-        except:
-            print("Xformers tidak tersedia, melanjutkan tanpa optimasi")
-    
-    return pipe
+# Project by Nymbo
 
-# Inisialisasi pipeline
-try:
-    pipe = init_pipeline()
-except Exception as e:
-    raise gr.Error(f"Gagal memuat model: {str(e)}. Pastikan token akses valid dan Anda memiliki izin.")
+API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
+API_TOKEN = os.getenv("HF_READ_TOKEN")
+headers = {"Authorization": f"Bearer {API_TOKEN}"}
+timeout = 100
 
 def convert_to_png(image):
-    """Konversi gambar ke format PNG"""
+    """Convert any image format to true PNG format"""
     png_buffer = io.BytesIO()
     if image.mode == 'RGBA':
+        # If image has alpha channel, save as PNG with transparency
         image.save(png_buffer, format='PNG', optimize=True)
     else:
+        # Convert to RGB first if not in RGB/RGBA mode
         if image.mode != 'RGB':
             image = image.convert('RGB')
         image.save(png_buffer, format='PNG', optimize=True)
     png_buffer.seek(0)
     return Image.open(png_buffer)
 
-def generate_image(
-    prompt,
-    negative_prompt="",
-    steps=35,
-    cfg_scale=7,
-    sampler="DPM++ 2M Karras",
-    seed=-1,
-    width=1024,
-    height=1024
-):
+def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", 
+          seed=-1, strength=0.7, width=1024, height=1024):
     if not prompt:
         return None
 
-    # Terjemahkan prompt jika bahasa Indonesia
+    key = random.randint(0, 999)
+    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
+    headers = {"Authorization": f"Bearer {API_TOKEN}"}
+    
+    # Translate prompt
     try:
-        translated_prompt = GoogleTranslator(source='id', target='en').translate(prompt)
-        prompt = f"{translated_prompt} | ultra detail, ultra quality, masterpiece"
-    except:
-        prompt = f"{prompt} | ultra detail, ultra quality, masterpiece"
+        prompt = GoogleTranslator(source='id', target='en').translate(prompt)
+        print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
+        prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
+    except Exception as e:
+        print(f"Translation error: {e}")
     
-    # Set generator dengan seed
-    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
-    if seed == -1:
-        seed = random.randint(1, 1000000000)
-    generator.manual_seed(seed)
+    print(f'\033[1mGeneration {key}:\033[0m {prompt}')
     
-    # Pemetaan sampler ke scheduler
-    sampler_config = {
-        "DPM++ 2M Karras": {"use_karras_sigmas": True},
-        "DPM++ SDE Karras": {"use_karras_sigmas": True},
-        "Euler": {},
-        "Euler a": {"use_karras_sigmas": False},
-        "Heun": {},
-        "DDIM": {}
+    payload = {
+        "inputs": prompt,
+        "is_negative": is_negative,
+        "steps": steps,
+        "cfg_scale": cfg_scale,
+        "seed": seed if seed != -1 else random.randint(1, 1000000000),
+        "strength": strength,
+        "parameters": {"width": width, "height": height}
     }
-    
+
     try:
-        # Generate gambar dengan LoRA
-        result = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
-            generator=generator,
-            cross_attention_kwargs={"scale": 1.0},  # Kontrol kekuatan LoRA
-            **sampler_config[sampler]
-        )
+        response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
+        response.raise_for_status()
+        
+        # Convert directly to PNG without intermediate format
+        img = Image.open(io.BytesIO(response.content))
+        png_img = convert_to_png(img)
         
-        # Konversi ke PNG
-        png_img = convert_to_png(result.images[0])
+        print(f'\033[1mGeneration {key} completed as PNG!\033[0m')
         return png_img
         
+    except requests.exceptions.RequestException as e:
+        print(f"API Error: {e}")
+        if hasattr(e, 'response') and e.response:
+            if e.response.status_code == 503:
+                raise gr.Error("503: Model is loading, please try again later")
+            raise gr.Error(f"{e.response.status_code}: {e.response.text}")
+        raise gr.Error("Network error occurred")
     except Exception as e:
-        raise gr.Error(f"Generasi gambar gagal: {str(e)}")
+        print(f"Image processing error: {e}")
+        raise gr.Error(f"Image processing failed: {str(e)}")
 
-# Antarmuka Gradio
+# Light theme CSS
 css = """
 #app-container {
     max-width: 800px;
@@ -152,7 +111,7 @@ h1 {
 """
 
 with gr.Blocks(theme=gr.themes.Default(primary_hue="green"), css=css) as app:
-    gr.HTML("<center><h1>FLUX.1-Dev with Uncensored LoRA</h1></center>")
+    gr.HTML("<center><h1>FLUX.1-Dev (PNG Output)</h1></center>")
     
     with gr.Column(elem_id="app-container"):
         with gr.Row():
@@ -178,6 +137,7 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="green"), css=css) as app:
                         steps = gr.Slider(35, label="Steps", minimum=10, maximum=100, step=1)
                         cfg = gr.Slider(7.0, label="CFG Scale", minimum=1.0, maximum=20.0, step=0.5)
                     with gr.Row():
+                        strength = gr.Slider(0.7, label="Strength", minimum=0.1, maximum=1.0, step=0.01)
                         seed = gr.Number(-1, label="Seed (-1 for random)")
                     method = gr.Radio(
                         ["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"],
@@ -191,13 +151,13 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="green"), css=css) as app:
             output_image = gr.Image(
                 type="pil",
                 label="Generated PNG Image",
-                format="png",
+                format="png",  # Explicitly set output format
                 elem_id="gallery"
             )
         
         generate_btn.click(
-            fn=generate_image,
-            inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, width, height],
+            fn=query,
+            inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height],
             outputs=output_image
         )