Update app.py

app.py

@@ -1,144 +1,16 @@
 import gradio as gr
-import numpy as np
-import random
-import torch
-import spaces
+import replicate
 import os
-import json
-import time
-
 from PIL import Image, ImageDraw
-import torch 
-import math
-
-from optimization import optimize_pipeline_
-from qwenimage.pipeline_qwen_image_edit import QwenImageEditPipeline
-from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
-from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
-
-from huggingface_hub import InferenceClient
-import math
-
-# --- Prompt Enhancement using Hugging Face InferenceClient ---
-def polish_prompt_hf(original_prompt, system_prompt):
-    """
-    Rewrites the prompt using a Hugging Face InferenceClient.
-    """
-    # Ensure HF_TOKEN is set
-    api_key = os.environ.get("HF_TOKEN")
-    if not api_key:
-        print("Warning: HF_TOKEN not set. Falling back to original prompt.")
-        return original_prompt
-
-    try:
-        # Initialize the client
-        client = InferenceClient(
-            provider="cerebras",
-            api_key=api_key,
-        )
-
-        # Format the messages for the chat completions API
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": original_prompt}
-        ]
-
-        # Call the API
-        completion = client.chat.completions.create(
-            model="Qwen/Qwen3-235B-A22B-Instruct-2507",
-            messages=messages,
-        )
-        
-        # Parse the response
-        result = completion.choices[0].message.content
-        
-        # Try to extract JSON if present
-        if '{"Rewritten"' in result:
-            try:
-                # Clean up the response
-                result = result.replace('```json', '').replace('```', '')
-                result_json = json.loads(result)
-                polished_prompt = result_json.get('Rewritten', result)
-            except:
-                polished_prompt = result
-        else:
-            polished_prompt = result
-            
-        polished_prompt = polished_prompt.strip().replace("\n", " ")
-        return polished_prompt
-        
-    except Exception as e:
-        print(f"Error during API call to Hugging Face: {e}")
-        # Fallback to original prompt if enhancement fails
-        return original_prompt
-
-
-def polish_prompt(prompt, img):
-    """
-    Main function to polish prompts for image editing using HF inference.
-    """
-    SYSTEM_PROMPT = '''
-# Edit Instruction Rewriter
-You are a professional edit instruction rewriter. Your task is to generate a precise, concise, and visually achievable professional-level edit instruction based on the user-provided instruction and the image to be edited.  
-
-Please strictly follow the rewriting rules below:
-
-## 1. General Principles
-- Keep the rewritten prompt **concise**. Avoid overly long sentences and reduce unnecessary descriptive language.  
-- If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary.  
-- Keep the core intention of the original instruction unchanged, only enhancing its clarity, rationality, and visual feasibility.  
-- All added objects or modifications must align with the logic and style of the edited input image's overall scene.  
-
-## 2. Task Type Handling Rules
-### 1. Add, Delete, Replace Tasks
-- If the instruction is clear (already includes task type, target entity, position, quantity, attributes), preserve the original intent and only refine the grammar.  
-- If the description is vague, supplement with minimal but sufficient details (category, color, size, orientation, position, etc.). For example:  
-    > Original: "Add an animal"  
-    > Rewritten: "Add a light-gray cat in the bottom-right corner, sitting and facing the camera"  
-- Remove meaningless instructions: e.g., "Add 0 objects" should be ignored or flagged as invalid.  
-- For replacement tasks, specify "Replace Y with X" and briefly describe the key visual features of X.  
-
-### 2. Text Editing Tasks
-- All text content must be enclosed in English double quotes " ". Do not translate or alter the original language of the text, and do not change the capitalization.  
-- **For text replacement tasks, always use the fixed template:**
-    - Replace "xx" to "yy".  
-    - Replace the xx bounding box to "yy".  
-- If the user does not specify text content, infer and add concise text based on the instruction and the input image's context. For example:  
-    > Original: "Add a line of text" (poster)  
-    > Rewritten: "Add text "LIMITED EDITION" at the top center with slight shadow"  
-- Specify text position, color, and layout in a concise way.  
-
-### 3. Human Editing Tasks
-- Maintain the person's core visual consistency (ethnicity, gender, age, hairstyle, expression, outfit, etc.).  
-- If modifying appearance (e.g., clothes, hairstyle), ensure the new element is consistent with the original style.  
-- **For expression changes, they must be natural and subtle, never exaggerated.**  
-- If deletion is not specifically emphasized, the most important subject in the original image (e.g., a person, an animal) should be preserved.
-    - For background change tasks, emphasize maintaining subject consistency at first.  
-- Example:  
-    > Original: "Change the person's hat"  
-    > Rewritten: "Replace the man's hat with a dark brown beret; keep smile, short hair, and gray jacket unchanged"  
-
-### 4. Style Transformation or Enhancement Tasks
-- If a style is specified, describe it concisely with key visual traits. For example:  
-    > Original: "Disco style"  
-    > Rewritten: "1970s disco: flashing lights, disco ball, mirrored walls, colorful tones"  
-- If the instruction says "use reference style" or "keep current style," analyze the input image, extract main features (color, composition, texture, lighting, art style), and integrate them concisely.  
-- **For coloring tasks, including restoring old photos, always use the fixed template:** "Restore old photograph, remove scratches, reduce noise, enhance details, high resolution, realistic, natural skin tones, clear facial features, no distortion, vintage photo restoration"  
-- If there are other changes, place the style description at the end.
-
-## 3. Rationality and Logic Checks
-- Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" should be logically corrected.  
-- Add missing key information: if position is unspecified, choose a reasonable area based on composition (near subject, empty space, center/edges).  
+import requests
+from io import BytesIO
+import time
+import tempfile
+import base64
+import numpy as np
 
-# Output Format
-Return only the rewritten instruction text directly, without JSON formatting or any other wrapper.
-'''
-    
-    # Note: We're not actually using the image in the HF version, 
-    # but keeping the interface consistent
-    full_prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {prompt}\n\nRewritten Prompt:"
-    
-    return polish_prompt_hf(full_prompt, SYSTEM_PROMPT)
+# Set up Replicate API key from environment variable
+os.environ['REPLICATE_API_TOKEN'] = os.getenv('REPLICATE_API_TOKEN')
 
 # --- Outpainting Functions ---
 def can_expand(source_width, source_height, target_width, target_height, alignment):
