Update app.py

app.py

@@ -6,209 +6,109 @@ from huggingface_hub import AsyncInferenceClient
 from translatepy import Translator
 from gradio_client import Client, handle_file
 from PIL import Image
-from lora_saver import main as backup_loras
-
-os.system('quent_models.py')
-
-from loras import loras
-from huggingface_hub import login
-from themes import IndonesiaTheme  # Import custom IndonesiaTheme
-from lorify import Lorify
-from css import css2
+from pathlib import Path
 MAX_SEED = np.iinfo(np.int32).max
 
-
 HF_TOKEN = os.getenv('HF_TOKEN')
 HF_TOKEN_UPSCALER = os.getenv('HF_TOKEN')
 
-qwen_client = Client("K00B404/HugChatWrap",hf_token=HF_TOKEN)
-loaded_loras=[]
-import os
-from pathlib import Path
-
-
 def load_local_loras(lora_directory="lora_models"):
-    """
-    Load loras from local safetensor files instead of remote API
-    Returns list of lora configurations
-    """
+    """Load loras from local safetensor files"""
     loras_list_custom = []
     
-    # Ensure directory exists
     if not os.path.exists(lora_directory):
         print(f"[-] Warning: Lora directory {lora_directory} not found")
         return []
         
-    # Get all safetensor files
     lora_files = list(Path(lora_directory).glob("*.safetensors"))
     
     for lora_file in lora_files:
-        # Get base name without extension
         lora_name = lora_file.stem
-        
-        # Create local lora path
         lora_path = str(lora_file.absolute())
-        
-        # Add to lora list
-        lora_config = {
+        loras_list_custom.append({
             "name": lora_name,
-            "path": lora_path,
-            "type": "local"
-        }
-        loras_list_custom.append(lora_config)
-        
+            "path": lora_path
+        })
+    
+    print(f"[-] Loaded {len(loras_list_custom)} local loras")
     return loras_list_custom
 
-    
-'''for lora in loras:
-    print(lora.get('repo'))
-    loaded_loras.append(lora.get('repo'))
-    
 # Function to enable LoRA if selected
-def enable_lora(lora_add, basemodel):
-    print(f"[-] Determining model: LoRA {'enabled' if lora_add else 'disabled'}, base model: {basemodel}")
-    return basemodel if not lora_add else lora_add
-'''
-
-
-# Replace the existing loras_list_custom with:
-loras_list_custom = load_local_loras()
-
-# Update the Gradio interface dropdown to use the local paths
-lora_model_choice = gr.Dropdown(
-    label="🎨 Select LoRA", 
-    choices=[lora["path"] for lora in loras_list_custom], 
-    value=loras_list_custom[0]["path"] if loras_list_custom else None
-)
-
-# Update the enable_lora function to handle local files
 def enable_lora(lora_path, basemodel):
     print(f"[-] Determining model: LoRA {'enabled' if lora_path else 'disabled'}, base model: {basemodel}")
     if not lora_path:
         return basemodel
         
-    # Load the local safetensor file
+    # Configure model with local lora
     return {
-        "base_model": basemodel,
-        "lora": {
-            "path": lora_path,
-            "scale": 1.0  # You can adjust this scale as needed
-        }
+        "model": basemodel,
+        "lora_weights": lora_path,
+        "lora_scale": 0.75  # Adjust this value as needed
     }
 
-
-
-
-
-
-
 # Function to generate image
-async def generate_image(prompt, model, lora_word, width, height, scales, steps, seed):
+async def generate_image(prompt, model_config, lora_word, width, height, scales, steps, seed):
     try:
         if seed == -1:
             seed = random.randint(0, MAX_SEED)
         seed = int(seed)
 
-        print(f"[-] Menerjemahkan prompt: {prompt}")
+        print(f"[-] Translating prompt: {prompt}")
         text = str(Translator().translate(prompt, 'English')) + "," + lora_word
         
