Update app.py

app.py

@@ -10,21 +10,20 @@ from diffusers import FluxPipeline
 from insightface.app import FaceAnalysis
 from insightface.model_zoo import get_model
 
-# --- GLOBAL CONFIG (CPU ONLY) ---
+# --- GLOBAL CONFIG ---
 MODEL_ID = "black-forest-labs/FLUX.1-dev"
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# We define the models as None globally and load them inside the GPU function
+# Initialize models as None for ZeroGPU lazy loading
 face_app = None
 swapper = None
 pipe = None
 
-def load_models():
-    """Initializes models inside the GPU-allocated context."""
+def load_models_on_gpu():
+    """Initializes models only when GPU is allocated."""
     global face_app, swapper, pipe
     
     if face_app is None:
-        # Use CPU provider initially to avoid startup crashes
         face_app = FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
         face_app.prepare(ctx_id=0, det_size=(640, 640))
 
@@ -39,7 +38,6 @@ def load_models():
             torch_dtype=torch.bfloat16,
             token=HF_TOKEN
         )
-        # Offloading helps manage ZeroGPU's 70GB VRAM efficiently
         pipe.enable_model_cpu_offload()
 
 def upscale_image(image):
@@ -52,28 +50,30 @@ def upscale_image(image):
 
 @spaces.GPU(duration=150)
 def generate_vton_final(face_image, body_type, height_ft):
-    if face_image is None: return None, "Upload face image."
+    if face_image is None: 
+        return None, "Please upload a face image."
     
-    # Trigger model loading within the GPU context
-    load_models()
+    # Ensure models are loaded in the GPU context
+    load_models_on_gpu()
 
     # 1. Face Analysis
     img_np = np.array(face_image)
     cv_img = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
     faces = face_app.get(cv_img)
-    if not faces: return None, "No face detected."
+    if not faces: 
+        return None, "No face detected in the upload."
     
     source_face = faces[0]
     gender = "man" if source_face.gender == 1 else "woman"
 
-    # 2. Simplified Prompt (Normal Pose)
+    # 2. Simplified Prompt (Normal Pose & Casual Clothes)
     profile_seed = int(hashlib.md5(f"{gender}-{body_type}".encode()).hexdigest(), 16) % (10**9)
     generator = torch.Generator("cuda").manual_seed(profile_seed)
 
     prompt = (
-        f"Full body photo of a {gender}, {body_type} build, {height_ft}ft tall. "
-        f"Standing in a relaxed normal pose, facing camera, neutral expression. "
-        f"Casual high-quality clothing, simple studio background, sharp focus, 8k ultra-hd."
+        f"A full body 8k professional photo of a {gender}, {body_type} build, {height_ft}ft tall. "
+        f"Standing in a relaxed, natural pose, facing the camera. "
+        f"Wearing stylish casual clothing, clean studio background, sharp focus, cinematic lighting."
     )
 
     # 3. Generation
@@ -91,25 +91,26 @@ def generate_vton_final(face_image, body_type, height_ft):
         res_cv = cv2.cvtColor(res_np, cv2.COLOR_RGB2BGR)
         target_faces = face_app.get(res_cv)
         if target_faces:
-            # Swap onto the largest detected face
+            # Sort to find the main person in the photo
             target_faces = sorted(target_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]), reverse=True)
             res_cv = swapper.get(res_cv, target_faces[0], source_face, paste_back=True)
             gen_img = Image.fromarray(cv2.cvtColor(res_cv, cv2.COLOR_BGR2RGB))
 
-    return upscale_image(gen_img), f"Status: Complete | Seed: {profile_seed}"
+    # 5. HD Upscale
+    return upscale_image(gen_img), f"Success | Seed: {profile_seed}"
 
-# --- GRADIO ---
-with gr.Blocks(css=".gradio-container {background-color: #f9f9f9}") as demo:
-    gr.Markdown("## 💎 HD Virtual Model Generator")
+# --- GRADIO INTERFACE ---
+with gr.Blocks() as demo:
+    gr.Markdown("# 💎 AI Virtual Model Engine")
     with gr.Row():
         with gr.Column():
-            face_in = gr.Image(type="pil", label="Face Photo")
-            body_in = gr.Radio(["slim", "muscular", "average"], value="average", label="Body Shape")
-            h_in = gr.Slider(4.5, 6.5, value=5.7, step=0.1, label="Height (ft)")
-            btn = gr.Button("Generate Model", variant="primary")
+            face_in = gr.Image(type="pil", label="Step 1: Upload Face")
+            body_in = gr.Radio(["slim", "muscular", "average"], value="average", label="Step 2: Body Build")
+            h_in = gr.Slider(4.5, 7.0, value=5.8, step=0.1, label="Step 3: Height (ft)")
+            btn = gr.Button("Generate High-Res Model", variant="primary")
         with gr.Column():
-            img_out = gr.Image(label="Result")
-            status = gr.Textbox(label="System Logs")
+            img_out = gr.Image(label="Final Result")
+            status = gr.Textbox(label="Logs")
 
     btn.click(generate_vton_final, [face_in, body_in, h_in], [img_out, status])