Update web.py

Added image format selection and Gradio cache clean up button and some info windows.

hdi1/web.py

@@ -2,11 +2,18 @@ import torch
 import gradio as gr
 import logging
 import os
+import tempfile
+import glob
 from datetime import datetime
+from PIL import Image
 from .nf4 import *
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Output directory for saving images
-OUTPUT_DIR = ".\outputs"
+OUTPUT_DIR = os.path.join("outputs")
 
 # Resolution options
 RESOLUTION_OPTIONS = [
@@ -25,72 +32,191 @@ SCHEDULER_OPTIONS = [
     "FlowUniPCMultistepScheduler"
 ]
 
+# Image format options
+IMAGE_FORMAT_OPTIONS = ["PNG", "JPEG", "WEBP"]
+
 # Parse resolution string to get height and width
 def parse_resolution(resolution_str):
-    return tuple(map(int, resolution_str.split("(")[0].strip().split(" × ")))
+    try:
+        return tuple(map(int, resolution_str.split("(")[0].strip().split(" × ")))
+    except (ValueError, IndexError) as e:
+        raise ValueError("Invalid resolution format") from e
 
-def gen_img_helper(model, prompt, res, seed, scheduler, guidance_scale, num_inference_steps, shift):
-    global pipe, current_model
+def clean_previous_temp_files():
+    """Delete temporary files from previous generations matching hdi1_* pattern and log Gradio temp files."""
+    temp_dir = tempfile.gettempdir()
+    patterns = [os.path.join(temp_dir, f"hdi1_*.{ext}") for ext in ["png", "jpeg", "webp"]]
+    deleted_files = []
+    
+    # Clean hdi1_* files
+    for pattern in patterns:
+        for temp_file in glob.glob(pattern):
+            try:
+                os.remove(temp_file)
+                deleted_files.append(temp_file)
+                logger.info(f"Deleted temporary file: {temp_file}")
+            except OSError as e:
+                logger.warning(f"Failed to delete temporary file {temp_file}: {str(e)}")
+    
+    # Log Gradio temp files (for monitoring)
+    gradio_temp_dir = os.path.join(temp_dir, "gradio")
+    if os.path.exists(gradio_temp_dir):
+        for root, _, files in os.walk(gradio_temp_dir):
+            for file in files:
+                if file.endswith((".png", ".jpeg", ".webp")):
+                    gradio_file = os.path.join(root, file)
+                    logger.info(f"Found Gradio temporary file: {gradio_file}")
+    
+    return deleted_files
 
-    # 1. Check if the model matches loaded model, load the model if not
-    if model != current_model:
-        print(f"Unloading model {current_model}...")
-        if pipe is not None:
-            del pipe
-            torch.cuda.empty_cache()
+def clean_all_temp_files():
+    """Manually clean hdi1_* and Gradio temporary files, with user confirmation."""
+    status_message = "Starting temporary file cleanup..."
+    logger.info(status_message)
+    
+    try:
+        # Clean hdi1_* files
+        deleted_files = clean_previous_temp_files()
         
-        print(f"Loading model {model}...")
-        pipe, _ = load_models(model)
-        current_model = model
-        print("Model loaded successfully!")
+        # Clean Gradio temp files
+        temp_dir = tempfile.gettempdir()
+        gradio_temp_dir = os.path.join(temp_dir, "gradio")
+        if os.path.exists(gradio_temp_dir):
+            for root, _, files in os.walk(gradio_temp_dir):
+                for file in files:
+                    if file.endswith((".png", ".jpeg", ".webp")):
+                        gradio_file = os.path.join(root, file)
+                        try:
+                            os.remove(gradio_file)
+                            deleted_files.append(gradio_file)
+                            logger.info(f"Deleted Gradio temporary file: {gradio_file}")
+                        except OSError as e:
+                            logger.warning(f"Failed to delete Gradio temporary file {gradio_file}: {str(e)}")
+        
+        status_message = f"Cleanup complete. Deleted {len(deleted_files)} files."
+        logger.info(status_message)
+        return status_message
+    except Exception as e:
+        error_message = f"Cleanup error: {str(e)}"
+        logger.error(error_message)
+        return error_message
 
-    # 2. Update scheduler
-    config = MODEL_CONFIGS[model]
-    scheduler_map = {
-        "FlashFlowMatchEulerDiscreteScheduler": FlashFlowMatchEulerDiscreteScheduler,
-        "FlowUniPCMultistepScheduler": FlowUniPCMultistepScheduler
-    }
-    scheduler_class = scheduler_map[scheduler]
-    device = pipe._execution_device
+def gen_img_helper(model, prompt, res, seed, scheduler, guidance_scale, num_inference_steps, shift, image_format):
+    global pipe, current_model
+    status_message = "Starting image generation..."
 
-    # Set scheduler with shift for flow-matching schedulers
-    pipe.scheduler = scheduler_class(num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False)
+    try:
+        # Clean up previous temporary files
+        status_message = "Cleaning up previous temporary files..."
+        logger.info(status_message)
+        clean_previous_temp_files()
+        status_message = "Previous temporary files cleaned."
 
-    # 3. Generate image
-    res = parse_resolution(res)
-    image, seed = generate_image(pipe, model, prompt, res, seed, guidance_scale, num_inference_steps)
-    
-    # 4. Save image locally
-    os.makedirs(OUTPUT_DIR, exist_ok=True)
-    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-    output_path = os.path.join(OUTPUT_DIR, f"output_{timestamp}.png")
-    image.save(output_path)
-    
-    return image, seed, f"Image saved to: {output_path}"
+        # Validate inputs
+        if not prompt or len(prompt.strip()) == 0:
+            raise ValueError("Prompt cannot be empty")
+        if not isinstance(seed, (int, float)) or seed < -1:
+            raise ValueError("Seed must be -1 or a non-negative integer")
+        if num_inference_steps < 1 or num_inference_steps > 100:
+            raise ValueError("Number of inference steps must be between 1 and 100")
+        if guidance_scale < 0 or guidance_scale > 10:
+            raise ValueError("Guidance scale must be between 0 and 10")
+        if shift < 1 or shift > 10:
+            raise ValueError("Shift must be between 1 and 10")
+
+        # 1. Check if the model matches loaded model, load the model if not
+        if model != current_model:
+            status_message = f"Unloading model {current_model}..."
+            logger.info(status_message)
+            if pipe is not None:
+                del pipe
+                torch.cuda.empty_cache()
+            
+            status_message = f"Loading model {model}..."
+            logger.info(status_message)
+            pipe, _ = load_models(model)
+            current_model = model
+            status_message = "Model loaded successfully!"
+            logger.info(status_message)
+
+        # 2. Update scheduler
+        config = MODEL_CONFIGS[model]
+        scheduler_map = {
+            "FlashFlowMatchEulerDiscreteScheduler": FlashFlowMatchEulerDiscreteScheduler,
+            "FlowUniPCMultistepScheduler": FlowUniPCMultistepScheduler
+        }
+        if scheduler not in scheduler_map:
+            raise ValueError(f"Invalid scheduler: {scheduler}")
+        scheduler_class = scheduler_map[scheduler]
+        device = pipe._execution_device
+
+        # Set scheduler with shift for flow-matching schedulers
+        pipe.scheduler = scheduler_class(num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False)
+
+        # 3. Generate image
+        status_message = "Generating image..."
+        logger.info(status_message)
+        res = parse_resolution(res)
+        image, seed = generate_image(pipe, model, prompt, res, seed, guidance_scale, num_inference_steps)
+        
+        # 4. Save image locally with selected format
+        status_message = "Saving image locally..."
+        logger.info(status_message)
+        os.makedirs(OUTPUT_DIR, exist_ok=True)
+        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        file_extension = image_format.lower()
+        output_path = os.path.join(OUTPUT_DIR, f"output_{timestamp}.{file_extension}")
+        if image_format == "JPEG":
+            image = image.convert("RGB")  # JPEG doesn't support RGBA
+        image.save(output_path, format=image_format)
+        logger.info(f"Image saved to {output_path}")
+        
+        # 5. Prepare image for download in selected format
+        status_message = "Preparing image for download..."
+        logger.info(status_message)
+        download_filename = f"generated_image_{timestamp}.{file_extension}"
+        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_extension}", prefix="hdi1_") as temp_file:
+            if image_format == "JPEG":
+                image = image.convert("RGB")  # Ensure JPEG compatibility
+            image.save(temp_file, format=image_format)
+            temp_file_path = temp_file.name
+        logger.info(f"Temporary file created at {temp_file_path}")
+        
+        status_message = "Image generation complete!"
+        logger.info(status_message)
+        return image, seed, f"Image saved to: {output_path}", temp_file_path, status_message
+
+    except Exception as e:
+        error_message = f"Error: {str(e)}"
+        logger.error(error_message)
+        return None, None, None, None, error_message
 
 def generate_image(pipe, model_type, prompt, resolution, seed, guidance_scale, num_inference_steps):
