Update app.py

app.py

@@ -58,40 +58,44 @@ edit_pipe = LongCatImageEditPipeline.from_pretrained(
 edit_pipe.to(device, torch.bfloat16)
 
 print(f"✅ Image Edit model loaded successfully on {device}")
-
-# # --- Core Functions ---
-# @spaces.GPU(duration=120)
-# def generate_image(
-#     prompt: str,
-#     width: int,
-#     height: int,
-#     seed: int,
-#     progress=gr.Progress()
-# ):
-#     """Generate image from text prompt"""
-#     if not prompt or prompt.strip() == "":
-#         raise gr.Error("Please enter a prompt")
-#     try:
-#         progress(0.1, desc="Preparing generation...")
-#         progress(0.2, desc="Generating image...")
-#         generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
-#         with torch.inference_mode():
-#             output = t2i_pipe(
-#                 prompt,
-#                 negative_prompt="",
-#                 height=height,
-#                 width=width,
-#                 guidance_scale=4.5,
-#                 num_inference_steps=50,
-#                 num_images_per_prompt=1,
-#                 generator=generator,
-#                 enable_cfg_renorm=True,
-#                 enable_prompt_rewrite=True
-#             )
-#         progress(1.0, desc="Done!")
-#         return output.images[0]
-#     except Exception as e:
-#         raise gr.Error(f"Error during image generation: {str(e)}")
+def load_lora_auto(pipe, lora_input):
+    lora_input = lora_input.strip()
+    if not lora_input:
+        return
+
+    # If it's just an ID like "author/model"
+    if "/" in lora_input and not lora_input.startswith("http"):
+        pipe.load_lora_weights(lora_input)
+        return
+
+    if lora_input.startswith("http"):
+        url = lora_input
+
+        # Repo page (no blob/resolve)
+        if "huggingface.co" in url and "/blob/" not in url and "/resolve/" not in url:
+            repo_id = urlparse(url).path.strip("/")
+            pipe.load_lora_weights(repo_id)
+            return
+
+        # Blob link → convert to resolve link
+        if "/blob/" in url:
+            url = url.replace("/blob/", "/resolve/")
+
+        # Download direct file
+        tmp_dir = tempfile.mkdtemp()
+        local_path = os.path.join(tmp_dir, os.path.basename(urlparse(url).path))
+
+        try:
+            print(f"Downloading LoRA from {url}...")
+            resp = requests.get(url, stream=True)
+            resp.raise_for_status()
+            with open(local_path, "wb") as f:
+                for chunk in resp.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            print(f"Saved LoRA to {local_path}")
+            pipe.load_lora_weights(local_path)
+        finally:
+            shutil.rmtree(tmp_dir, ignore_errors=True)
 
 @spaces.GPU(duration=120)
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=4, num_inference_steps=28, lora_id=None, lora_scale=0.95, progress=gr.Progress(track_tqdm=True)):