-        print(f"[-] Generating image with prompt: {text}, model: {model}")
-        client = AsyncInferenceClient()
-        image = await client.text_to_image(prompt=text, height=height, width=width, guidance_scale=scales, num_inference_steps=steps, model=model)
+        print(f"[-] Generating image with prompt: {text}")
+        client = AsyncInferenceClient(token=HF_TOKEN)
+        
+        # Handle both simple model string and lora config
+        if isinstance(model_config, dict):
+            print(f"[-] Using model with LoRA: {model_config}")
+            image = await client.text_to_image(
+                prompt=text,
+                height=height,
+                width=width,
+                guidance_scale=scales,
+                num_inference_steps=steps,
+                model=model_config["model"],
+                lora_weights=model_config["lora_weights"],
+                lora_scale=model_config["lora_scale"]
+            )
+        else:
+            print(f"[-] Using base model: {model_config}")
+            image = await client.text_to_image(
+                prompt=text,
+                height=height,
+                width=width,
+                guidance_scale=scales,
+                num_inference_steps=steps,
+                model=model_config
+            )
+            
         return image, seed
     except Exception as e:
         print(f"[-] Error generating image: {e}")
         return None, None
 
-# Function to upscale image
-def get_upscale_finegrain(prompt, img_path, upscale_factor):
-    try:
-        print(f"[-] Memulai proses upscaling dengan faktor {upscale_factor} untuk gambar {img_path}")
-        client = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN)
-        result = client.predict(
-            input_image=handle_file(img_path), 
-            prompt=prompt, 
-            negative_prompt="worst quality, low quality, normal quality",
-            upscale_factor=upscale_factor,
-            controlnet_scale=0.6,
-            controlnet_decay=1,
-            condition_scale=6,
-            denoise_strength=0.35, 
-            num_inference_steps=18,
-            solver="DDIM", 
-            api_name="/process"
-        )
-        print(f"[-] Proses upscaling berhasil.")
-        return result[1]  # Return upscale image path
-    except Exception as e:
-        print(f"[-] Error scaling image: {e}")
-        return None
-        
-# Main function to generate images and optionally upscale
-async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora):
-    print(f"[-] Memulai generasi gambar dengan prompt: {prompt}")
-    
-    model = enable_lora(lora_model, basemodel) if process_lora else basemodel
-    print(f"[-] Menggunakan model: {model}")
-
-    image, seed = await generate_image(prompt, model, "", width, height, scales, steps, seed)
-    
-    if image is None:
-        print("[-] Image generation failed.")
-        return []
-
-    image_path = "temp_image.jpg"
-    print(f"[-] Menyimpan gambar sementara di: {image_path}")
-    image.save(image_path, format="JPEG")
-
-    upscale_image_path = None
-    if process_upscale:
-        print(f"[-] Memproses upscaling dengan faktor: {upscale_factor}")
-        upscale_image_path = get_upscale_finegrain(prompt, image_path, upscale_factor)
-        if upscale_image_path is not None and os.path.exists(upscale_image_path):
-            print(f"[-] Proses upscaling selesai. Gambar tersimpan di: {upscale_image_path}")
-            return [image_path, upscale_image_path]  # Return both images
-        else:
-            print("[-] Upscaling gagal, jalur gambar upscale tidak ditemukan.")
-
-    return [image_path]
-
-base_models=[
-    "black-forest-labs/FLUX.1-schnell", 
-    "black-forest-labs/FLUX.1-DEV", 
-    "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro", 
-    #"city96/FLUX.1-dev-gguf"
-] 
-
-loras_list_custom=[
-    "Keltezaa/anal-riding-missionary",
-    "Keltezaa/Fingering",
-    "Keltezaa/Spreading",
-    "Keltezaa/flux-prone-ass-spread-hd",
-    "Keltezaa/Flux_P",
-    "Shakker-Labs/FLUX.1-dev-LoRA-add-details", 
-    "XLabs-AI/flux-RealismLora",       
-]# + loaded_loras # add loras loaded from file
- 
-backup_loras(loras_list_custom)
+# Load local loras
+local_loras = load_local_loras()
 
 # Creating Gradio interface
 with gr.Blocks(css=css2, theme=IndonesiaTheme()) as WallpaperFluxMaker:
