Update app.py

app.py

@@ -8,11 +8,7 @@ from PIL import Image
 from typing import Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
-from huggingface_hub import snapshot_download
 
-# --------------------------
-# Theme Setup
-# --------------------------
 colors.steel_blue = colors.Color(
     name="steel_blue",
     c50="#EBF3F8",
@@ -81,127 +77,129 @@ class SteelBlueTheme(Soft):
 
 steel_blue_theme = SteelBlueTheme()
 
-# --------------------------
-# Device Setup
-# --------------------------
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-print("Using device:", device)
-
-# --------------------------
-# Automatic download paths
-# --------------------------
-MODEL_DIR = "./models/Qwen-Image-Edit-2509"
-TRANSFORMER_DIR = "./models/Qwen-Image-Edit-Rapid-AIO/transformer"
-LORA_DIRS = {
-    "anime": "./loras/anime",
-    "multiple-angles": "./loras/multiple-angles",
-    "light-restoration": "./loras/light-restoration",
-    "relight": "./loras/relight",
-    "multi-angle-lighting": "./loras/multi-angle-lighting",
-    "edit-skin": "./loras/edit-skin",
-    "next-scene": "./loras/next-scene",
-    "upscale-image": "./loras/upscale-image",
-}
+print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
+print("torch.__version__ =", torch.__version__)
+print("torch.version.cuda =", torch.version.cuda)
+print("cuda available:", torch.cuda.is_available())
+print("cuda device count:", torch.cuda.device_count())
+if torch.cuda.is_available():
+    print("current device:", torch.cuda.current_device())
+    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
 
-LORA_WEIGHTS = {
-    "anime": "Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors",
-    "multiple-angles": "镜头转换.safetensors",
-    "light-restoration": "移除光影.safetensors",
-    "relight": "Qwen-Edit-Relight.safetensors",
-    "multi-angle-lighting": "多角度灯光-251116.safetensors",
-    "edit-skin": "qwen-edit-skin_1.1_000002750.safetensors",
-    "next-scene": "next-scene_lora-v2-3000.safetensors",
-    "upscale-image": "qwen-edit-enhance_64-v3_000001000.safetensors",
-}
-
-# Function to download repos if missing
-from huggingface_hub import snapshot_download
-
-def ensure_model(repo_id, local_dir):
-    if not os.path.exists(local_dir):
-        print(f"Downloading {repo_id} to {local_dir}...")
-        snapshot_download(repo_id=repo_id, local_dir=local_dir)
-    else:
-        print(f"{local_dir} already exists.")
-
-# Download main model and transformer
-ensure_model("Qwen/Qwen-Image-Edit-2509", MODEL_DIR)
-ensure_model("linoyts/Qwen-Image-Edit-Rapid-AIO", TRANSFORMER_DIR)
-
-# Download all LoRAs
-ensure_model("autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime", LORA_DIRS["anime"])
-ensure_model("dx8152/Qwen-Edit-2509-Multiple-angles", LORA_DIRS["multiple-angles"])
-ensure_model("dx8152/Qwen-Image-Edit-2509-Light_restoration", LORA_DIRS["light-restoration"])
-ensure_model("dx8152/Qwen-Edit-2509-Relight", LORA_DIRS["relight"])
-ensure_model("dx8152/Qwen-Edit-2509-Multi-Angle-Lighting", LORA_DIRS["multi-angle-lighting"])
-ensure_model("tlennon-ie/qwen-edit-skin", LORA_DIRS["edit-skin"])
-ensure_model("lovis93/next-scene-qwen-image-lora-2509", LORA_DIRS["next-scene"])
-ensure_model("vafipas663/Qwen-Edit-2509-Upscale-LoRA", LORA_DIRS["upscale-image"])
+print("Using device:", device)
 
-# --------------------------
-# Import pipeline after ensuring downloads
-# --------------------------
 from diffusers import FlowMatchEulerDiscreteScheduler
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
-# --------------------------
-# Pipeline setup
-# --------------------------
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 pipe = QwenImageEditPlusPipeline.from_pretrained(
-    MODEL_DIR,
+    "Qwen/Qwen-Image-Edit-2509",
     transformer=QwenImageTransformer2DModel.from_pretrained(
-        TRANSFORMER_DIR,
+        "linoyts/Qwen-Image-Edit-Rapid-AIO", # [transformer weights extracted from: Phr00t/Qwen-Image-Edit-Rapid-AIO]
         subfolder='transformer',
         torch_dtype=dtype,
-        device_map='cuda' if torch.cuda.is_available() else None
+        device_map='cuda'
     ),
     torch_dtype=dtype
 ).to(device)
 
-# Load LoRAs locally
-for adapter, lora_path in LORA_DIRS.items():
-    pipe.load_lora_weights(lora_path, weight_name=LORA_WEIGHTS[adapter], adapter_name=adapter)
+pipe.load_lora_weights("autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime",
+                       weight_name="Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors",
+                       adapter_name="anime")
+pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multiple-angles",
+                       weight_name="镜头转换.safetensors",
+                       adapter_name="multiple-angles")
+pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Light_restoration",
+                       weight_name="移除光影.safetensors",
+                       adapter_name="light-restoration")
+pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Relight",
+                       weight_name="Qwen-Edit-Relight.safetensors",
+                       adapter_name="relight")
+pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multi-Angle-Lighting",
+                       weight_name="多角度灯光-251116.safetensors",
+                       adapter_name="multi-angle-lighting")
+pipe.load_lora_weights("tlennon-ie/qwen-edit-skin",
+                       weight_name="qwen-edit-skin_1.1_000002750.safetensors",
+                       adapter_name="edit-skin")
+pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509",
+                       weight_name="next-scene_lora-v2-3000.safetensors",
+                       adapter_name="next-scene")
+pipe.load_lora_weights("vafipas663/Qwen-Edit-2509-Upscale-LoRA",
+                       weight_name="qwen-edit-enhance_64-v3_000001000.safetensors",
+                       adapter_name="upscale-image")
 
 pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
