save4

app.py

@@ -3,62 +3,93 @@ from transformers import AutoProcessor, AutoModelForVision2Seq
 from huggingface_hub import login
 import gradio as gr
 import os
+import gc
 
 # ----------------------------
 # AUTHENTICATION
 # ----------------------------
-# Option 1: Use HF token from environment variable (recommended for Spaces)
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
 else:
-    # Option 2: Interactive login (for local testing)
     print("No HF_TOKEN found. Please log in manually.")
     login()
 
 # ----------------------------
 # CONFIG
 # ----------------------------
-MODEL_NAME = "reverseforward/inferencemodel"  # change this to your repo name
+MODEL_NAME = "reverseforward/inferencemodel"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE = torch.float16  # use float16 on A10G
+DTYPE = torch.float16
+
+# Clear cache before loading
+gc.collect()
+if DEVICE == "cuda":
+    torch.cuda.empty_cache()
 
 # ----------------------------
-# LOAD MODEL
+# LOAD MODEL (with error handling)
 # ----------------------------
-print("Loading model...")
-model = AutoModelForVision2Seq.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=DTYPE,
-    device_map="auto",
-    token=HF_TOKEN,  # Pass token explicitly
-)
-processor = AutoProcessor.from_pretrained(
-    MODEL_NAME,
-    token=HF_TOKEN,
-)
-print("Model loaded successfully.")
+print(f"Loading model on {DEVICE}...")
+try:
+    model = AutoModelForVision2Seq.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=DTYPE,
+        device_map="auto",
+        token=HF_TOKEN,
+        low_cpu_mem_usage=True,  # Reduce memory usage
+    )
+    processor = AutoProcessor.from_pretrained(
+        MODEL_NAME,
+        token=HF_TOKEN,
+    )
+    print("✓ Model loaded successfully.")
+except Exception as e:
+    print(f"✗ Error loading model: {e}")
+    raise
 
 # ----------------------------
 # INFERENCE FUNCTION
 # ----------------------------
 def chat_with_image(image, text):
-    if image is None or text.strip() == "":
-        return "Please provide both an image and text input."
+    try:
+        if image is None or text.strip() == "":
+            return "Please provide both an image and text input."
 
-    # Prepare inputs for Qwen3-VL
-    inputs = processor(text=[text], images=[image], return_tensors="pt").to(DEVICE, DTYPE)
+        # Clear memory before inference
+        gc.collect()
+        if DEVICE == "cuda":
+            torch.cuda.empty_cache()
 
-    # Generate output
-    with torch.inference_mode():
-        generated_ids = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            temperature=0.7,
-        )
+        # Prepare inputs
+        inputs = processor(
+            text=[text],
+            images=[image],
+            return_tensors="pt"
+        ).to(DEVICE, DTYPE)
 
-    output = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return output.strip()
+        # Generate output
+        with torch.inference_mode():
+            generated_ids = model.generate(
+                **inputs,
+                max_new_tokens=256,
+                temperature=0.7,
+                do_sample=True,
+            )
+
+        output = processor.batch_decode(
+            generated_ids,
+            skip_special_tokens=True
+        )[0]
+        
+        # Clean up
+        del inputs, generated_ids
+        gc.collect()
+        
+        return output.strip()
+    
+    except Exception as e:
+        return f"Error during inference: {str(e)}"
 
 
 # ----------------------------
@@ -74,16 +105,14 @@ demo = gr.Interface(
     fn=chat_with_image,
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
-        gr.Textbox(label="Enter Instruction or Question"),
+        gr.Textbox(label="Enter Instruction or Question", lines=3),
     ],
-    outputs=gr.Textbox(label="Model Output"),
+    outputs=gr.Textbox(label="Model Output", lines=5),
     title=title,
     description=description,
-    examples=[
-        ["examples/cat.jpg", "Describe this image."],
-        ["examples/room.jpg", "How many chairs are visible?"],
-    ],
+
+    allow_flagging="never",  # Disable flagging to reduce overhead
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(show_error=True)