-    # Parse resolution
-    width, height = resolution
+    try:
+        # Parse resolution
+        width, height = resolution
 
-    # Handle seed
-    if seed == -1:
-        seed = torch.randint(0, 1000000, (1,)).item()
-    
-    generator = torch.Generator("cuda").manual_seed(seed)
-    
-    # Common parameters
-    params = {
-        "prompt": prompt,
-        "height": height,
-        "width": width,
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": num_inference_steps,
-        "num_images_per_prompt": 1,
-        "generator": generator
-    }
+        # Handle seed
+        if seed == -1:
+            seed = torch.randint(0, 1000000, (1,)).item()
+        
+        generator = torch.Generator("cuda").manual_seed(seed)
+        
+        # Common parameters
+        params = {
+            "prompt": prompt,
+            "height": height,
+            "width": width,
+            "guidance_scale": guidance_scale,
+            "num_inference_steps": num_inference_steps,
+            "num_images_per_prompt": 1,
+            "generator": generator
+        }
 
-    images = pipe(**params).images
-    return images[0], seed
+        images = pipe(**params).images
+        return images[0], seed
+    except Exception as e:
+        raise RuntimeError(f"Image generation failed: {str(e)}") from e
 
 if __name__ == "__main__":
     logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
@@ -102,6 +228,7 @@ if __name__ == "__main__":
     # Create Gradio interface
     with gr.Blocks(title="HiDream-I1-nf4 Dashboard") as demo:
         gr.Markdown("# HiDream-I1-nf4 Dashboard")