-
-# --------------------------
-# Misc constants
-# --------------------------
 MAX_SEED = np.iinfo(np.int32).max
 
-# --------------------------
-# Helper functions
-# --------------------------
 def update_dimensions_on_upload(image):
     if image is None:
         return 1024, 1024
-    w, h = image.size
-    if w > h:
-        new_w = 1024
-        new_h = int(1024 * h / w)
+    
+    original_width, original_height = image.size
+    
+    if original_width > original_height:
+        new_width = 1024
+        aspect_ratio = original_height / original_width
+        new_height = int(new_width * aspect_ratio)
     else:
-        new_h = 1024
-        new_w = int(1024 * w / h)
-    return (new_w // 8) * 8, (new_h // 8) * 8
+        new_height = 1024
+        aspect_ratio = original_width / original_height
+        new_width = int(new_height * aspect_ratio)
+        
+    # Ensure dimensions are multiples of 8
+    new_width = (new_width // 8) * 8
+    new_height = (new_height // 8) * 8
+    
+    return new_width, new_height
 
 @spaces.GPU(duration=30)
-def infer(input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
+def infer(
+    input_image,
+    prompt,
+    lora_adapter,
+    seed,
+    randomize_seed,
+    guidance_scale,
+    steps,
+    progress=gr.Progress(track_tqdm=True)
+):
     if input_image is None:
         raise gr.Error("Please upload an image to edit.")
+
+    if lora_adapter == "Photo-to-Anime":
+        pipe.set_adapters(["anime"], adapter_weights=[1.0])
+    elif lora_adapter == "Multiple-Angles":
+        pipe.set_adapters(["multiple-angles"], adapter_weights=[1.0])
+    elif lora_adapter == "Light-Restoration":
+        pipe.set_adapters(["light-restoration"], adapter_weights=[1.0])
+    elif lora_adapter == "Relight":
+        pipe.set_adapters(["relight"], adapter_weights=[1.0])
+    elif lora_adapter == "Multi-Angle-Lighting":
+        pipe.set_adapters(["multi-angle-lighting"], adapter_weights=[1.0])
+    elif lora_adapter == "Edit-Skin":
+        pipe.set_adapters(["edit-skin"], adapter_weights=[1.0])
+    elif lora_adapter == "Next-Scene":
+        pipe.set_adapters(["next-scene"], adapter_weights=[1.0])
+    elif lora_adapter == "Upscale-Image":
+        pipe.set_adapters(["upscale-image"], adapter_weights=[1.0])
+        
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+
     generator = torch.Generator(device=device).manual_seed(seed)
     negative_prompt = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
+
     original_image = input_image.convert("RGB")
+    
+    # Use the new function to update dimensions
     width, height = update_dimensions_on_upload(original_image)
 
-    # Set the adapter
-    pipe.set_adapters([lora_adapter.lower().replace("-", "_")], adapter_weights=[1.0])
-
     result = pipe(
         image=original_image,
         prompt=prompt,
@@ -212,54 +210,75 @@ def infer(input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scal
         generator=generator,
         true_cfg_scale=guidance_scale,
     ).images[0]
+
     return result, seed
 
 @spaces.GPU(duration=30)
 def infer_example(input_image, prompt, lora_adapter):
-    return infer(input_image, prompt, lora_adapter, 0, True, 1.0, 4)
+    input_pil = input_image.convert("RGB")
+    guidance_scale = 1.0
+    steps = 4
+    result, seed = infer(input_pil, prompt, lora_adapter, 0, True, guidance_scale, steps)
+    return result, seed
+
 
-# --------------------------
-# Gradio interface
-# --------------------------
 css="""
-#col-container {margin:0 auto; max-width:960px;}
-#main-title h1 {font-size:2.1em !important;}
+#col-container {
+    margin: 0 auto;
+    max-width: 960px;
+}
+#main-title h1 {font-size: 2.1em !important;}
 """
 
 with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# **Qwen-Image-Edit-2509-LoRAs-Fast**", elem_id="main-title")
         gr.Markdown("Perform diverse image edits using specialized [LoRA](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image-Edit-2509) adapters for the [Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) model.")
+
         with gr.Row(equal_height=True):
             with gr.Column():
                 input_image = gr.Image(label="Upload Image", type="pil", height=290)
-                prompt = gr.Text(label="Edit Prompt", show_label=True, placeholder="e.g., transform into anime..")
+                
+                prompt = gr.Text(
+                    label="Edit Prompt",
+                    show_label=True,
+                    placeholder="e.g., transform into anime..",
+                )
+
                 run_button = gr.Button("Edit Image", variant="primary")
+
             with gr.Column():
                 output_image = gr.Image(label="Output Image", interactive=False, format="png", height=353)
-                lora_adapter = gr.Dropdown(
-                    label="Choose Editing Style",
-                    choices=list(LORA_DIRS.keys()),
-                    value="anime"
-                )
+                
+                with gr.Row():
+                    lora_adapter = gr.Dropdown(
+                        label="Choose Editing Style",
+                        choices=["Photo-to-Anime", "Multiple-Angles", "Light-Restoration", "Multi-Angle-Lighting", "Upscale-Image", "Relight", "Next-Scene", "Edit-Skin"],
+                        value="Photo-to-Anime"
+                    )
                 with gr.Accordion("Advanced Settings", open=False, visible=False):
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                     guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
                     steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=4)
-
+        
         gr.Examples(
             examples=[
-                ["examples/1.jpg", "Transform into anime.", "anime"],
-                ["examples/5.jpg", "Remove shadows and relight the image using soft lighting.", "light-restoration"],
-                ["examples/4.jpg", "Use a subtle golden-hour filter with smooth light diffusion.", "relight"],
-                ["examples/2.jpeg", "Rotate the camera 45 degrees to the left.", "multiple-angles"],
-                ["examples/7.jpg", "Light source from the Right Rear", "multi-angle-lighting"],
-                ["examples/10.jpeg", "Upscale the image.", "upscale-image"],
-                ["examples/7.jpg", "Light source from the Below", "multi-angle-lighting"],
-                ["examples/2.jpeg", "Switch the camera to a top-down right corner view.", "multiple-angles"],
-                ["examples/9.jpg", "The camera moves slightly forward as sunlight breaks through the clouds, casting a soft glow around the character's silhouette in the mist. Realistic cinematic style, atmospheric depth.", "next-scene"],
-                ["examples/8.jpg", "Make the subjects skin details more prominent and natural.", "edit-skin"],
+                ["examples/1.jpg", "Transform into anime.", "Photo-to-Anime"],
+                ["examples/5.jpg", "Remove shadows and relight the image using soft lighting.", "Light-Restoration"],
+                ["examples/4.jpg", "Use a subtle golden-hour filter with smooth light diffusion.", "Relight"],
+                ["examples/2.jpeg", "Rotate the camera 45 degrees to the left.", "Multiple-Angles"],
+                ["examples/7.jpg", "Light source from the Right Rear", "Multi-Angle-Lighting"],
+                ["examples/10.jpeg", "Upscale the image.", "Upscale-Image"],
+                ["examples/7.jpg", "Light source from the Below", "Multi-Angle-Lighting"],
+                ["examples/2.jpeg", "Switch the camera to a top-down right corner view.", "Multiple-Angles"],
+                ["examples/9.jpg", "The camera moves slightly forward as sunlight breaks through the clouds, casting a soft glow around the character's silhouette in the mist. Realistic cinematic style, atmospheric depth.", "Next-Scene"],
+                ["examples/8.jpg", "Make the subjects skin details more prominent and natural.", "Edit-Skin"],
+                ["examples/6.jpg", "Switch the camera to a bottom-up view.", "Multiple-Angles"],
+                ["examples/6.jpg", "Rotate the camera 180 degrees upside down.", "Multiple-Angles"],
+                ["examples/4.jpg", "Rotate the camera 45 degrees to the right.", "Multiple-Angles"],
+                ["examples/4.jpg", "Switch the camera to a top-down view.", "Multiple-Angles"],
+                ["examples/4.jpg", "Switch the camera to a wide-angle lens.", "Multiple-Angles"],
             ],
             inputs=[input_image, prompt, lora_adapter],
             outputs=[output_image, seed],
@@ -275,4 +294,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=30).launch(css=css, theme=steel_blue_theme, mcp_server=True, ssr_mode=False, show_error=True)
+    demo.queue(max_size=30).launch(css=css, theme=steel_blue_theme, mcp_server=True, ssr_mode=False, show_error=True)