-    # Displaying the application title
     gr.HTML('<div id="banner">✨ Flux MultiMode Generator + Upscaler ✨</div>')
 
     with gr.Column(elem_id="col-container"):
-        # Output section (replacing ImageSlider with gr.Gallery)
         with gr.Row():
-            output_res = gr.Gallery(label="⚡ Flux / Upscaled Image ⚡", elem_id="output-res", columns=2, height="auto")
+            output_res = gr.Gallery(
+                label="⚡ Flux / Upscaled Image ⚡", 
+                elem_id="output-res",
+                columns=2,
+                height="auto"
+            )
 
-        # User input section split into two columns
         with gr.Row():
-            # Column 1: Input prompt, LoRA, and base model
             with gr.Column(scale=1, elem_id="col-left"):
                 prompt = gr.Textbox(
                     label="📜 Description", 
-                    placeholder="Tuliskan prompt Anda dalam bahasa apapun, yang akan langsung diterjemahkan ke bahasa Inggris.",
+                    placeholder="Write your prompt in any language, it will be translated to English.",
                     elem_id="textbox-prompt"
                 )
 
@@ -218,32 +118,38 @@ with gr.Blocks(css=css2, theme=IndonesiaTheme()) as WallpaperFluxMaker:
                     value=base_models[0]
                 )
 
+                # Updated to use local lora paths
                 lora_model_choice = gr.Dropdown(
-                    label="🎨 Select LoRA", 
-                    choices=loras_list_custom, 
-                    value=loras_list_custom[0]
+                    label="🎨 Select LoRA",
+                    choices=[lora["path"] for lora in local_loras],
+                    value=local_loras[0]["path"] if local_loras else None
                 )
 
-                process_lora = gr.Checkbox(label="🎨 Aktifkan LoRA")
-                process_upscale = gr.Checkbox(label="🔍 Aktifkan Peningkatan Resolusi")
-                upscale_factor = gr.Radio(label="🔍 Faktor Peningkatan Resolusi", choices=[2, 4, 8], value=2)
+                process_lora = gr.Checkbox(label="🎨 Enable LoRA")
+                process_upscale = gr.Checkbox(label="🔍 Enable Upscaling")
+                upscale_factor = gr.Radio(
+                    label="🔍 Upscale Factor",
+                    choices=[2, 4, 8],
+                    value=2
+                )
 
-            # Column 2: Advanced options (always open)
             with gr.Column(scale=1, elem_id="col-right"):
-                with gr.Accordion(label="⚙️ Opsi Lanjutan", open=True):
-                    width = gr.Slider(label="Lebar", minimum=512, maximum=1280, step=8, value=1280)
-                    height = gr.Slider(label="Tinggi", minimum=512, maximum=1280, step=8, value=768)
-                    scales = gr.Slider(label="Skala", minimum=1, maximum=20, step=1, value=8)
-                    steps = gr.Slider(label="Langkah", minimum=1, maximum=100, step=1, value=8)
+                with gr.Accordion(label="⚙️ Advanced Options", open=True):
+                    width = gr.Slider(label="Width", minimum=512, maximum=1280, step=8, value=1280)
+                    height = gr.Slider(label="Height", minimum=512, maximum=1280, step=8, value=768)
+                    scales = gr.Slider(label="Scale", minimum=1, maximum=20, step=1, value=8)
+                    steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=8)
                     seed = gr.Number(label="Seed", value=-1)
 
-        # Button to generate image
-        btn = gr.Button("🚀 Buat Gambar", elem_id="generate-btn")
-
-        # Running the `gen` function when "Generate" button is pressed
-        btn.click(fn=gen, inputs=[
-            prompt, basemodel_choice, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model_choice, process_lora
-        ], outputs=output_res)
+        btn = gr.Button("🚀 Generate Image", elem_id="generate-btn")
+        
+        btn.click(
+            fn=gen,
+            inputs=[
+                prompt, basemodel_choice, width, height, scales, steps, seed,
+                upscale_factor, process_upscale, lora_model_choice, process_lora
+            ],
+            outputs=output_res
+        )
 
-# Launching the Gradio app
 WallpaperFluxMaker.queue(api_open=True).launch(show_api=True)