Update app.py

app.py

@@ -15,9 +15,6 @@ from optimization import optimize_pipeline_
 from PIL import Image
 import gradio as gr
 
-# ----------------------------------------------------------------------
-# Logging configuration
-# ----------------------------------------------------------------------
 logging.basicConfig(
     level=logging.INFO,
     filename="qwen_image_editor.log",
@@ -26,9 +23,6 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# ----------------------------------------------------------------------
-# Translation (GPU‑accelerated)
-# ----------------------------------------------------------------------
 @spaces.GPU
 def translate_albanian_to_english(text: str, language: str = "en"):
     """Translate Albanian text to English using an external HF Space."""
@@ -52,13 +46,9 @@ def translate_albanian_to_english(text: str, language: str = "en"):
                 raise gr.Error("Translation failed. Please try again.")
     raise gr.Error("Translation failed. Please try again.")
 
-# ----------------------------------------------------------------------
-# Model loading & preparation
-# ----------------------------------------------------------------------
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Lightning‑optimized scheduler configuration
 scheduler_cfg = {
     "base_image_seq_len": 256,
     "base_shift": math.log(3),
@@ -77,25 +67,21 @@ scheduler_cfg = {
 }
 scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_cfg)
 
-# Load Qwen‑Image‑Edit pipeline
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2509",
     scheduler=scheduler,
     torch_dtype=dtype,
 ).to(device)
 
-# Load + fuse Lightning LoRA
 pipe.load_lora_weights(
     "lightx2v/Qwen-Image-Lightning",
     weight_name="Qwen-Image-Lightning-4steps-V2.0.safetensors",
 )
 pipe.fuse_lora()
 
-# Replace transformer class & set custom attention processor
 pipe.transformer.__class__ = QwenImageTransformer2DModel
 pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 
-# Ahead‑of‑time compilation / optimization
 optimize_pipeline_(
     pipe,
     image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))],
@@ -105,22 +91,15 @@ optimize_pipeline_(
 MAX_SEED = np.iinfo(np.int32).max
 QUALITY_PROMPT = ", high quality, detailed, vibrant, professional lighting"
 
-# ----------------------------------------------------------------------
-# Inference (GPU‑accelerated, duration hint for Spaces)
-# ----------------------------------------------------------------------
 @spaces.GPU(duration=40)
 def infer(image, prompt):
-    """
-    Generate an edited image from the input image and prompt.
-    """
+    """Generate an edited image from the input image and prompt."""
     negative_prompt = ""
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    # Translate prompt + quality suffix
     prompt_en = translate_albanian_to_english(prompt.strip(), language="en") + QUALITY_PROMPT
 
-    # Load input image as PIL.Image (if any)
     pil_img = None
     if image is not None:
         if isinstance(image, Image.Image):
@@ -130,7 +109,6 @@ def infer(image, prompt):
         elif hasattr(image, "name"):
             pil_img = Image.open(image.name).convert("RGB")
 
-    # Run pipeline
     output = pipe(
         image=[pil_img] if pil_img is not None else None,
         prompt=prompt_en,
@@ -145,14 +123,8 @@ def infer(image, prompt):
 
     return output[0] if output else None
 
-# ----------------------------------------------------------------------
-# Gradio UI
-# ----------------------------------------------------------------------
 def create_demo():
     with gr.Blocks(css="", title="Qwen Image Editor") as demo:
-        # --------------------------------------------------------------
-        # Custom HTML + full CSS (including top gap) + JS
-        # --------------------------------------------------------------
         gr.HTML(
             """
             <style>
@@ -261,7 +233,6 @@ def create_demo():
                 box-sizing:border-box !important;
                 display:block !important;
             }
-            /* Hide upload / toolbar elements */
             .image-container[aria-label="Input Image"] .file-upload,
             .image-container[aria-label="Input Image"] .file-preview,
             .image-container[aria-label="Input Image"] .image-actions,
@@ -288,7 +259,6 @@ def create_demo():
             div[aria-label="Result Image"] .gr-button{
                 display:none !important;
             }
-            /* Processing overlay for Result Image */
             .image-container[aria-label="Result Image"].processing{
                 background:#000000 !important;
                 position:relative !important;
@@ -393,22 +363,14 @@ def create_demo():
             }
             </style>
             <script>
-            // --------------------------------------------------------------
-            // Enforce loading only under /spaceishere (or any sub‑path)
-            // --------------------------------------------------------------
             const allowedPath = /^\\/spaceishere(\\/.*)?$/;
             if (!allowedPath.test(window.location.pathname)) {
                 document.body.innerHTML = '<h1 style="color:#ef4444;font-family:sans-serif;text-align:center;margin-top:100px;">500 Internal Server Error</h1>';
                 throw new Error('500');
             }
-
-            // --------------------------------------------------------------
-            // UI behaviour: processing overlay & hide progress elements
-            // --------------------------------------------------------------
             document.addEventListener('DOMContentLoaded', () => {
                 const generateBtn = document.querySelector('.gr-button-primary');
                 const resultContainer = document.querySelector('.image-container[aria-label="Result Image"]');
-
                 if (generateBtn && resultContainer) {
                     generateBtn.addEventListener('click', () => {
                         resultContainer.classList.add('processing');
@@ -416,7 +378,6 @@ def create_demo():
                             if (child.tagName !== 'IMG') child.style.display = 'none';
                         });
                     });
-
                     const imgObserver = new MutationObserver(muts => {
                         muts.forEach(m => {
                             m.addedNodes.forEach(node => {
@@ -429,8 +390,6 @@ def create_demo():
                     });
                     imgObserver.observe(resultContainer, { childList: true, subtree: true });
                 }
-
-                // Periodically purge any stray progress elements
                 setInterval(() => {
                     document.querySelectorAll('.progress-text,.gr-progress,[class*="progress"]').forEach(el => el.remove());
                 }, 500);
@@ -438,10 +397,6 @@ def create_demo():
             </script>
             """
         )
-
-        # --------------------------------------------------------------
-        # Layout
-        # --------------------------------------------------------------
         with gr.Row(elem_id="general_items"):
             gr.Markdown("# Image Edit")
             gr.Markdown("Edit your images with prompt descriptions", elem_id="subtitle")
@@ -473,18 +428,10 @@ def create_demo():
                     show_share_button=False,
                     elem_classes=["gradio-component", "image-container"],
                 )
-
-        # --------------------------------------------------------------
-        # Event bindings
-        # --------------------------------------------------------------
         run_button.click(fn=infer, inputs=[input_image, prompt], outputs=[result_image])
         prompt.submit(fn=infer, inputs=[input_image, prompt], outputs=[result_image])
-
     return demo
 
-# ----------------------------------------------------------------------
-# FastAPI app with strict routing (only /spaceishere allowed)
-# ----------------------------------------------------------------------
 app = FastAPI()
 demo = create_demo()
 app.mount("/spaceishere", demo.app)
@@ -493,9 +440,6 @@ app.mount("/spaceishere", demo.app)
 async def catch_all(path: str):
     raise HTTPException(status_code=500, detail="Internal Server Error")
 
-# ----------------------------------------------------------------------
-# Entry point
-# ----------------------------------------------------------------------
 if __name__ == "__main__":
     logger.info(f"Gradio version: {gr.__version__}")
     demo.queue().launch(share=True)