+        gr.Markdown("**Note**: Use the 'Download Image' link below to download the image in your selected format (PNG, JPEG, or WEBP). Downloading from the image preview's download button is WEBP format.")
         
         with gr.Row():
             with gr.Column():
@@ -165,17 +292,32 @@ if __name__ == "__main__":
                     info="Scheduler shift parameter for flow-matching schedulers. Use 1.0–5.0; 3.0 is a good default."
                 )
                 
+                image_format = gr.Radio(
+                    choices=IMAGE_FORMAT_OPTIONS,
+                    value="PNG",
+                    label="Image Format",
+                    info="Select the format for the saved and downloaded image."
+                )
+                
                 generate_btn = gr.Button("Generate Image")
+                cleanup_btn = gr.Button("Clean Temporary Files")
                 
             with gr.Column():
+                status_message = gr.Textbox(label="Status", value="Ready", interactive=False)
                 output_image = gr.Image(label="Generated Image", type="pil")
                 seed_used = gr.Number(label="Seed Used", interactive=False)
                 save_path = gr.Textbox(label="Saved Image Path", interactive=False)
+                download_file = gr.File(label="Download Image", interactive=False, file_types=[".png", ".jpeg", ".webp"])
         
         generate_btn.click(
             fn=gen_img_helper,
-            inputs=[model_type, prompt, resolution, seed, scheduler, guidance_scale, num_inference_steps, shift],
-            outputs=[output_image, seed_used, save_path]
+            inputs=[model_type, prompt, resolution, seed, scheduler, guidance_scale, num_inference_steps, shift, image_format],
+            outputs=[output_image, seed_used, save_path, download_file, status_message]
+        )
+        cleanup_btn.click(
+            fn=clean_all_temp_files,
+            inputs=[],
+            outputs=[status_message]
         )
 
     demo.launch(share=True, pwa=True)