Update app.py

app.py

@@ -1,65 +1,63 @@
 import gradio as gr
-from diffusers import StableDiffusionPipeline
-import torch
-from PIL import Image, ImageDraw
-import langid
-from transformers import pipeline  # Optional fallback for lang detection
-
-# Load Stable Diffusion model (use GPU if available)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
-pipe = pipe.to(device)
-
-# Language mapping
-lang_map = {
-    "si": "Sinhala",
-    "en": "English",
-    "ta": "Tamil"
-}
-
-def detect_language(text):
-    lang, _ = langid.classify(text)
-    if lang in lang_map:
-        return lang_map[lang]
-    return "English"  # Default fallback
-
-def generate_design(text, main_image=None, logo=None):
-    # Detect language
-    language = detect_language(text)
-
-    # Create prompt for design generation
-    prompt = f"Beautiful social media post graphic design with the text '{text}' in {language} font, elegant layout, modern shapes, vibrant colors, clean background, high resolution"
-
-    # Generate base image (512x512 for social media square post)
-    image = pipe(prompt).images[0]
-
-    # If main image uploaded, resize and paste centered
-    if main_image:
-        main_img = Image.open(main_image).resize((300, 300))  # Resize to fit
-        image.paste(main_img, ( (image.width - main_img.width) // 2, (image.height - main_img.height) // 2 ))
-
-    # If logo uploaded, resize and paste in bottom-right corner
-    if logo:
-        logo_img = Image.open(logo).resize((100, 100))  # Small logo size
-        image.paste(logo_img, (image.width - logo_img.width - 10, image.height - logo_img.height - 10), logo_img if logo_img.mode == 'RGBA' else None)
-
-    # Add subtle border or shapes (optional enhancement with PIL)
-    draw = ImageDraw.Draw(image)
-    draw.rectangle([(10, 10), (image.width-10, image.height-10)], outline="white", width=5)
-
-    return image
-
-# Gradio interface
+import os, zipfile, tempfile
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import json
+
+# ---------------------------
+# 1️⃣ Load Local Model
+# ---------------------------
+model_name = "TheBloke/WizardCoder-7B-GPTQ"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
+
+# ---------------------------
+# 2️⃣ Website Generation Function
+# ---------------------------
+def generate_website(prompt):
+    # Generate multi-file code as JSON
+    instruction = f"""
+Generate a full website project structure based on the user prompt. 
+Provide the response in JSON format where key = file path, value = file content. 
+Include folders: css/, js/, assets/, db/ if needed. User prompt: {prompt}
+"""
+    gen = generator(instruction, max_new_tokens=2048)[0]['generated_text']
+
+    # Parse JSON safely
+    try:
+        files_dict = json.loads(gen)
+    except:
+        files_dict = {"index.html": "<!-- Failed to generate code -->"}
+
+    # Create temp folder and save files
+    temp_dir = tempfile.mkdtemp()
+    for fname, content in files_dict.items():
+        file_path = os.path.join(temp_dir, fname)
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+
+    # Zip the folder
+    zip_path = os.path.join(temp_dir, "website.zip")
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(temp_dir):
+            for file in files:
+                if file != os.path.basename(zip_path):
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, temp_dir)
+                    zipf.write(file_path, arcname=arcname)
+
+    return zip_path
+
+# ---------------------------
+# 3️⃣ Gradio UI
+# ---------------------------
 iface = gr.Interface(
-    fn=generate_design,
-    inputs=[
-        gr.Textbox(label="Enter Text for the Post"),
-        gr.Image(type="filepath", label="Upload Main Image (e.g., Product/Human) - Optional"),
-        gr.Image(type="filepath", label="Upload Business Logo - Optional")
-    ],
-    outputs=gr.Image(label="Generated Social Media Post Design"),
-    title="AI Graphic Design Tool for Social Media",
-    description="Input text, optionally upload images/logo. Detects Sinhala/English/Tamil and generates a design."
+    fn=generate_website,
+    inputs=gr.Textbox(lines=4, placeholder="Describe your website: PHP/HTML/JS, admin panel, MySQL..."),
+    outputs=gr.File(label="Download Generated Website ZIP"),
+    title="Local AI Multi-file Website Generator",
+    description="Single Python file app. Generates multi-file website locally with folders (css/, js/, assets/, db/). Download as ZIP."
 )
 
 iface.launch()