Upload app.py with huggingface_hub

app.py

@@ -1,67 +1,47 @@
 import gradio as gr
-import torch
 import os
-from diffusers import AmusedPipeline
-from PIL import Image
 from huggingface_hub import login
+import requests
+import base64
+from PIL import Image
+import io
 
 # Login to Hugging Face
 if os.getenv("HF_TOKEN"):
     login(token=os.getenv("HF_TOKEN"))
 
-# Global variables
-pipe = None
-
-def load_model():
-    """Load the model (cached after first load)"""
-    global pipe
-    
-    if pipe is not None:
-        return pipe
-    
+def generate_image_mlx(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, seed):
+    """Generate image using MLX via Hugging Face Inference API"""
     try:
-        print("Loading aMUSEd (lightweight model)...")
-        # Load aMUSEd - much smaller model designed for CPU
-        pipe = AmusedPipeline.from_pretrained(
-            "amused/amused-512",
-            torch_dtype=torch.float32,  # Use float32 for CPU
-            use_safetensors=True
-        )
-        
-        # Move to device
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        pipe = pipe.to(device)
-        print(f"Model loaded on {device}")
-        
-        return pipe
+        print(f"Generating image with prompt: {prompt}")
         
-    except Exception as e:
-        print(f"Error loading model: {str(e)}")
-        raise e
-
-def generate_image(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, seed):
-    """Generate image using aMUSEd"""
-    try:
-        # Load model if not already loaded
-        pipe = load_model()
+        # Use Hugging Face Inference API with MLX
+        api_url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
+        headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
         
-        print(f"Generating image with prompt: {prompt}")
-        # Generate image
-        with torch.no_grad():
-            generator = torch.Generator().manual_seed(seed) if seed else None
-            image = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale,
-                width=width,
-                height=height,
-                generator=generator
-            ).images[0]
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "negative_prompt": negative_prompt,
+                "num_inference_steps": num_inference_steps,
+                "guidance_scale": guidance_scale,
+                "width": width,
+                "height": height,
+                "seed": seed if seed else None
+            }
+        }
         
-        print("Image generated successfully!")
-        return image
+        response = requests.post(api_url, headers=headers, json=payload)
         
+        if response.status_code == 200:
+            # Convert response to PIL Image
+            image = Image.open(io.BytesIO(response.content))
+            print("Image generated successfully with MLX!")
+            return image
+        else:
+            print(f"API request failed: {response.status_code} - {response.text}")
+            return None
+            
     except Exception as e:
         print(f"Error generating image: {str(e)}")
         return None
@@ -69,8 +49,8 @@ def generate_image(prompt, negative_prompt, num_inference_steps, guidance_scale,
 # Create Gradio interface
 with gr.Blocks(title="Silicon Photonics Image Generator") as demo:
     gr.Markdown("# 🔬 Silicon Photonics Image Generator")
-    gr.Markdown("Generate technical diagrams and visualizations for silicon photonics using aMUSEd (lightweight model).")
-    gr.Markdown("**🚀 Now using aMUSEd for fast CPU inference!**")
+    gr.Markdown("Generate technical diagrams and visualizations for silicon photonics using MLX-powered inference.")
+    gr.Markdown("**🚀 Powered by MLX for ultra-fast Apple Silicon performance!**")
     
     with gr.Row():
         with gr.Column():
@@ -86,12 +66,12 @@ with gr.Blocks(title="Silicon Photonics Image Generator") as demo:
             )
             
             with gr.Row():
-                num_inference_steps = gr.Slider(1, 20, value=10, label="Inference Steps")
+                num_inference_steps = gr.Slider(1, 50, value=20, label="Inference Steps")
                 guidance_scale = gr.Slider(1.0, 20.0, value=7.0, label="Guidance Scale")
             
             with gr.Row():
-                width = gr.Slider(256, 512, value=512, step=64, label="Width")
-                height = gr.Slider(256, 512, value=512, step=64, label="Height")
+                width = gr.Slider(512, 1024, value=1024, step=64, label="Width")
+                height = gr.Slider(512, 1024, value=1024, step=64, label="Height")
                 seed = gr.Number(value=42, label="Seed (0 for random)")
             
             generate_btn = gr.Button("🎨 Generate Image", variant="primary")
@@ -116,7 +96,7 @@ with gr.Blocks(title="Silicon Photonics Image Generator") as demo:
     
     # Event handlers
     generate_btn.click(
-        fn=generate_image,
+        fn=generate_image_mlx,
         inputs=[prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, seed],
         outputs=[output_image]
     )