Update app.py

app.py

@@ -2,24 +2,16 @@ import os
 import random
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline, UNet2DConditionModel
+from diffusers import DiffusionPipeline
 from transformers import CLIPTextModel, CLIPTokenizer
 
 # Configuration - Using Flux Model
 MODEL_ID = "CompVis/Flux-Pro"
-LORA_ID = "flux/lora-weights"
 MODEL_CACHE = "model_cache"
 os.makedirs(MODEL_CACHE, exist_ok=True)
 
 def get_pipeline():
-    # Load Flux components
-    unet = UNet2DConditionModel.from_pretrained(
-        MODEL_ID,
-        subfolder="unet",
-        cache_dir=MODEL_CACHE,
-        torch_dtype=torch.float32
-    )
-    
+    # Load Flux model components
     text_encoder = CLIPTextModel.from_pretrained(
         MODEL_ID,
         subfolder="text_encoder",
@@ -35,7 +27,6 @@ def get_pipeline():
     # Create pipeline
     pipe = DiffusionPipeline.from_pretrained(
         MODEL_ID,
-        unet=unet,
         text_encoder=text_encoder,
         tokenizer=tokenizer,
         cache_dir=MODEL_CACHE,
@@ -43,14 +34,6 @@ def get_pipeline():
         safety_checker=None
     )
     
-    # Load LoRA weights
-    lora_path = hf_hub_download(
-        LORA_ID,
-        "flux_lora.safetensors",
-        cache_dir=MODEL_CACHE
-    )
-    pipe.unet.load_attn_procs(lora_path)
-    
     # CPU optimizations
     pipe = pipe.to("cpu")
     pipe.enable_attention_slicing()