Update app.py

app.py

@@ -1,17 +1,16 @@
-import os
 import gc
 import random
 import threading
 import traceback
 
 import gradio as gr
-import torch
 import spaces
+import torch
 
 from diffusers import QwenImage21Pipeline
 
 # ============================================================
-# CONFIGURATION
+# CONFIG
 # ============================================================
 
 BASE_MODEL = "Qwen/Qwen-Image-2.1"
@@ -25,7 +24,7 @@ pipe = None
 generation_lock = threading.Lock()
 
 # ============================================================
-# MODEL LOADING
+# LOAD PIPELINE
 # ============================================================
 
 def load_pipeline():
@@ -34,7 +33,7 @@ def load_pipeline():
     if pipe is not None:
         return pipe
 
-    print("Loading Qwen Image 2.1...", flush=True)
+    print("Loading Qwen-Image-2.1...", flush=True)
 
     pipeline = QwenImage21Pipeline.from_pretrained(
         BASE_MODEL,
@@ -42,7 +41,7 @@ def load_pipeline():
         low_cpu_mem_usage=True,
     )
 
-    print("Loading Pixelship LoRA...", flush=True)
+    print("Loading Pixelship 2.0 LoRA...", flush=True)
 
     pipeline.load_lora_weights(
         LORA_REPO,
@@ -58,18 +57,19 @@ def load_pipeline():
 
     pipe = pipeline
 
-    print("Pixelship 2 ready!", flush=True)
+    print("Pipeline ready.", flush=True)
 
     return pipe
 
 
 # ============================================================
-# GENERATION
+# GENERATE / EDIT
 # ============================================================
 
 @spaces.GPU(duration=120)
 def generate(
     prompt,
+    input_image,
     negative_prompt,
     width,
     height,
@@ -80,13 +80,13 @@ def generate(
     random_seed,
     progress=gr.Progress(),
 ):
-    if not prompt or not prompt.strip():
-        raise gr.Error("Please enter a prompt.")
+    if not prompt.strip():
+        raise gr.Error("Prompt is required.")
 
     with generation_lock:
 
         try:
-            progress(0, desc="Loading Pixelship 2")
+            progress(0.0, desc="Loading model")
 
             pipeline = load_pipeline()
 
@@ -106,74 +106,84 @@ def generate(
 
             kwargs = {
                 "prompt": prompt,
-                "width": int(width),
-                "height": int(height),
                 "num_inference_steps": int(steps),
                 "true_cfg_scale": float(true_cfg),
                 "generator": generator,
             }
 
+            if input_image is not None:
+                kwargs["image"] = input_image
+            else:
+                kwargs["width"] = int(width)
+                kwargs["height"] = int(height)
+
             if float(true_cfg) > 1.0:
-                kwargs["negative_prompt"] = negative_prompt or ""
+                kwargs["negative_prompt"] = (
+                    negative_prompt or ""
+                )
 
-            progress(0.1, desc="Generating image")
+            progress(0.15, desc="Generating")
 
             with torch.no_grad():
                 result = pipeline(**kwargs)
 
             image = result.images[0]
 
-            progress(1, desc="Generation complete")
+            mode = (
+                "Image Editing"
+                if input_image is not None
+                else "Text Generation"
+            )
+
+            progress(1.0, desc="Complete")
 
             return (
                 image,
-                f"Seed: {seed} | Steps: {steps} | "
-                f"True CFG: {true_cfg} | "
-                f"LoRA Strength: {lora_strength}"
+                f"{mode} | Seed={seed} | "
+                f"Steps={steps} | "
+                f"CFG={true_cfg} | "
+                f"LoRA={lora_strength}"
             )
 
         except torch.cuda.OutOfMemoryError:
+
             gc.collect()
 
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
 
             raise gr.Error(
-                "GPU ran out of memory. "
-                "Try reducing the resolution."
+                "Out of VRAM. Try a smaller image."
             )
 
         except Exception as e:
+
             traceback.print_exc()
+
             raise gr.Error(
-                f"Generation failed: {e}"
+                f"Generation failed:\n{e}"
             )
 
 
 # ============================================================
-# INTERFACE
+# UI
 # ============================================================
 
 CSS = """
 .gradio-container {
-    max-width: 1500px !important;
+    max-width: 1600px !important;
     margin: auto !important;
 }
 
 #prompt-box textarea {
-    min-height: 260px !important;
+    min-height: 240px !important;
     font-family: monospace !important;
-    font-size: 14px !important;
 }
 
 #generate-button {
     font-size: 18px !important;
     font-weight: 700 !important;
 }
-
-#output-image {
-    min-height: 550px;
-}
 """
 
 with gr.Blocks(
@@ -186,32 +196,49 @@ with gr.Blocks(
         """
 # 🎨 Pixelship 2
 
-### Powered by Qwen Image 2.1
+### Qwen Image 2.1 + Pixelship 2.0
+
+Supports:
 
-Generate images using the Pixelship 2 LoRA.
+✅ Text → Image
 
-**LoRA loaded directly from Hugging Face • ZeroGPU**
+✅ Image → Image Editing
+
+✅ Adjustable LoRA Strength
+
+✅ ZeroGPU Compatible
 """
     )
 
     with gr.Row():
 
+        # ========================================================
+        # LEFT
+        # ========================================================
+
         with gr.Column(scale=5):
 
             prompt = gr.Textbox(
                 label="Prompt",
-                placeholder="Describe your image in as much detail as you want...",
                 lines=12,
                 max_lines=40,
                 max_length=None,
                 elem_id="prompt-box",
+                placeholder=(
+                    "Describe what you want..."
+                ),
+            )
+
+            input_image = gr.Image(
+                label="Image To Edit (Optional)",
+                type="pil",
             )
 
             negative_prompt = gr.Textbox(
                 label="Negative Prompt",
-                placeholder="Optional negative prompt...",
                 lines=3,
                 max_length=None,
+                placeholder="Optional",
             )
 
             with gr.Accordion(
@@ -222,16 +249,16 @@ Generate images using the Pixelship 2 LoRA.
                 with gr.Row():
 
                     width = gr.Slider(
-                        minimum=256,
-                        maximum=2048,
+                        256,
+                        2048,
                         value=1024,
                         step=64,
                         label="Width",
                     )
 
                     height = gr.Slider(
-                        minimum=256,
-                        maximum=2048,
+                        256,
+                        2048,
                         value=1024,
                         step=64,
                         label="Height",
@@ -240,24 +267,24 @@ Generate images using the Pixelship 2 LoRA.
                 with gr.Row():
 
                     steps = gr.Slider(
-                        minimum=1,
-                        maximum=100,
+                        1,
+                        100,
                         value=30,
                         step=1,
-                        label="Inference Steps",
+                        label="Steps",
                     )
 
                     true_cfg = gr.Slider(
-                        minimum=1.0,
-                        maximum=10.0,
+                        1.0,
+                        10.0,
                         value=1.0,
                         step=0.1,
-                        label="True CFG Scale",
+                        label="True CFG",
                     )
 
                 lora_strength = gr.Slider(
-                    minimum=0.0,
-                    maximum=2.0,
+                    0.0,
+                    2.0,
                     value=0.75,
                     step=0.05,
                     label="LoRA Strength",
@@ -277,25 +304,27 @@ Generate images using the Pixelship 2 LoRA.
                     )
 
             generate_button = gr.Button(
-                "Generate Image",
+                "Generate",
                 variant="primary",
                 size="lg",
                 elem_id="generate-button",
             )
 
+        # ========================================================
+        # RIGHT
+        # ========================================================
+
         with gr.Column(scale=5):
 
             output_image = gr.Image(
-                label="Generated Image",
+                label="Output",
                 type="pil",
                 format="png",
-                interactive=False,
-                height=650,
-                elem_id="output-image",
+                height=700,
             )
 
             status = gr.Textbox(
-                label="Generation Status",
+                label="Status",
                 interactive=False,
             )
 
@@ -303,6 +332,7 @@ Generate images using the Pixelship 2 LoRA.
         fn=generate,
         inputs=[
             prompt,
+            input_image,
             negative_prompt,
             width,
             height,
@@ -319,9 +349,8 @@ Generate images using the Pixelship 2 LoRA.
         concurrency_limit=1,
     )
 
-
 # ============================================================
-# QUEUE / LAUNCH
+# QUEUE
 # ============================================================
 
 demo.queue(
@@ -329,6 +358,10 @@ demo.queue(
     default_concurrency_limit=1,
 )
 
+# ============================================================
+# LAUNCH
+# ============================================================
+
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",