Update app.py

app.py

@@ -1,7 +1,836 @@
+import logging
+import os
+import random
+import time
+import traceback
+from io import BytesIO
+import io
+import base64
+from openai import OpenAI
+import uuid
+import requests
 import gradio as gr
+import requests
+from PIL import Image, PngImagePlugin
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# API Configuration
+API_TOKEN = os.environ.get("HIDREAM_API_TOKEN")
+API_REQUEST_URL = os.environ.get("API_REQUEST_URL")
+API_RESULT_URL = os.environ.get("API_RESULT_URL")
+API_IMAGE_URL = os.environ.get("API_IMAGE_URL")
+API_VERSION = os.environ.get("API_VERSION")
+API_MODEL_NAME = os.environ.get("API_MODEL_NAME")
+OSS_IMAGE_BUCKET = os.environ.get("OSS_IMAGE_BUCKET")
+OSS_MEDIA_BUCKET = os.environ.get("OSS_MEDIA_BUCKET")
+OSS_TOKEN_URL = os.environ.get("OSS_TOKEN_URL")
+MAX_RETRY_COUNT = int(os.environ.get("MAX_RETRY_COUNT", "3"))
+POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", "1"))
+MAX_POLL_TIME = int(os.environ.get("MAX_POLL_TIME", "300"))
+
+def get_oss_token(is_image=True, prefix="p_"):
+    head = {
+        "Cookie": os.environ.get("OSS_AUTH_COOKIE", "")
+    }
+    if is_image:
+        filename = f"p_{uuid.uuid4()}" if prefix == "p_" else f"j_{uuid.uuid4()}"
+        bucket = OSS_IMAGE_BUCKET
+    else:
+        filename = f"{uuid.uuid4()}.mp4"
+        bucket = OSS_MEDIA_BUCKET
+    token_url = f"{OSS_TOKEN_URL}{bucket}?filename={filename}"
+    req = requests.get(token_url, headers=head)
+    if req.status_code == 200 and req.json()["code"] == 0:
+        return req.json()["result"], filename
+    else:
+        print(req.status_code, req.text)
+        return None, None
+
+def upload_to_gcs(signed_url: str, file_io, is_image=True):
+    if is_image:
+        headers = {
+            "Content-Type": "image/png",  # ensure content-type matches the signed url
+        }
+    else:
+        headers = {
+            "Content-Type": "video/mp4",  # ensure content-type matches the signed url
+        }
+    # with open(file_path, "rb") as f:
+    #     response = requests.put(signed_url, data=f, headers=headers)
+    response = requests.put(signed_url, data=file_io, headers=headers)
+    if response.status_code == 200:
+        print("✅ Upload success")
+    else:
+        print(f"❌ Upload failed, status code: {response.status_code}, response content: {response.text}")
+
+# Instruction refinement prompt
+INSTRUCTION_PROMPT = """Your Role: You are an analytical assistant. Your task is to process a source image and a corresponding editing instruction, assuming the instruction accurately describes a desired transformation. You will 1) describe the source image, 2) output the editing instruction (potentially refined for clarity based on the source image context), and 3) describe the *imagined* result of applying that instruction.
+
+Input:
+1. Source Image: The original 'before' image.
+2. Source Instruction: A text instruction describing the edit to be performed on the Source Image. You *must assume* this instruction is accurate and feasible for the purpose of this task.
+
+Task Breakdown:
+1.  **Describe Source Image:** Generate a description (e.g., key subject, setting) of the Source Image by analyzing it. This will be the first line of your output.
+
+2.  **Output Editing Instruction:** This step determines the second line of your output.
+    * **Assumption:** The provided Source Instruction *accurately* describes the desired edit.
+    * **Goal:** Output a concise, single-line instruction based on the Source Instruction.
+    * **Refinement based on Source Image:** While the Source Instruction is assumed correct, analyze the Source Image to see if the instruction needs refinement for specificity. If the Source Image contains multiple similar objects and the Source Instruction is potentially ambiguous (e.g., "change the car color" when there are three cars), refine the instruction to be specific, using positional qualifiers (e.g., 'the left car', 'the bird on the top branch'), size ('the smaller dog', 'the largest building'), or other distinguishing visual features apparent in the Source Image. If the Source Instruction is already specific or if there's no ambiguity in the Source Image context, you can use it directly or with minor phrasing adjustments for naturalness. The *core meaning* of the Source Instruction must be preserved.
+    * **Output:** Present the resulting specific, single-line instruction as the second line.
+
+3.  **Describe Imagined Target Image:** Based *only* on the Source Image description (Line 1) and the Editing Instruction (Line 2), generate a description of the *imagined outcome*.
+    * Describe the scene from Line 1 *as if* the instruction from Line 2 has been successfully applied. Conceptualize the result of the edit on the source description.
+    * This description must be purely a logical prediction based on applying the instruction (Line 2) to the description in Line 1. Do *not* invent details not implied by the instruction or observed in the source image beyond the specified edit. This will be the third line of your output.
+
+Output Format:
+* Your response *must* consist of exactly three lines.
+* Do not include any other explanations, comments, introductory phrases, labels (like "Line 1:"), or formatting.
+* Your output should be in English.
+
+Line 1: [Description of the Source Image]
+Line 2: [The specific, single-line editing instruction based on the Source Instruction and Source Image context]
+Line 3: [Description of the Imagined Target Image based on Lines 1 & 2]
+
+Now, please generate the three-line output based on the Source Image and the Source Instruction: {source_instruction}
+"""
+
+def filter_response(src_instruction):
+    try:
+        src_instruction = src_instruction.strip().split("\n")
+        src_instruction = [k.strip() for k in src_instruction if k.strip()]
+        src_instruction = [k for k in src_instruction if len(k) > 0]
+        if len(src_instruction) != 3:
+            return ""
+        instruction = src_instruction[1]
+        target_description = src_instruction[2]
+        instruction = instruction.strip().strip(".")
+        inst_format = "Editing Instruction: {}. Target Image Description: {}"
+        return inst_format.format(instruction, target_description)
+    except:
+        return ""
+
+def refine_instruction(src_image, src_instruction):
+    MAX_TOKENS_RESPONSE = 500 # Limit response tokens as output format is structured
+    client = OpenAI()
+    src_image = src_image.convert("RGB")
+    src_image_buffer = io.BytesIO()
+    src_image.save(src_image_buffer, format="JPEG")
+    src_image_buffer.seek(0)
+    src_base64 = base64.b64encode(src_image_buffer.read()).decode('utf-8')
+    encoded_str = f"data:image/jpeg;base64,{src_base64}"
+    image_content = [
+        {"type": "image_url", "image_url": {"url": encoded_str,}},
+    ]
+    instruction_text = INSTRUCTION_PROMPT.format(source_instruction=src_instruction)
+    message_content = [
+        {"type": "text", "text": instruction_text},
+        *image_content # Unpack the list of image dictionaries
+    ]
+    completion = client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {"role": "system", "content": "You are a professional digital artist."},
+                {"role": "user", "content": message_content}
+            ],
+            max_tokens=MAX_TOKENS_RESPONSE, # Good practice to set max tokens
+            temperature=0.2 # Lower temperature for more deterministic output
+        )
+    evaluation_result = completion.choices[0].message.content
+    refined_instruction = filter_response(evaluation_result)
+    return refined_instruction
+
+# Resolution options
+ASPECT_RATIO_OPTIONS = ["1:1", "3:4", "4:3", "9:16", "16:9"]
+
+# Log configuration details
+logger.info(f"API configuration loaded: REQUEST_URL={API_REQUEST_URL}, RESULT_URL={API_RESULT_URL}, VERSION={API_VERSION}, MODEL={API_MODEL_NAME}")
+logger.info(f"OSS configuration: IMAGE_BUCKET={OSS_IMAGE_BUCKET}, MEDIA_BUCKET={OSS_MEDIA_BUCKET}, TOKEN_URL={OSS_TOKEN_URL}")
+logger.info(f"Retry configuration: MAX_RETRY_COUNT={MAX_RETRY_COUNT}, POLL_INTERVAL={POLL_INTERVAL}s, MAX_POLL_TIME={MAX_POLL_TIME}s")
+
+
+class APIError(Exception):
+    """Custom exception for API-related errors"""
+    pass
+
+
+def create_request(prompt, image, guidance_scale=5.0, image_guidance_scale=4.0, seed=-1):
+    """
+    Create an image editing request to the API.
+    
+    Args:
+        prompt (str): Text prompt describing the image edit
+        image (PIL.Image): Input image to edit
+        guidance_scale (float): Strength of instruction following
+        image_guidance_scale (float): Strength of image preservation
+        seed (int): Seed for reproducibility, -1 for random
+        
+    Returns:
+        tuple: (task_id, seed) - Task ID if successful and the seed used
+        
+    Raises:
+        APIError: If the API request fails
+    """
+    logger.info(f"Starting create_request with prompt='{prompt[:50]}...', guidance_scale={guidance_scale}, image_guidance_scale={image_guidance_scale}, seed={seed}")
+    image_io = io.BytesIO()
+    image = image.convert("RGB")
+    image.save(image_io, format="PNG")
+    image_io.seek(0)
+    token_url, filename = get_oss_token(is_image=True)
+    upload_to_gcs(token_url, image_io, is_image=True)
+
+    if not prompt or not prompt.strip():
+        logger.error("Empty prompt provided to create_request")
+        raise ValueError("Prompt cannot be empty")
+
+    if not image:
+        logger.error("No image provided to create_request")
+        raise ValueError("Image cannot be empty")
+
+    # Generate random seed if not provided
+    if seed == -1:
+        seed = random.randint(1, 1000000)
+        logger.info(f"Generated random seed: {seed}")
+
+    # Validate seed
+    try:
+        seed = int(seed)
+        if seed < -1 or seed > 1000000:
+            logger.info(f"Invalid seed value: {seed}, forcing to 8888")
+            seed = 8888
+    except (TypeError, ValueError) as e:
+        logger.error(f"Seed validation failed: {str(e)}")
+        raise ValueError(f"Seed must be an integer but got {seed}")
+
+    headers = {
+        "Authorization": f"Bearer {API_TOKEN}",
+        "X-accept-language": "en",
+        "X-source": "api",
+        "Content-Type": "application/json",
+    }
+
+    generate_data = {
+        "module": "image_edit",
+        "images": [filename, ],
+        "prompt": prompt,
+        "params": {
+            "seed": seed,
+            "custom_params": {
+                "sample_steps": 28,
+                "guidance_scale": guidance_scale,
+                "image_guidance_scale": image_guidance_scale
+            },
+        },
+        "version": API_VERSION,
+    }
+
+    retry_count = 0
+    while retry_count < MAX_RETRY_COUNT:
+        try:
+            logger.info(f"Sending API request [attempt {retry_count+1}/{MAX_RETRY_COUNT}] for prompt: '{prompt[:50]}...'")
+            response = requests.post(API_REQUEST_URL, json=generate_data, headers=headers, timeout=10)
+            
+            # Log response status code
+            logger.info(f"API request response status: {response.status_code}")
+            
+            response.raise_for_status()
+
+            result = response.json()
+            if not result or "result" not in result:
+                logger.error(f"Invalid API response format: {str(result)}")
+                raise APIError(f"Invalid response format from API when sending request: {str(result)}")
+
+            task_id = result.get("result", {}).get("task_id")
+            if not task_id:
+                logger.error(f"No task ID in API response: {str(result)}")
+                raise APIError(f"No task ID returned from API: {str(result)}")
+
+            logger.info(f"Successfully created task with ID: {task_id}, seed: {seed}")
+            return task_id, seed
+
+        except requests.exceptions.Timeout:
+            retry_count += 1
+            logger.warning(f"Request timed out. Retrying ({retry_count}/{MAX_RETRY_COUNT})...")
+            time.sleep(1)
+
+        except requests.exceptions.HTTPError as e:
+            status_code = e.response.status_code
+            error_message = f"HTTP error {status_code}"
+            
+            try:
+                error_detail = e.response.json()
+                error_message += f": {error_detail}"
+                logger.error(f"API response error content: {error_detail}")
+            except:
+                logger.error(f"Could not parse API error response as JSON. Raw content: {e.response.content[:500]}")
+
+            if status_code == 401:
+                logger.error(f"Authentication failed with API token. Status code: {status_code}")
+                raise APIError("Authentication failed. Please check your API token.")
+            elif status_code == 429:
+                retry_count += 1
+                wait_time = min(2 ** retry_count, 10)  # Exponential backoff
+                logger.warning(f"Rate limit exceeded. Waiting {wait_time}s before retry ({retry_count}/{MAX_RETRY_COUNT})...")
+                time.sleep(wait_time)
+            elif 400 <= status_code < 500:
+                try:
+                    error_detail = e.response.json()
+                    error_message += f": {error_detail.get('message', 'Client error')}"
+                except:
+                    pass
+                logger.error(f"Client error: {error_message}, Prompt: '{prompt[:50]}...', Status: {status_code}")
+                raise APIError(error_message)
+            else:
+                retry_count += 1
+                logger.warning(f"Server error: {error_message}. Retrying ({retry_count}/{MAX_RETRY_COUNT})...")
+                time.sleep(1)
+
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Request error: {str(e)}")
+            logger.debug(f"Request error details: {traceback.format_exc()}")
+            raise APIError(f"Failed to connect to API: {str(e)}")
+
+        except Exception as e:
+            logger.error(f"Unexpected error in create_request: {str(e)}")
+            logger.error(f"Full traceback: {traceback.format_exc()}")
+            raise APIError(f"Unexpected error: {str(e)}")
+
+    logger.error(f"Failed to create request after {MAX_RETRY_COUNT} retries for prompt: '{prompt[:50]}...'")
+    raise APIError(f"Failed after {MAX_RETRY_COUNT} retries")
+
+
+def get_results(task_id):
+    """
+    Check the status of an image generation task.
+    
+    Args:
+        task_id (str): The task ID to check
+        
+    Returns:
+        dict: Task result information
+        
+    Raises:
+        APIError: If the API request fails
+    """
+    logger.debug(f"Checking status for task ID: {task_id}")
+    
+    if not task_id:
+        logger.error("Empty task ID provided to get_results")
+        raise ValueError("Task ID cannot be empty")
+
+    url = f"{API_RESULT_URL}?task_id={task_id}"
+    headers = {
+        "Authorization": f"Bearer {API_TOKEN}",
+        "X-accept-language": "en",
+    }
+
+    try:
+        response = requests.get(url, headers=headers, timeout=10)
+        logger.debug(f"Status check response code: {response.status_code}")
+        
+        response.raise_for_status()
+        result = response.json()
+
+        if not result or "result" not in result:
+            logger.warning(f"Invalid response format from API when checking task {task_id}: {str(result)}")
+            raise APIError(f"Invalid response format from API when checking task {task_id}: {str(result)}")
+
+        return result
+
+    except requests.exceptions.Timeout:
+        logger.warning(f"Request timed out when checking task {task_id}")
+        return None
+
+    except requests.exceptions.HTTPError as e:
+        status_code = e.response.status_code
+        logger.warning(f"HTTP error {status_code} when checking task {task_id}")
+        
+        try:
+            error_content = e.response.json()
+            logger.error(f"Error response content: {error_content}")
+        except:
+            logger.error(f"Could not parse error response as JSON. Raw content: {e.response.content[:500]}")
+            
+        if status_code == 401:
+            logger.error(f"Authentication failed when checking task {task_id}")
+            raise APIError(f"Authentication failed. Please check your API token when checking task {task_id}")
+        elif 400 <= status_code < 500:
+            try:
+                error_detail = e.response.json()
+                error_message = f"HTTP error {status_code}: {error_detail.get('message', 'Client error')}"
+            except:
+                error_message = f"HTTP error {status_code}"
+            logger.error(error_message)
+            return None
+        else:
+            logger.warning(f"Server error {status_code} when checking task {task_id}")
+            return None
+
+    except requests.exceptions.RequestException as e:
+        logger.warning(f"Network error when checking task {task_id}: {str(e)}")
+        logger.debug(f"Network error details: {traceback.format_exc()}")
+        return None
+
+    except Exception as e:
+        logger.error(f"Unexpected error when checking task {task_id}: {str(e)}")
+        logger.error(f"Full traceback: {traceback.format_exc()}")
+        return None
+
+
+def download_image(image_url):
+    """
+    Download an image from a URL and return it as a PIL Image.
+    Converts WebP to PNG format while preserving original image data.
+    
+    Args:
+        image_url (str): URL of the image
+        
+    Returns:
+        PIL.Image: Downloaded image object converted to PNG format
+        
+    Raises:
+        APIError: If the download fails
+    """
+    logger.info(f"Starting download_image from URL: {image_url}")
+    
+    if not image_url:
+        logger.error("Empty image URL provided to download_image")
+        raise ValueError("Image URL cannot be empty when downloading image")
+
+    retry_count = 0
+    while retry_count < MAX_RETRY_COUNT:
+        try:
+            logger.info(f"Downloading image [attempt {retry_count+1}/{MAX_RETRY_COUNT}] from {image_url}")
+            response = requests.get(image_url, timeout=15)
+            
+            logger.debug(f"Image download response status: {response.status_code}, Content-Type: {response.headers.get('Content-Type')}, Content-Length: {response.headers.get('Content-Length')}")
+            
+            response.raise_for_status()
+
+            # Open the image from response content
+            image = Image.open(BytesIO(response.content))
+            logger.info(f"Image opened successfully. Format: {image.format}, Size: {image.size[0]}x{image.size[1]}, Mode: {image.mode}")
+
+            # Get original metadata before conversion
+            original_metadata = {}
+            for key, value in image.info.items():
+                if isinstance(key, str) and isinstance(value, str):
+                    original_metadata[key] = value
+            
+            logger.debug(f"Original image metadata: {original_metadata}")
+
+            # Convert to PNG regardless of original format (WebP, JPEG, etc.)
+            if image.format != 'PNG':
+                logger.info(f"Converting image from {image.format} to PNG format")
+                png_buffer = BytesIO()
+
+                # If the image has an alpha channel, preserve it, otherwise convert to RGB
+                if 'A' in image.getbands():
+                    logger.debug("Preserving alpha channel in image conversion")
+                    image_to_save = image
+                else:
+                    logger.debug("Converting image to RGB mode")
+                    image_to_save = image.convert('RGB')
+
+                image_to_save.save(png_buffer, format='PNG')
+                png_buffer.seek(0)
+                image = Image.open(png_buffer)
+                logger.debug(f"Image converted to PNG. New size: {image.size[0]}x{image.size[1]}, Mode: {image.mode}")
+
+                # Preserve original metadata
+                for key, value in original_metadata.items():
+                    image.info[key] = value
+                logger.debug("Original metadata preserved in converted image")
+
+            logger.info(f"Successfully downloaded and processed image: {image.size[0]}x{image.size[1]}")
+            return image
+
+        except requests.exceptions.Timeout:
+            retry_count += 1
+            logger.warning(f"Download timed out. Retrying ({retry_count}/{MAX_RETRY_COUNT})...")
+            time.sleep(1)
+
+        except requests.exceptions.HTTPError as e:
+            status_code = e.response.status_code
+            logger.error(f"HTTP error {status_code} when downloading image from {image_url}")
+            
+            try:
+                error_content = e.response.text[:500]
+                logger.error(f"Error response content: {error_content}")
+            except:
+                logger.error("Could not read error response content")
+                
+            if 400 <= status_code < 500:
+                error_message = f"HTTP error {status_code} when downloading image"
+                logger.error(error_message)
+                raise APIError(error_message)
+            else:
+                retry_count += 1
+                logger.warning(f"Server error {status_code}. Retrying ({retry_count}/{MAX_RETRY_COUNT})...")
+                time.sleep(1)
+
+        except requests.exceptions.RequestException as e:
+            retry_count += 1
+            logger.warning(f"Network error during image download: {str(e)}. Retrying ({retry_count}/{MAX_RETRY_COUNT})...")
+            logger.debug(f"Network error details: {traceback.format_exc()}")
+            time.sleep(1)
+
+        except Exception as e:
+            logger.error(f"Error processing image from {image_url}: {str(e)}")
+            logger.error(f"Full traceback: {traceback.format_exc()}")
+            raise APIError(f"Failed to process image: {str(e)}")
+
+    logger.error(f"Failed to download image from {image_url} after {MAX_RETRY_COUNT} retries")
+    raise APIError(f"Failed to download image after {MAX_RETRY_COUNT} retries")
+
+
+def add_metadata_to_image(image, metadata):
+    """
+    Add metadata to a PIL image.
+    
+    Args:
+        image (PIL.Image): The image to add metadata to
+        metadata (dict): Metadata to add to the image
+        
+    Returns:
+        PIL.Image: Image with metadata
+    """
+    logger.debug(f"Adding metadata to image: {metadata}")
+    
+    if not image:
+        logger.error("Null image provided to add_metadata_to_image")
+        return None
+    try:
+        # Get any existing metadata
+        existing_metadata = {}
+        for key, value in image.info.items():
+            if isinstance(key, str) and isinstance(value, str):
+                existing_metadata[key] = value
+        
+        logger.debug(f"Existing image metadata: {existing_metadata}")
+
+        # Merge with new metadata (new values override existing ones)
+        all_metadata = {**existing_metadata, **metadata}
+        logger.debug(f"Combined metadata: {all_metadata}")
+
+        # Create a new metadata dictionary for PNG
+        meta = PngImagePlugin.PngInfo()
+
+        # Add each metadata item
+        for key, value in all_metadata.items():
+            meta.add_text(key, str(value))
+
+        # Save with metadata to a buffer
+        buffer = BytesIO()
+        image.save(buffer, format='PNG', pnginfo=meta)
+        logger.debug("Image saved to buffer with metadata")
+
+        # Reload the image from the buffer
+        buffer.seek(0)
+        result_image = Image.open(buffer)
+        logger.debug("Image reloaded from buffer with metadata")
+        return result_image
+
+    except Exception as e:
+        logger.error(f"Failed to add metadata to image: {str(e)}")
+        logger.error(f"Full traceback: {traceback.format_exc()}")
+        return image  # Return original image if metadata addition fails
+
+# Create Gradio interface
+def create_ui():
+    logger.info("Creating Gradio UI")
+    with gr.Blocks(title="HiDream-E1-Full Image Editor", theme=gr.themes.Base()) as demo:
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=1):
+                gr.Markdown("""
+                # HiDream-E1-Full Image Editor
+                
+                Edit images using natural language instructions with state-of-the-art AI  [🤗 HuggingFace](https://huggingface.co/HiDream-ai/HiDream-E1-Full) | [GitHub](https://github.com/HiDream-ai/HiDream-E1) | [Twitter](https://x.com/vivago_ai)
+                
+                <span style="color: #FF5733; font-weight: bold">For more features and to experience the full capabilities of our product, please visit [https://vivago.ai/](https://vivago.ai/).</span>
+                """)
+
+        with gr.Row():
+            # Input column
+            with gr.Column(scale=1):
+                input_image = gr.Image(
+                    type="pil", 
+                    label="Input Image", 
+                    height=400,
+                    show_download_button=True,
+                    show_label=True,
+                    scale=1,
+                    container=True,
+                    image_mode="RGB"
+                )
+                
+                instruction = gr.Textbox(
+                    label="Editing Instruction",
+                    placeholder="e.g., convert the image into a Ghibli style",
+                    lines=3
+                )
+                
+                gr.Markdown("""
+                <div style="padding: 8px; margin-bottom: 10px; background-color: #E2F0FF; border-left: 5px solid #2E86DE; color: #2C3E50;">
+                    <strong>Note:</strong> For optimal results, we recommend using the <strong>Refine Instruction</strong> button which formats your input into:
+                    <br><em>"Editing Instruction: [your instruction]. Target Image Description: [expected result]"</em>
+                </div>
+                """)
+                
+                with gr.Row():
+                    refine_btn = gr.Button("Refine Instruction")
+                    generate_btn = gr.Button("Generate", variant="primary", size="lg")
+                
+                with gr.Accordion("Advanced Settings", open=True):
+                    gr.Markdown("""
+                    <div style="padding: 8px; margin: 15px 0; background-color: #FFF3CD; border-left: 5px solid #FFDD57; color: #856404;">
+                        <strong>Important:</strong> Adjust these parameters based on your editing needs:
+                        <ul>
+                            <li>For style changes, use higher image preservation strength (e.g., 3.0-4.0)</li>
+                            <li>For local edits like adding, deleting, replacing elements, use lower image preservation strength (e.g., 2.0-3.0)</li>
+                        </ul>
+                    </div>
+                    """)
+                    with gr.Row():
+                        guidance_scale = gr.Slider(
+                            minimum=1.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=5.0,
+                            label="Instruction Following Strength"
+                        )
+                        image_guidance_scale = gr.Slider(
+                            minimum=1.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=3.0,
+                            label="Image Preservation Strength"
+                        )
+                    
+                    seed = gr.Number(
+                        label="Seed (use -1 for random)",
+                        value=82706,
+                        precision=0
+                    )
+                    
+
+                
+                progress = gr.Progress(track_tqdm=False)
+
+            # Output column
+            with gr.Column(scale=1):
+                output_image = gr.Image(
+                    label="Generated Image", 
+                    type="pil", 
+                    height=400,
+                    interactive=False,
+                    show_download_button=True,
+                    scale=1,
+                    container=True,
+                    image_mode="RGB"
+                )
+                
+                with gr.Accordion("Image Information", open=False):
+                    image_info = gr.JSON(label="Details")
+
+        def refine_instruction_ui(image, instruction):
+            if not image or not instruction:
+                return instruction
+            try:
+                refined = refine_instruction(image, instruction)
+                if len(refined) > 0:
+                    return refined
+                else:
+                    logger.warning("Instruction refinement service returned empty result")
+                    gr.Warning("Instruction refinement service is currently not working. Please try again later.")
+                    return instruction
+            except Exception as e:
+                logger.error(f"Error refining instruction: {str(e)}")
+                gr.Warning("Instruction refinement service is currently not working. Please try again later.")
+                return instruction
+
+        # Generate function with progress updates
+        def generate_with_progress(image, instruction, seed, guidance_scale, image_guidance_scale, progress=gr.Progress()):
+            logger.info(f"Starting image generation with instruction='{instruction[:50]}...', seed={seed}")
+            
+            try:
+                if not image:
+                    logger.error("No image provided in UI")
+                    return None, None
+
+                if not instruction.strip():
+                    logger.error("Empty instruction provided in UI")
+                    return None, None
+
+                # Create request
+                logger.info("Creating API request")
+                task_id, used_seed = create_request(
+                    prompt=instruction,
+                    image=image,
+                    guidance_scale=guidance_scale,
+                    image_guidance_scale=image_guidance_scale,
+                    seed=seed
+                )
+
+                # Poll for results
+                start_time = time.time()
+                last_completion_ratio = 0
+                progress(0, desc="Initializing...")
+                logger.info(f"Starting to poll for results for task ID: {task_id}")
+
+                while time.time() - start_time < MAX_POLL_TIME:
+                    result = get_results(task_id)
+                    if not result:
+                        time.sleep(POLL_INTERVAL)
+                        continue
+
+                    sub_results = result.get("result", {}).get("sub_task_results", [])
+                    if not sub_results:
+                        time.sleep(POLL_INTERVAL)
+                        continue
+
+                    status = sub_results[0].get("task_status")
+                    logger.debug(f"Task status for ID {task_id}: {status}")
+
+                    # Get and display completion ratio
+                    completion_ratio = sub_results[0].get('task_completion', 0) * 100
+                    if completion_ratio != last_completion_ratio:
+                        # Only update UI when completion ratio changes
+                        last_completion_ratio = completion_ratio
+                        progress(completion_ratio / 100, desc=f"Generating image")
+                        logger.info(f"Generation progress - Task ID: {task_id}, Completion: {completion_ratio:.1f}%")
+
+                    # Check task status
+                    if status == 1:  # Success
+                        logger.info(f"Task completed successfully - Task ID: {task_id}")
+                        progress(1.0, desc="Generation complete")
+                        image_name = sub_results[0].get("image")
+                        if not image_name:
+                            logger.error(f"No image name in successful response. Response: {sub_results[0]}")
+                            return None, None
+
+                        image_url = f"{API_IMAGE_URL}{image_name}.png"
+                        logger.info(f"Downloading image - Task ID: {task_id}, URL: {image_url}")
+                        image = download_image(image_url)
+
+                        if image:
+                            # Add metadata to the image
+                            logger.info(f"Adding metadata to image - Task ID: {task_id}")
+                            metadata = {
+                                "prompt": instruction,
+                                "seed": str(used_seed),
+                                "model": API_MODEL_NAME,
+                                "guidance_scale": str(guidance_scale),
+                                "image_guidance_scale": str(image_guidance_scale),
+                                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+                                "generated_by": "HiDream-E1-Full Editor"
+                            }
+
+                            image_with_metadata = add_metadata_to_image(image, metadata)
+
+                            # Create info for display
+                            info = {
+                                "model": API_MODEL_NAME,
+                                "prompt": instruction,
+                                "seed": used_seed,
+                                "guidance_scale": guidance_scale,
+                                "image_guidance_scale": image_guidance_scale,
+                                "generated_at": time.strftime("%Y-%m-%d %H:%M:%S")
+                            }
+
+                            logger.info(f"Image generation complete - Task ID: {task_id}")
+                            return image_with_metadata, info
+                        else:
+                            logger.error(f"Failed to download image - Task ID: {task_id}, URL: {image_url}")
+                            return None, None
+
+                    elif status in {3, 4}:  # Failed or Canceled
+                        error_msg = sub_results[0].get("task_error", "Unknown error")
+                        logger.error(f"Task failed - Task ID: {task_id}, Status: {status}, Error: {error_msg}")
+                        return None, None
+
+                    time.sleep(POLL_INTERVAL)
+
+                logger.error(f"Timeout waiting for task completion - Task ID: {task_id}, Max time: {MAX_POLL_TIME}s")
+                return None, None
+
+            except Exception as e:
+                logger.error(f"Error during image generation: {str(e)}")
+                logger.error(f"Full traceback: {traceback.format_exc()}")
+                return None, None
+
+        # Set up event handlers
+        refine_btn.click(
+            fn=refine_instruction_ui,
+            inputs=[input_image, instruction],
+            outputs=[instruction]
+        )
+
+        generate_btn.click(
+            fn=generate_with_progress,
+            inputs=[input_image, instruction, seed, guidance_scale, image_guidance_scale],
+            outputs=[output_image, image_info]
+        )
+
+        # Define a combined function to refine instruction and then generate image
+        def refine_and_generate(image, instruction, seed, guidance_scale, image_guidance_scale, progress=gr.Progress()):
+            try:
+                # First refine the instruction
+                if not image or not instruction:
+                    return None, None, instruction
+                    
+                logger.info(f"Refining instruction: '{instruction[:50]}...'")
+                refined_instruction = refine_instruction_ui(image, instruction)
+                
+                if not refined_instruction or refined_instruction.strip() == "":
+                    logger.warning("Instruction refinement failed, using original instruction")
+                    refined_instruction = instruction
+                    gr.Warning("Instruction refinement failed, using original instruction instead.")
+                else:
+                    logger.info(f"Instruction refined to: '{refined_instruction[:50]}...'")
+                
+                # Then generate with the refined instruction
+                progress(0.2, desc="Instruction refined, generating image...")
+                generated_image, image_info = generate_with_progress(image, refined_instruction, seed, guidance_scale, image_guidance_scale, progress)
+                return generated_image, image_info, refined_instruction
+            except Exception as e:
+                logger.error(f"Error in refine_and_generate: {str(e)}")
+                logger.error(f"Full traceback: {traceback.format_exc()}")
+                gr.Warning(f"An error occurred during processing: {str(e)}")
+                return None, None, instruction
+
+        # Examples
+        gr.Examples(
+            examples=[
+                ["assets/test_1.png", "convert the image into a Ghibli style",3, 5, 4],
+                ["assets/test_1.png", "change the image into Disney Pixar style",3, 5, 4],
+                ["assets/test_1.png", "turn to sketch style",3, 5, 4],
+                ["assets/test_1.png", "add a sunglasses to the girl",3, 5, 2],
+                ["assets/test_1.png", "change the background to a sunset",3, 5, 2],
+                ["assets/test_2.jpg", "convert this image into a ink sketch image",3, 5, 2],
+                ["assets/test_2.jpg", "add butterfly'",3, 5, 2],
+                ["assets/test_2.jpg", "remove the wooden sign'",3, 5, 2],
+            ],
+            inputs=[input_image, instruction, seed, guidance_scale, image_guidance_scale],
+            outputs=[output_image, image_info, instruction],
+            fn=refine_and_generate,
+            cache_examples=True,
+            cache_mode = "lazy"
+        )
+
+    logger.info("Gradio UI created successfully")
+    return demo
+
+# Launch app
+if __name__ == "__main__":
+    logger.info("Starting HiDream-E1-Full Image Generator application")
+    demo = create_ui()
+    logger.info("Launching Gradio interface with queue")
+    demo.queue(max_size=50, default_concurrency_limit=8).launch(show_api=False)
+    logger.info("Application shutdown")