Demo version

app.py

@@ -3,61 +3,52 @@ import zipfile
 import io
 import os
 from datetime import datetime
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-import torch
-
-# Global model variables
-model = None
-tokenizer = None
-
-def load_model():
-    """Load the small AI model for code generation."""
-    global model, tokenizer
-    if model is None:
-        try:
-            # Use a very small model that can run on Hugging Face Spaces free tier
-            model_name = "distilgpt2"
-
-            print("Loading AI model...")
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModelForCausalLM.from_pretrained(model_name)
-
-            # Set padding token to EOS token for GPT-2
-            tokenizer.pad_token = tokenizer.eos_token
-
-            print("Model loaded successfully!")
-        except Exception as e:
-            print(f"Error loading model: {e}")
-            return False
-    return True
+import requests
 
-def generate_code_with_ai(prompt, max_length=500):
-    """Generate code using the AI model."""
-    if not load_model():
-        return "// Error: Could not load AI model\n// Using fallback template instead"
+# Hugging Face API configuration
+HF_API_URL = "https://api-inference.huggingface.co/models/microsoft/DialoGPT-small"
+HF_TOKEN = os.environ.get("HF_TOKEN")  # Will be set in Hugging Face Spaces
 
+def generate_code_with_ai(prompt, max_length=500):
+    """Generate code using the Hugging Face Inference API."""
     try:
         # Create a more specific prompt for Flutter/Dart code
-        full_prompt = f"Write Flutter Dart code for: {prompt}\n\n"
+        full_prompt = f"""Write clean Flutter Dart code for this requirement: {prompt}
+
+Generate proper Dart classes, methods, and Flutter widgets. Focus on:
+- Clean, readable code
+- Proper Flutter/Dart syntax
+- Material Design patterns
+- Error handling where appropriate
 
-        inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True, max_length=100)
+```dart
+"""
 
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs["input_ids"],
-                attention_mask=inputs["attention_mask"],
-                max_length=max_length,
-                num_return_sequences=1,
-                temperature=0.7,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-            )
+        headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
 
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        payload = {
+            "inputs": full_prompt,
+            "parameters": {
+                "max_length": max_length,
+                "temperature": 0.7,
+                "do_sample": True,
+                "return_full_text": False,
+                "num_return_sequences": 1,
+            }
+        }
 
-        # Extract the code part (remove the prompt)
-        code = generated_text[len(full_prompt):].strip()
+        response = requests.post(HF_API_URL, headers=headers, json=payload)
+        response.raise_for_status()
+
+        result = response.json()
+
+        if isinstance(result, list) and len(result) > 0:
+            generated_text = result[0].get("generated_text", "")
+        else:
+            generated_text = str(result)
+
+        # Clean up the generated text
+        code = generated_text.strip()
 
         # Basic cleanup - try to extract Dart-like code
         if "```dart" in code:
@@ -65,11 +56,14 @@ def generate_code_with_ai(prompt, max_length=500):
         elif "```" in code:
             code = code.split("```")[1].split("```")[0]
 
+        # Remove any remaining markdown formatting
+        code = code.replace("```", "").strip()
+
         return code if code else "// AI-generated placeholder code"
 
     except Exception as e:
         print(f"Error generating code with AI: {e}")
-        return "// Error generating code\n// Using fallback template"
+        return "// Error: AI service unavailable\n// Using fallback template"
 
 def generate_flutter_code(description):
     """
@@ -721,10 +715,10 @@ def create_interface():
             - **README.md** - Setup and usage instructions
             
             ### ⚠️ Important Notes:
-            - This uses a **small AI model** (DistilGPT-2) running directly on Hugging Face Spaces
+            - This uses the **Hugging Face Inference API** for AI-powered code generation
             - The AI generates Flutter/Dart code based on your description
             - Always review and test generated code before production use
-            - The model may produce varying quality results - iterate on your descriptions
+            - The AI service may occasionally be slow or unavailable - try again if needed
             - Customize the generated code to match your specific requirements
             
             ### 🔧 Next Steps After Download: