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
# Set up logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# 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.
[Description of the Source Image]
[The specific, single-line editing instruction based on the Source Instruction and Source Image context]
[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 ""
import httpx
# Create a custom httpx client with verification disabled
insecure_client = httpx.Client(
verify=False, # THIS DISABLES SSL VERIFICATION - SECURITY RISK
timeout=httpx.Timeout(60.0, connect=10.0)
)
def refine_instruction(src_image, src_instruction):
MAX_TOKENS_RESPONSE = 500 # Limit response tokens as output format is structured
client = OpenAI(http_client=insecure_client)
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)
For more features and to experience the full capabilities of our product, please visit [https://vivago.ai/](https://vivago.ai/).
""")
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("""
Note: For optimal results, we recommend using the Refine Instruction button which formats your input into:
"Editing Instruction: [your instruction]. Target Image Description: [expected result]"
""")
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("""
Important: Adjust these parameters based on your editing needs:
- For style changes, use higher image preservation strength (e.g., 3.0-4.0)
- For local edits like adding, deleting, replacing elements, use lower image preservation strength (e.g., 2.0-3.0)
- If you notice visual artifacts or distortions in the generated image, try reduce the image preservation strength value.
""")
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",82706, 5, 4],
["assets/test_1.png", "change the image into Disney Pixar style",82706, 5, 4],
["assets/test_1.png", "add a sunglasses to the girl",82706, 5, 2],
["assets/test_2.jpg", "convert this image into a ink sketch image",82706, 5, 2],
["assets/test_2.jpg", "add butterfly",82706, 5, 2],
["assets/test_2.jpg", "remove the wooden sign",82706, 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")