@@ -246,8 +118,11 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     
     return background, mask
 
-def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
-    """Creates a preview showing the mask overlay."""
+def preview_outpaint(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
+    """Creates a preview showing the mask overlay for outpainting."""
+    if not image:
+        return None
+    
     background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
     
     # Create a preview image showing the mask
@@ -265,248 +140,508 @@ def preview_image_and_mask(image, width, height, overlap_percentage, resize_opti
     
     return preview
 
-# --- Model Loading ---
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=dtype).to(device)
-pipe.transformer.__class__ = QwenImageTransformer2DModel
-pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
-
-# --- Ahead-of-time compilation ---
-optimize_pipeline_(pipe, image=Image.new("RGB", (1024, 1024)), prompt="prompt")
-
-# --- UI Constants and Helpers ---
-MAX_SEED = np.iinfo(np.int32).max
-
-def clear_result():
-    """Clears the result image."""
-    return gr.update(value=None)
-
-def update_history(new_image, history):
-    """Updates the history gallery with the new image."""
-    time.sleep(0.5)  # Small delay to ensure image is ready
-    if history is None:
-        history = []
-    if new_image is not None:
-        # Convert to list if needed (Gradio sometimes returns tuples)
-        if not isinstance(history, list):
-            history = list(history) if history else []
-        history.insert(0, new_image)
-    # Keep only the last 20 images in history
-    history = history[:20]
-    return history
-
-def use_history_as_input(evt: gr.SelectData, history):
-    """Sets the selected history image as the new input image."""
-    if history and evt.index < len(history):
-        return gr.update(value=history[evt.index][0])
-    return gr.update()
-
-def use_output_as_input(output_image):
-    """Sets the generated output as the new input image."""
-    if output_image is not None:
-        return gr.update(value=output_image)
-    return gr.update()
-
-def preload_presets(target_ratio, ui_width, ui_height):
-    """Updates the width and height sliders based on the selected aspect ratio."""
-    if target_ratio == "9:16":
-        changed_width = 720
-        changed_height = 1280
-        return changed_width, changed_height, gr.update()
-    elif target_ratio == "16:9":
-        changed_width = 1280
-        changed_height = 720
-        return changed_width, changed_height, gr.update()
-    elif target_ratio == "1:1":
-        changed_width = 1024
-        changed_height = 1024
-        return changed_width, changed_height, gr.update()
-    elif target_ratio == "Custom":
-        return ui_width, ui_height, gr.update(open=True)
-
-def select_the_right_preset(user_width, user_height):
-    if user_width == 720 and user_height == 1280:
-        return "9:16"
-    elif user_width == 1280 and user_height == 720:
-        return "16:9"
-    elif user_width == 1024 and user_height == 1024:
-        return "1:1"
-    else:
-        return "Custom"
-
-def toggle_custom_resize_slider(resize_option):
-    return gr.update(visible=(resize_option == "Custom"))
-
-# --- Main Inference Function (with outpainting preprocessing) ---
-@spaces.GPU(duration=120)
-def infer(
-    image,
-    prompt,
-    width,
-    height,
-    overlap_percentage,
-    resize_option,
-    custom_resize_percentage,
-    alignment,
-    overlap_left,
-    overlap_right,
-    overlap_top,
-    overlap_bottom,
-    seed=42,
-    randomize_seed=False,
-    true_guidance_scale=4.0,
-    num_inference_steps=50,
-    rewrite_prompt=True,
-    progress=gr.Progress(track_tqdm=True),
-):
+# --- Image Upload Functions ---
+def upload_image_to_hosting(image):
     """
-    Generates an outpainted image using the Qwen-Image-Edit pipeline.
+    Upload image to multiple hosting services with fallback
     """
-    # Hardcode the negative prompt as requested
-    negative_prompt = " "
+    # Method 1: Try imgbb.com (most reliable)
+    try:
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        buffered.seek(0)
+        img_base64 = base64.b64encode(buffered.getvalue()).decode()
+        
+        response = requests.post(
+            "https://api.imgbb.com/1/upload",
+            data={
+                'key': '6d207e02198a847aa98d0a2a901485a5',
+                'image': img_base64,
+            }
+        )
+        
+        if response.status_code == 200:
+            data = response.json()
+            if data.get('success'):
+                return data['data']['url']
+    except:
+        pass
+    
+    # Method 2: Try 0x0.st (simple and reliable)
+    try:
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        buffered.seek(0)
+        
+        files = {'file': ('image.png', buffered, 'image/png')}
+        response = requests.post("https://0x0.st", files=files)
+        
+        if response.status_code == 200:
+            return response.text.strip()
+    except:
+        pass
+    
+    # Method 3: Fallback to base64
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    buffered.seek(0)
+    img_base64 = base64.b64encode(buffered.getvalue()).decode()
+    return f"data:image/png;base64,{img_base64}"
+
+def upscale_image(image):
+    """
+    Upscale the generated image using Real-ESRGAN (mandatory)
+    """
+    if not image:
+        return None, "No image to upscale"
     
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    # Set up the generator for reproducibility
-    generator = torch.Generator(device=device).manual_seed(seed)
+    if not os.getenv('REPLICATE_API_TOKEN'):
+        return None, "Please set REPLICATE_API_TOKEN"
     
-    print(f"Original Prompt: '{prompt}'")
-    print(f"Negative Prompt: '{negative_prompt}'")
-    print(f"Seed: {seed}, Steps: {num_inference_steps}")
-
-    if rewrite_prompt:
-        prompt = polish_prompt(prompt, image)
-        print(f"Rewritten Prompt: {prompt}")
+    try:
+        # Upload image to hosting
+        image_url = upload_image_to_hosting(image)
+        
+        # Run Real-ESRGAN model
+        output = replicate.run(
+            "nightmareai/real-esrgan:f121d640bd286e1fdc67f9799164c1d5be36ff74576ee11c803ae5b665dd46aa",
+            input={
+                "image": image_url,
+                "scale": 4  # 4x upscaling as default
+            }
+        )
+        
+        if output is None:
+            return None, "No output received from upscaler"
+        
+        # Get the upscaled image
+        try:
+            if hasattr(output, 'read'):
+                img_data = output.read()
+                img = Image.open(BytesIO(img_data))
+                return img, "🔍 Upscaled 4x successfully!"
+        except:
+            pass
+        
+        try:
+            if hasattr(output, 'url'):
+                output_url = output.url()
+                response = requests.get(output_url, timeout=30)
+                if response.status_code == 200:
+                    img = Image.open(BytesIO(response.content))
+                    return img, "🔍 Upscaled 4x successfully!"
+        except:
+            pass
+        
+        output_url = None
+        if isinstance(output, str):
+            output_url = output
+        elif isinstance(output, list) and len(output) > 0:
+            output_url = output[0]
+        
+        if output_url:
+            response = requests.get(output_url, timeout=30)
+            if response.status_code == 200:
+                img = Image.open(BytesIO(response.content))
+                return img, "🔍 Upscaled 4x successfully!"
+        
+        return None, "Could not process upscaled output"
+        
+    except Exception as e:
+        return None, f"Upscale error: {str(e)[:100]}"
 
+def apply_outpainting_to_image(image, outpaint_prompt, target_width, target_height, 
+                               overlap_percentage, resize_option, custom_resize_percentage,
+                               alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
+    """
+    Apply outpainting to an image by preparing it with white margins
+    """
+    if not image:
+        return None
+    
+    # Check if expansion is possible
+    if not can_expand(image.width, image.height, target_width, target_height, alignment):
+        alignment = "Middle"
+    
     # Prepare the image with white margins for outpainting
     outpaint_image, mask = prepare_image_and_mask(
-        image, width, height, overlap_percentage,
+        image, target_width, target_height, overlap_percentage,
         resize_option, custom_resize_percentage, alignment,
         overlap_left, overlap_right, overlap_top, overlap_bottom
     )
     
-    # Check if expansion is possible
-    if not can_expand(image.width, image.height, width, height, alignment):
-        alignment = "Middle"
-        outpaint_image, mask = prepare_image_and_mask(
-            image, width, height, overlap_percentage,
-            resize_option, custom_resize_percentage, "Middle",
-            overlap_left, overlap_right, overlap_top, overlap_bottom
-        )
+    return outpaint_image
+
+def process_images(prompt, image1, image2=None, enable_outpaint=False, outpaint_prompt="",
+                  target_width=1280, target_height=720, overlap_percentage=10,
+                  resize_option="Full", custom_resize_percentage=50,
+                  alignment="Middle", overlap_left=True, overlap_right=True,
+                  overlap_top=True, overlap_bottom=True):
+    """
+    Process uploaded images with Replicate API, apply optional outpainting, and mandatory upscaling
+    """
+    if not image1:
+        return None, "Please upload at least one image"
+    
+    if not os.getenv('REPLICATE_API_TOKEN'):
+        return None, "Please set REPLICATE_API_TOKEN"
     
-    print(f"Outpaint dimensions: {outpaint_image.size}")
+    try:
+        # Step 1: Apply outpainting if enabled
+        if enable_outpaint:
+            # Apply outpainting to image1
+            image1 = apply_outpainting_to_image(
+                image1, outpaint_prompt, target_width, target_height,
+                overlap_percentage, resize_option, custom_resize_percentage,
+                alignment, overlap_left, overlap_right, overlap_top, overlap_bottom
+            )
+            
+            # Apply outpainting to image2 if it exists
+            if image2:
+                image2 = apply_outpainting_to_image(
+                    image2, outpaint_prompt, target_width, target_height,
+                    overlap_percentage, resize_option, custom_resize_percentage,
+                    alignment, overlap_left, overlap_right, overlap_top, overlap_bottom
+                )
+            
+            # Update the prompt if outpainting is enabled
+            if outpaint_prompt:
+                prompt = f"replace the white margins. {outpaint_prompt}. {prompt}"
+        
+        # Step 2: Upload images and process with Nano Banana
+        image_urls = []
+        
+        url1 = upload_image_to_hosting(image1)
+        image_urls.append(url1)
+        
+        if image2:
+            url2 = upload_image_to_hosting(image2)
+            image_urls.append(url2)
+        
+        # Run the Nano Banana model
+        output = replicate.run(
+            "google/nano-banana",
+            input={
+                "prompt": prompt,
+                "image_input": image_urls
+            }
+        )
+        
+        if output is None:
+            return None, "No output received"
+        
+        # Get the generated image
+        generated_image = None
+        try:
+            if hasattr(output, 'read'):
+                img_data = output.read()
+                generated_image = Image.open(BytesIO(img_data))
+        except:
+            pass
+        
+        if not generated_image:
+            try:
+                if hasattr(output, 'url'):
+                    output_url = output.url()
+                    response = requests.get(output_url, timeout=30)
+                    if response.status_code == 200:
+                        generated_image = Image.open(BytesIO(response.content))
+            except:
+                pass
+        
+        if not generated_image:
+            output_url = None
+            if isinstance(output, str):
+                output_url = output
+            elif isinstance(output, list) and len(output) > 0:
+                output_url = output[0]
+            
+            if output_url:
+                response = requests.get(output_url, timeout=30)
+                if response.status_code == 200:
+                    generated_image = Image.open(BytesIO(response.content))
+        
+        if not generated_image:
+            return None, "Could not process output"
+        
+        # Step 3: Apply mandatory upscaling
+        upscaled_image, upscale_status = upscale_image(generated_image)
+        
+        if upscaled_image:
+            return upscaled_image, f"✨ Generated and {upscale_status}"
+        else:
+            # If upscaling fails, return the generated image with a warning
+            return generated_image, "✨ Generated (upscaling failed, returning original)"
+        
+    except Exception as e:
+        return None, f"Error: {str(e)[:100]}"
 
-    # Generate the image with outpainting preprocessing
-    result_image = pipe(
-        outpaint_image,  # Use the preprocessed image with white margins
-        prompt="replace the white margins. "+ prompt,
-        negative_prompt=negative_prompt,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-        true_cfg_scale=true_guidance_scale,
-    ).images[0]
+def toggle_outpaint_options(enable):
+    """Toggle visibility of outpainting options"""
+    return gr.update(visible=enable)
 
-    return result_image, seed
+def preload_presets(target_ratio):
+    """Updates the width and height based on the selected aspect ratio."""
+    if target_ratio == "9:16":
+        return 720, 1280
+    elif target_ratio == "16:9":
+        return 1280, 720
+    elif target_ratio == "1:1":
+        return 1024, 1024
+    else:  # Custom
+        return 1280, 720
 
-# --- Examples and UI Layout ---
-# You can add examples here if you have sample images
-# examples = [
-#     ["path/to/example1.jpg", "extend the landscape", 1280, 720, "Middle"],
-#     ["path/to/example2.jpg", "add more sky", 1024, 1024, "Top"],
-# ]
+def toggle_custom_resize_slider(resize_option):
+    return gr.update(visible=(resize_option == "Custom"))
 
+# Enhanced CSS with modern, minimal design
 css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 1024px;
+.gradio-container {
+    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
+    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
+    min-height: 100vh;
+}
+.header-container {
+    background: linear-gradient(135deg, #ffd93d 0%, #ffb347 100%);
+    padding: 2.5rem;
+    border-radius: 24px;
+    margin-bottom: 2.5rem;
+    box-shadow: 0 20px 60px rgba(255, 179, 71, 0.25);
+}
+.logo-text {
+    font-size: 3.5rem;
+    font-weight: 900;
+    color: #2d3436;
+    text-align: center;
+    margin: 0;
+    letter-spacing: -2px;
 }
-#logo-title {
+.subtitle {
+    color: #2d3436;
     text-align: center;
+    font-size: 1rem;
+    margin-top: 0.5rem;
+    opacity: 0.8;
+}
+.main-content {
+    background: rgba(255, 255, 255, 0.95);
+    backdrop-filter: blur(20px);
+    border-radius: 24px;
+    padding: 2.5rem;
+    box-shadow: 0 10px 40px rgba(0, 0, 0, 0.08);
+}
+.gr-button-primary {
+    background: linear-gradient(135deg, #ffd93d 0%, #ffb347 100%) !important;
+    border: none !important;
+    color: #2d3436 !important;
+    font-weight: 700 !important;
+    font-size: 1.1rem !important;
+    padding: 1.2rem 2rem !important;
+    border-radius: 14px !important;
+    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
+    text-transform: uppercase;
+    letter-spacing: 1px;
+    width: 100%;
+    margin-top: 1rem !important;
+}
+.gr-button-primary:hover {
+    transform: translateY(-3px) !important;
+    box-shadow: 0 15px 40px rgba(255, 179, 71, 0.35) !important;
+}
+.gr-button-secondary {
+    background: linear-gradient(135deg, #74b9ff 0%, #0984e3 100%) !important;
+    border: none !important;
+    color: white !important;
+    font-weight: 600 !important;
+    font-size: 0.95rem !important;
+    padding: 0.8rem 1.5rem !important;
+    border-radius: 12px !important;
+    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
+}
+.gr-button-secondary:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 10px 30px rgba(9, 132, 227, 0.3) !important;
+}
+.gr-input, .gr-textarea {
+    background: #ffffff !important;
+    border: 2px solid #e1e8ed !important;
+    border-radius: 14px !important;
+    color: #2d3436 !important;
+    font-size: 1rem !important;
+    padding: 0.8rem 1rem !important;
 }
-#logo-title img {
-    width: 400px;
+.gr-input:focus, .gr-textarea:focus {
+    border-color: #ffd93d !important;
+    box-shadow: 0 0 0 4px rgba(255, 217, 61, 0.15) !important;
+}
+.gr-form {
+    background: transparent !important;
+    border: none !important;
+}
+.gr-panel {
+    background: #ffffff !important;
+    border: 2px solid #e1e8ed !important;
+    border-radius: 16px !important;
+    padding: 1.5rem !important;
+}
+.gr-box {
+    border-radius: 14px !important;
+    border-color: #e1e8ed !important;
+}
+label {
+    color: #636e72 !important;
+    font-weight: 600 !important;
+    font-size: 0.85rem !important;
+    text-transform: uppercase;
+    letter-spacing: 0.5px;
+    margin-bottom: 0.5rem !important;
+}
+.status-text {
+    font-family: 'SF Mono', 'Monaco', monospace;
+    color: #00b894;
+    font-size: 0.9rem;
+}
+.image-container {
+    border-radius: 14px !important;
+    overflow: hidden;
+    border: 2px solid #e1e8ed !important;
+    background: #fafbfc !important;
 }
-#edit_text{margin-top: -62px !important}
 .preview-container {
-    border: 1px solid #e0e0e0;
-    border-radius: 8px;
-    padding: 10px;
-    margin-top: 10px;
+    border: 2px dashed #ff6b6b !important;
+    border-radius: 14px !important;
+    padding: 1rem !important;
+    background: rgba(255, 107, 107, 0.05) !important;
+}
+footer {
+    display: none !important;
 }
-.gallery-container {
-    margin-top: 20px;
+/* Equal sizing for all image containers */
+.image-upload {
+    min-height: 200px !important;
+    max-height: 200px !important;
+}
+.output-image {
+    min-height: 420px !important;
+    max-height: 420px !important;
+}
+/* Ensure consistent spacing */
+.gr-row {
+    gap: 1rem !important;
+}
+.gr-column {
+    gap: 1rem !important;
+}
+/* Outpainting options styling */
+.outpaint-section {
+    background: rgba(116, 185, 255, 0.1) !important;
+    border: 2px solid #74b9ff !important;
+    border-radius: 14px !important;
+    padding: 1rem !important;
+    margin-top: 1rem !important;
 }
 """
 
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
+with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
+    with gr.Column(elem_classes="header-container"):
         gr.HTML("""
-        <div id="logo-title">
-            <img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/qwen_image_edit_logo.png" alt="Qwen-Image Edit Logo" width="400" style="display: block; margin: 0 auto;">
-            <h2 style="font-style: italic;color: #5b47d1;margin-top: -27px !important;margin-left: 133px;">Outpaint [Fast]</h2>
-        </div>
-        """)
-        gr.Markdown("""
-        
-        Outpaint images with Qwen Image Edit. [Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series. 
-        
-        This demo uses the [Qwen-Image-Lightning](https://huggingface.co/lightx2v/Qwen-Image-Lightning) LoRA with AoT compilation and FA3 for accelerated 8-step inference.
-        Try on [Qwen Chat](https://chat.qwen.ai/), or [download model](https://huggingface.co/Qwen/Qwen-Image-Edit) to run locally with ComfyUI or diffusers.
-        """)
-        
-        
-        with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(label="Input Image", type="pil")
-                
-                prompt = gr.Text(
-                    label="Prompt",
-                    info="Describe what should appear in the extended areas",
-                    value="extend the image naturally", 
+            <h1 class="logo-text">🍌 Nano Banana Upscale(X4)</h1>
+            <p class="subtitle">AI-Powered Image Style Transfer with Outpainting & Auto-Upscaling</p>
+                <div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-top: 20px;">
+                    <a href="https://huggingface.co/spaces/openfree/Nano-Banana-Upscale" target="_blank">
+                        <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=UPSCALE&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Nano Banana Upscale">
+                    </a>
+                    <a href="https://huggingface.co/spaces/openfree/Free-Nano-Banana" target="_blank">
+                        <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=FREE&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Free Nano Banana">
+                    </a>
+                    <a href="https://huggingface.co/spaces/aiqtech/Nano-Banana-API" target="_blank">
+                        <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=API&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Nano Banana API">
+                    </a>
+                    <a href="https://huggingface.co/spaces/ginigen/Nano-Banana-Video" target="_blank">
+                        <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=VIDEO&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Nano Banana VIDEO">
+                    </a>                    
+                    <a href="https://discord.gg/openfreeai" target="_blank">
+                        <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="Discord Openfree AI">
+                    </a>
+                </div>
+            """)
+    
+    with gr.Column(elem_classes="main-content"):
+        with gr.Row(equal_height=True):
+            # Left Column - Inputs
+            with gr.Column(scale=1):
+                prompt = gr.Textbox(
+                    label="Style Description",
+                    placeholder="Describe your style...",
+                    lines=3,
+                    value="Make the sheets in the style of the logo. Make the scene natural.",
+                    elem_classes="prompt-input"
                 )
                 
-                with gr.Row():
-                    target_ratio = gr.Radio(
-                        label="Target Ratio",
-                        choices=["9:16", "16:9", "1:1", "Custom"],
-                        value="16:9",
-                        scale=2
+                with gr.Row(equal_height=True):
+                    image1 = gr.Image(
+                        label="Primary Image",
+                        type="pil",
+                        height=200,
+                        elem_classes="image-container image-upload"
                     )
-                    alignment_dropdown = gr.Dropdown(
-                        choices=["Middle", "Left", "Right", "Top", "Bottom"],
-                        value="Middle",
-                        label="Alignment"
+                    image2 = gr.Image(
+                        label="Secondary Image (Optional)",
+                        type="pil",
+                        height=200,
+                        elem_classes="image-container image-upload"
                     )
                 
-                run_button = gr.Button("run", variant="primary")
+                # Outpainting Options
+                enable_outpaint = gr.Checkbox(
+                    label="🎨 Enable Outpainting (Expand Image)",
+                    value=False
+                )
                 
-                with gr.Accordion("Outpainting Settings", open=False) as settings_panel:
+                with gr.Column(visible=False, elem_classes="outpaint-section") as outpaint_options:
+                    outpaint_prompt = gr.Textbox(
+                        label="Outpaint Prompt",
+                        placeholder="Describe what should appear in the extended areas",
+                        value="extend the image naturally",
+                        lines=2
+                    )
+                    
                     with gr.Row():
-                        width_slider = gr.Slider(
+                        target_ratio = gr.Radio(
+                            label="Target Ratio",
+                            choices=["9:16", "16:9", "1:1", "Custom"],
+                            value="16:9"
+                        )
+                        alignment_dropdown = gr.Dropdown(
+                            choices=["Middle", "Left", "Right", "Top", "Bottom"],
+                            value="Middle",
+                            label="Alignment"
+                        )
+                    
+                    with gr.Row():
+                        target_width = gr.Slider(
                             label="Target Width",
                             minimum=512,
                             maximum=2048,
                             step=8,
-                            value=1280,
+                            value=1280
                         )
-                        height_slider = gr.Slider(
+                        target_height = gr.Slider(
                             label="Target Height",
                             minimum=512,
                             maximum=2048,
                             step=8,
-                            value=720,
+                            value=720
                         )
                     
-                    with gr.Group():
+                    with gr.Accordion("Advanced Outpaint Settings", open=False):
                         overlap_percentage = gr.Slider(
                             label="Mask overlap (%)",
                             minimum=1,
                             maximum=50,
                             value=10,
                             step=1,
-                            info="Controls the blending area between original and new content"
+                            info="Controls the blending area"
                         )
                         
                         with gr.Row():
@@ -515,228 +650,116 @@ with gr.Blocks(css=css) as demo:
                         with gr.Row():
                             overlap_left = gr.Checkbox(label="Overlap Left", value=True)
                             overlap_bottom = gr.Checkbox(label="Overlap Bottom", value=True)
+                        
+                        with gr.Row():
+                            resize_option = gr.Radio(
+                                label="Resize input image",
+                                choices=["Full", "50%", "33%", "25%", "Custom"],
+                                value="Full"
+                            )
+                            custom_resize_percentage = gr.Slider(
+                                label="Custom resize (%)",
+                                minimum=1,
+                                maximum=100,
+                                step=1,
+                                value=50,
+                                visible=False
+                            )
                     
-                    with gr.Row():
-                        resize_option = gr.Radio(
-                            label="Resize input image",
-                            choices=["Full", "50%", "33%", "25%", "Custom"],
-                            value="Full",
-                            info="How much of the target canvas the original image should occupy"
-                        )
-                        custom_resize_percentage = gr.Slider(
-                            label="Custom resize (%)",
-                            minimum=1,
-                            maximum=100,
-                            step=1,
-                            value=50,
-                            visible=False
-                        )
-                    
-                    preview_button = gr.Button("👁️ Preview alignment and mask", variant="secondary")
-
-                with gr.Accordion("Advanced Settings", open=False):
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=0,
+                    preview_outpaint_btn = gr.Button(
+                        "👁️ Preview Outpaint Mask",
+                        variant="secondary"
                     )
-
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-                    with gr.Row():
-                        true_guidance_scale = gr.Slider(
-                            label="True guidance scale",
-                            minimum=1.0,
-                            maximum=10.0,
-                            step=0.1,
-                            value=1.0
-                        )
-
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=1,
-                            maximum=28,
-                            step=1,
-                            value=8,
-                        )
-
-                    rewrite_prompt = gr.Checkbox(
-                        label="Enhance prompt (using HF Inference)", 
-                        value=True
-                    )
-
-            with gr.Column():
-                result = gr.Image(label="Result", type="pil", interactive=False)
                 
-                use_as_input_button = gr.Button("🔄 Use as Input Image", visible=False, variant="secondary")
+                generate_btn = gr.Button(
+                    "Generate Magic with Auto-Upscale ✨",
+                    variant="primary",
+                    size="lg"
+                )
+            
+            # Right Column - Output
+            with gr.Column(scale=1):
+                output_image = gr.Image(
+                    label="Generated & Upscaled Result",
+                    type="pil",
+                    height=420,
+                    elem_classes="image-container output-image"
+                )
                 
-                with gr.Column(visible=False) as preview_container:
-                    preview_image = gr.Image(label="Preview (red area will be generated)", type="pil")
+                status = gr.Textbox(
+                    label="Status",
+                    interactive=False,
+                    lines=1,
+                    elem_classes="status-text",
+                    value="Ready to generate..."
+                )
                 
-                gr.Markdown("---")
+                outpaint_preview = gr.Image(
+                    label="Outpaint Preview (red area will be generated)",
+                    type="pil",
+                    visible=False,
+                    elem_classes="preview-container"
+                )
                 
-                with gr.Row():
-                    gr.Markdown("### 📜 History")
-                    clear_history_button = gr.Button("🗑️ Clear History", size="sm", variant="stop")
+                gr.Markdown("""
+                ### 📌 Features:
+                - **Style Transfer**: Apply artistic styles to your images
+                - **Outpainting** (Optional): Expand your images with AI
+                - **Auto 4x Upscaling**: All outputs are automatically upscaled for maximum quality
                 
-                history_gallery = gr.Gallery(
-                    label="Click any image to use as input", 
-                    columns=4, 
-                    rows=2,
-                    object_fit="contain", 
-                    height="auto",
-                    interactive=False,
-                    show_label=True,
-                    elem_classes=["gallery-container"]
-                )
-
-    # Event handlers
-    use_as_input_button.click(
-        fn=use_output_as_input,
-        inputs=[result],
-        outputs=[input_image],
-        show_api=False
-    )
+                ### 💡 Tips:
+                - Upload 1-2 images to apply style transfer
+                - Enable outpainting to expand image boundaries
+                - All generated images are automatically upscaled 4x
+                """)
     
-    history_gallery.select(
-        fn=use_history_as_input,
-        inputs=[history_gallery],
-        outputs=[input_image],
-        show_api=False
-    )
-    
-    clear_history_button.click(
-        fn=lambda: [],
-        inputs=None,
-        outputs=history_gallery,
-        show_api=False
+    # Event handlers
+    enable_outpaint.change(
+        fn=toggle_outpaint_options,
+        inputs=[enable_outpaint],
+        outputs=[outpaint_options]
     )
     
     target_ratio.change(
         fn=preload_presets,
-        inputs=[target_ratio, width_slider, height_slider],
-        outputs=[width_slider, height_slider, settings_panel],
-        queue=False,
-    )
-    
-    width_slider.change(
-        fn=select_the_right_preset,
-        inputs=[width_slider, height_slider],
-        outputs=[target_ratio],
-        queue=False,
-    )
-    
-    height_slider.change(
-        fn=select_the_right_preset,
-        inputs=[width_slider, height_slider],
-        outputs=[target_ratio],
-        queue=False,
+        inputs=[target_ratio],
+        outputs=[target_width, target_height]
     )
     
     resize_option.change(
         fn=toggle_custom_resize_slider,
         inputs=[resize_option],
-        outputs=[custom_resize_percentage],
-        queue=False,
+        outputs=[custom_resize_percentage]
     )
     
-    preview_button.click(
-        fn=lambda: gr.update(visible=True),
-        inputs=None,
-        outputs=[preview_container],
-        queue=False,
-    ).then(
-        fn=preview_image_and_mask,
+    preview_outpaint_btn.click(
+        fn=preview_outpaint,
         inputs=[
-            input_image, width_slider, height_slider, overlap_percentage,
+            image1, target_width, target_height, overlap_percentage,
             resize_option, custom_resize_percentage, alignment_dropdown,
             overlap_left, overlap_right, overlap_top, overlap_bottom
         ],
-        outputs=preview_image,
-        queue=False,
-    )
-
-    # Main generation pipeline with result clearing, history update, and button visibility
-    run_button.click(
-        fn=clear_result,
-        inputs=None,
-        outputs=result,
-        show_api=False
-    ).then(
-        fn=infer,
-        inputs=[
-            input_image,
-            prompt,
-            width_slider,
-            height_slider,
-            overlap_percentage,
-            resize_option,
-            custom_resize_percentage,
-            alignment_dropdown,
-            overlap_left,
-            overlap_right,
-            overlap_top,
-            overlap_bottom,
-            seed,
-            randomize_seed,
-            true_guidance_scale,
-            num_inference_steps,
-            rewrite_prompt,
-        ],
-        outputs=[result, seed],
+        outputs=[outpaint_preview]
     ).then(
         fn=lambda: gr.update(visible=True),
-        inputs=None,
-        outputs=use_as_input_button,
-        show_api=False
-    ).then(
-        fn=update_history,
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-        show_api=False
+        outputs=[outpaint_preview]
     )
     
-    # Also trigger on prompt submit
-    prompt.submit(
-        fn=clear_result,
-        inputs=None,
-        outputs=result,
-        show_api=False
-    ).then(
-        fn=infer,
+    generate_btn.click(
+        fn=process_images,
         inputs=[
-            input_image,
-            prompt,
-            width_slider,
-            height_slider,
-            overlap_percentage,
-            resize_option,
-            custom_resize_percentage,
-            alignment_dropdown,
-            overlap_left,
-            overlap_right,
-            overlap_top,
-            overlap_bottom,
-            seed,
-            randomize_seed,
-            true_guidance_scale,
-            num_inference_steps,
-            rewrite_prompt,
+            prompt, image1, image2, enable_outpaint, outpaint_prompt,
+            target_width, target_height, overlap_percentage,
+            resize_option, custom_resize_percentage, alignment_dropdown,
+            overlap_left, overlap_right, overlap_top, overlap_bottom
         ],
-        outputs=[result, seed],
-    ).then(
-        fn=lambda: gr.update(visible=True),
-        inputs=None,
-        outputs=use_as_input_button,
-        show_api=False
-    ).then(
-        fn=update_history,
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-        show_api=False
+        outputs=[output_image, status]
     )
 
+# Launch
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        share=True,
+        server_name="0.0.0.0",
+        server_port=7860
+    )