Update app.py

app.py

@@ -11,6 +11,7 @@ MAX_SEED = np.iinfo(np.int32).max
 
 HF_TOKEN = os.getenv('HF_TOKEN')
 HF_TOKEN_UPSCALER = os.getenv('HF_TOKEN')
+
 css2="""
 /* Apply dark theme (black background) */
 body {
@@ -157,7 +158,6 @@ body {
 
 """
 
-
 # Define base models
 base_models = [
     "black-forest-labs/FLUX.1-schnell", 
@@ -165,14 +165,13 @@ base_models = [
     "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro"
 ]
 
-
 def load_local_loras(lora_directory="lora_models"):
     """Load loras from local safetensor files"""
     loras_list_custom = []
     
     if not os.path.exists(lora_directory):
-        print(f"[-] Warning: Lora directory {lora_directory} not found")
-        return []
+        os.makedirs(lora_directory)
+        print(f"[-] Created lora directory: {lora_directory}")
         
     lora_files = list(Path(lora_directory).glob("*.safetensors"))
     
@@ -200,6 +199,30 @@ def enable_lora(lora_path, basemodel):
         "lora_scale": 0.75  # Adjust this value as needed
     }
 
+# Function to upscale image
+def get_upscale_finegrain(prompt, img_path, upscale_factor):
+    try:
+        print(f"[-] Starting upscaling process with factor {upscale_factor} for image {img_path}")
+        client = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER)
+        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"[-] Upscaling successful.")
+        return result[1]  # Return upscale image path
+    except Exception as e:
+        print(f"[-] Error scaling image: {e}")
+        return None
+
 # Function to generate image
 async def generate_image(prompt, model_config, lora_word, width, height, scales, steps, seed):
     try:
@@ -242,11 +265,40 @@ async def generate_image(prompt, model_config, lora_word, width, height, scales,
         print(f"[-] Error generating image: {e}")
         return None, 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"[-] Starting image generation with prompt: {prompt}")
+    
+    model = enable_lora(lora_model, basemodel) if process_lora else basemodel
+    print(f"[-] Using 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"[-] Saving temporary image to: {image_path}")
+    image.save(image_path, format="JPEG")
+
+    upscale_image_path = None
+    if process_upscale:
+        print(f"[-] Processing upscaling with factor: {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"[-] Upscaling complete. Image saved at: {upscale_image_path}")
+            return [image_path, upscale_image_path]  # Return both images
+        else:
+            print("[-] Upscaling failed, upscaled image path not found.")
+
+    return [image_path]
+
 # Load local loras
 local_loras = load_local_loras()
 
 # Creating Gradio interface
-with gr.Blocks(css=css2) as WallpaperFluxMaker:
+with gr.Blocks(css=css2, theme=IndonesiaTheme()) as WallpaperFluxMaker:
     gr.HTML('<div id="banner">✨ Flux MultiMode Generator + Upscaler ✨</div>')
 
     with gr.Column(elem_id="col-container"):