Update app.py

app.py

@@ -1,187 +1,73 @@
+import os
 import gradio as gr
-import numpy as np
-import random
-from diffusers import DiffusionPipeline
-import torch
-
-from diffusers import DiffusionPipeline
 import torch
-from huggingface_hub import login
-
-import os
-from huggingface_hub import login
+from transformers import AutoTokenizer, AutoModel
+from safetensors.torch import load_file
 
-# Access the token from environment variables
+# Load the Hugging Face API token
 token = os.getenv("HUGGINGFACE_API_TOKEN")
 if not token:
-    raise ValueError("API token is missing! Please set it in the Hugging Face Space Secrets.")
+    raise ValueError("HUGGINGFACE_API_TOKEN is not set. Please add it in the Secrets section of your Space.")
 
-# Log in to Hugging Face
-login(token)
-
-
-# Model details
+# Configure device and data type
 device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-# Load model
-pipe = DiffusionPipeline.from_pretrained(
-    "Grandediw/lora_model",
-    torch_dtype=torch_dtype,
-    use_auth_token=True  # Enables private model access
-)
-pipe = pipe.to(device)
-
-
-# Constants
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
 
-# Inference function
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+# Load the tokenizer and model
+model_repo = "Grandediw/lora_model"
+tokenizer = AutoTokenizer.from_pretrained(model_repo, use_auth_token=True)
+base_model = AutoModel.from_pretrained(model_repo, use_auth_token=True)
 
-    generator = torch.Generator(device).manual_seed(seed)
+# Load LoRA adapter weights
+lora_weights_path = "adapter_model.safetensors"  # Ensure this file is present in the same directory
+lora_weights = load_file(lora_weights_path)
 
-    # Generate the image
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+# Apply LoRA weights to the base model
+for name, param in base_model.named_parameters():
+    if name in lora_weights:
+        param.data += lora_weights[name].to(device, dtype=param.dtype)
 
-    return image, seed
+# Move the model to the device
+base_model = base_model.to(device)
 
-# Example prompts
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
+# Inference function
+def infer(prompt, negative_prompt=None):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = base_model(**inputs)
+    return outputs.last_hidden_state.mean(dim=1).cpu().detach().numpy()  # Placeholder return
 
-# Improved CSS for better styling
+# Gradio Interface
 css = """
 #interface-container {
     margin: 0 auto;
     max-width: 700px;
-    padding: 10px;
-    box-shadow: 0px 4px 10px rgba(0, 0, 0, 0.1);
+    padding: 15px;
     border-radius: 10px;
     background-color: #f9f9f9;
+    box-shadow: 0px 4px 10px rgba(0, 0, 0, 0.1);
 }
 #header {
     text-align: center;
     font-size: 1.5em;
+    font-weight: bold;
     margin-bottom: 20px;
     color: #333;
 }
-#advanced-settings {
-    background-color: #f1f1f1;
-    padding: 10px;
-    border-radius: 8px;
-}
 """
 
-# Gradio interface
 with gr.Blocks(css=css) as demo:
     with gr.Box(elem_id="interface-container"):
-        gr.Markdown(
-            """
-            <div id="header">🖼️ Text-to-Image Generator</div>
-            Generate high-quality images from your text prompts with the fine-tuned LoRA model.
-            """
-        )
-
-        # Main input row
-        with gr.Row():
-            prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="Describe the image you want to create...",
-                lines=2,
-            )
-            run_button = gr.Button("Generate Image", variant="primary")
-
-        # Output image display
-        result = gr.Image(label="Generated Image").style(height="512px")
-
-        # Advanced settings
-        with gr.Accordion("Advanced Settings", open=False, elem_id="advanced-settings"):
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                placeholder="What to exclude from the image...",
-            )
-            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-            seed = gr.Number(label="Seed", value=0, interactive=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Image Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=64,
-                    value=512,
-                )
-                height = gr.Slider(
-                    label="Image Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=64,
-                    value=512,
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance Scale",
-                    minimum=0.0,
-                    maximum=20.0,
-                    step=0.1,
-                    value=7.5,
-                )
-                num_inference_steps = gr.Slider(
-                    label="Steps",
-                    minimum=10,
-                    maximum=100,
-                    step=5,
-                    value=50,
-                )
+        gr.Markdown("<div id='header'>LoRA Model Inference</div>")
+        
+        # Input for prompt and run button
+        prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
+        run_button = gr.Button("Generate Output", variant="primary")
 
-        # Examples
-        gr.Examples(
-            examples=examples,
-            inputs=[prompt],
-            outputs=[result],
-            label="Try these prompts",
-        )
+        # Display output
+        output = gr.Textbox(label="Output")
 
-    # Event handler
-    run_button.click(
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
+        # Connect button with inference
+        run_button.click(fn=infer, inputs=[prompt], outputs=[output])
 
+# Launch the app
 if __name__ == "__main__":
     demo.launch()