Merge branch 'main' of https://huggingface.co/spaces/Onise/qwen-image-edit-aio-lora

app.py

@@ -30,6 +30,7 @@ import os
 import base64
 from io import BytesIO
 import json
+import tempfile
 
 SYSTEM_PROMPT = '''
 # Edit Instruction Rewriter
@@ -142,7 +143,7 @@ def polish_prompt_hf(original_prompt, img_list):
                 # If img is a PIL Image
                 if hasattr(img, 'save'):  # Check if it's a PIL Image
                     buffered = BytesIO()
-                    img.save(buffered, format="PNG")
+                    img.save(buffered, format="PNG", compress_level=0, optimize=False)
                     img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
                     image_url = f"data:image/png;base64,{img_base64}"
                 # If img is already a file path (string)
@@ -217,7 +218,7 @@ def polish_prompt_hf(original_prompt, img_list):
 def encode_image(pil_image):
     import io
     buffered = io.BytesIO()
-    pil_image.save(buffered, format="PNG")
+    pil_image.save(buffered, format="PNG", compress_level=0, optimize=False)
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 # --- Model Loading ---
@@ -398,17 +399,13 @@ torch.cuda.empty_cache()
 # optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
 
 # --- Custom LoRA loading (QWEN_MCNL) ---
-# If you want to temporarily disable this LoRA, comment out this block.
-try:
-    print("loading custom LoRA: Daverrrr75/QWEN_MCNL ...")
-    pipe.load_lora_weights(
-        "Daverrrr75/QWEN_MCNL",
-        weight_name="qwen_MCNL_v1.0.safetensors",
-    )
-    pipe.fuse_lora()
-    print("custom LoRA fused successfully.")
-except Exception as e:
-    print(f"Warning: failed to load/fuse custom LoRA Daverrrr75/QWEN_MCNL: {e}")
+print("loading custom LoRA: Daverrrr75/QWEN_MCNL ...")
+pipe.load_lora_weights(
+    "Daverrrr75/QWEN_MCNL",
+    weight_name="qwen_MCNL_v1.0.safetensors",
+)
+pipe.fuse_lora()
+print("custom LoRA fused successfully.")
 
 
 # --- UI Constants and Helpers ---
@@ -504,8 +501,33 @@ def infer(
             except Exception:
                 continue
 
-    if height==256 and width==256:
-        height, width = None, None
+    # Adjust dimensions to preserve input resolution and ensure quality
+    if len(pil_images) > 0:
+        input_width, input_height = pil_images[0].size
+        aspect_ratio = input_width / input_height
+    else:
+        aspect_ratio = 1.0  # Default square if no image
+
+    # Preserve input resolution by default - use input dimensions directly
+    if height is None and width is None:
+        # Use input dimensions directly, rounded to required multiple
+        width = round(input_width / 32) * 32
+        height = round(input_height / 32) * 32
+    elif height is not None and width is None:
+        width = round(height * aspect_ratio / 32) * 32
+    elif width is not None and height is None:
+        height = round(width / aspect_ratio / 32) * 32
+    else:
+        # Both provided: adjust width to match height * aspect_ratio
+        width = round(height * aspect_ratio / 32) * 32
+
+    # Quality safeguard: ensure minimum size to avoid excessive quality loss
+    min_size = 512
+    if height < min_size or width < min_size:
+        # Scale up while preserving ratio if too small
+        scale_factor = max(min_size / height, min_size / width)
+        height = round(height * scale_factor / 32) * 32
+        width = round(width * scale_factor / 32) * 32
     print(f"Calling pipeline with prompt: '{prompt}'")
     print(f"Negative Prompt: '{negative_prompt}'")
     print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
@@ -527,8 +549,15 @@ def infer(
         num_images_per_prompt=num_images_per_prompt,
     ).images
 
-    # Return images, seed, and make button visible
-    return image, seed, gr.update(visible=True)
+    # Save images as PNG with 100% quality (lossless, no compression)
+    output_paths = []
+    for i, img in enumerate(image):
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+            img.save(tmp_file.name, format="PNG", compress_level=0, optimize=False)
+            output_paths.append(tmp_file.name)
+
+    # Return paths, seed, and make button visible
+    return output_paths, seed, gr.update(visible=True)
 
 # --- Examples and UI Layout ---
 examples = []
@@ -568,7 +597,7 @@ with gr.Blocks(css=css) as demo:
                                           interactive=True)
 
             with gr.Column():
-                result = gr.Gallery(label="Result", show_label=False, type="pil", interactive=False)
+                result = gr.Gallery(label="Result", show_label=False, type="filepath", interactive=False)
                 # Add this button right after the result gallery - initially hidden
                 use_output_btn = gr.Button("↗️ Use as input", variant="secondary", size="sm", visible=False)
 

qwenimage/pipeline_qwenimage_edit_plus.py

@@ -627,8 +627,18 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
             [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
             returning a tuple, the first element is a list with the generated images.
         """
+        # Preserve input resolution - use input dimensions directly
         image_size = image[-1].size if isinstance(image, list) else image.size
-        calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
+        input_width, input_height = image_size[0], image_size[1]
+        aspect_ratio = input_width / input_height
+        
+        # Use input dimensions if not specified, rounded to required multiple
+        if height is None and width is None:
+            calculated_width = round(input_width / (self.vae_scale_factor * 2)) * (self.vae_scale_factor * 2)
+            calculated_height = round(input_height / (self.vae_scale_factor * 2)) * (self.vae_scale_factor * 2)
+        else:
+            calculated_width, calculated_height = calculate_dimensions(1024 * 1024, aspect_ratio)
+        
         height = height or calculated_height
         width = width or calculated_width
 
@@ -677,7 +687,11 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
                 condition_width, condition_height = calculate_dimensions(
                     CONDITION_IMAGE_SIZE, image_width / image_height
                 )
-                vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height)
+                # Use input dimensions for VAE processing to preserve quality
+                # Round to required multiple for VAE compatibility
+                multiple_of = self.vae_scale_factor * 2
+                vae_width = round(image_width / multiple_of) * multiple_of
+                vae_height = round(image_height / multiple_of) * multiple_of
                 condition_image_sizes.append((condition_width, condition_height))
                 vae_image_sizes.append((vae_width, vae_height))
                 condition_images.append(self.image_processor.resize(img, condition_height, condition_width))