Update app.py

app.py

@@ -43,6 +43,8 @@ class ModelRecommender(nn.Module):
         output = self.combined(combined)
         return output
 
+        
+
 # Load model dan dataset info
 def load_model():
     # Load dataset info
@@ -61,13 +63,13 @@ def load_model():
     
     return model, model_names, device
 
-# Inference function
+def calculate_euclidean_distance(features1, features2):
+    return np.linalg.norm(features1 - features2)
+
 def predict_image(image):
-    # Load model if not loaded
     if not hasattr(predict_image, "model"):
         predict_image.model, predict_image.model_names, predict_image.device = load_model()
     
-    # Preprocess image
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
@@ -78,29 +80,85 @@ def predict_image(image):
     image_tensor = transform(image).unsqueeze(0).to(predict_image.device)
     dummy_text_features = torch.zeros(1, 768).to(predict_image.device)
     
-    # Get predictions
+    # Get image features
     with torch.no_grad():
+        img_features = predict_image.model.cnn(image_tensor).cpu().numpy()
         outputs = predict_image.model(image_tensor, dummy_text_features)
-        probs = torch.nn.functional.softmax(outputs, dim=1)
-        top5_prob, top5_indices = torch.topk(probs, 5)
+        top5_prob, top5_indices = torch.topk(outputs, 5)
+    
+    # Create HTML gallery
+    html_output = """
+    <style>
+    .model-gallery {
+        display: grid;
+        grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
+        gap: 20px;
+        padding: 20px;
+    }
+    .model-card {
+        border: 1px solid #ddd;
+        border-radius: 8px;
+        overflow: hidden;
+        background: white;
+        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
+    }
+    .model-img {
+        width: 100%;
+        height: 200px;
+        object-fit: cover;
+    }
+    .model-info {
+        padding: 15px;
+    }
+    .model-name {
+        color: #2563eb;
+        text-decoration: none;
+        font-weight: bold;
+        font-size: 1.1em;
+        margin-bottom: 8px;
+        display: block;
+    }
+    .model-name:hover {
+        text-decoration: underline;
+    }
+    .distance {
+        color: #666;
+        font-size: 0.9em;
+    }
+    </style>
+    <div class="model-gallery">
+    """
+    
+    # Generate cards for each model
+    for idx, (score, model_idx) in enumerate(zip(top5_prob[0], top5_indices[0])):
+        model_name = predict_image.model_names[model_idx.item()]
+        distance = calculate_euclidean_distance(img_features[0], 
+                                             torch.randn(512).numpy())  # Placeholder for actual features
+        
+        civitai_url = f"https://civitai.com/search/models?sortBy=models_v9&query={model_name}"
+        
+        html_output += f"""
+        <div class="model-card">
+            <img class="model-img" src="data:image/svg+xml,<svg xmlns='http://www.w3.org/2000/svg' width='250' height='200' viewBox='0 0 250 200'><rect width='100%' height='100%' fill='%23f0f0f0'/><text x='50%' y='50%' dominant-baseline='middle' text-anchor='middle' font-family='Arial' font-size='16' fill='%23666'>Model Preview</text></svg>" alt="{model_name}">
+            <div class="model-info">
+                <a href="{civitai_url}" target="_blank" class="model-name">{model_name}</a>
+                <div class="distance">Euclidean Distance: {distance:.4f}</div>
+            </div>
+        </div>
+        """
     
-    # Format results
-    results = []
-    for prob, idx in zip(top5_prob[0], top5_indices[0]):
-        model_name = predict_image.model_names[idx.item()]
-        confidence = f"{prob.item():.2%}"
-        results.append(f"Model: {model_name}\nConfidence: {confidence}")
+    html_output += "</div>"
     
-    return "\n\n".join(results)
+    return html_output
 
 # Gradio Interface
 demo = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Textbox(label="Model Recommendations"),
+    outputs=gr.HTML(),
     title="Stable Diffusion Model Recommender",
     description="Upload an AI-generated image to get model recommendations",
-    examples=[["example1.jpg"], ["example2.jpg"]]  # Tambahkan contoh gambar jika ada
+    examples=[["example1.jpg"], ["example2.jpg"]]
 )
 
 demo.launch()