import spaces
import torch
import gradio as gr
from PIL import Image
import random
from diffusers import (
    DiffusionPipeline,
    AutoencoderKL,
    StableDiffusionControlNetPipeline,
    ControlNetModel,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionImg2ImgPipeline,
    StableDiffusionControlNetImg2ImgPipeline,
    DPMSolverMultistepScheduler,
    EulerDiscreteScheduler
)
import tempfile
import time
import os

BASE_MODEL = "SG161222/Realistic_Vision_V5.1_noVAE"

# Initialize both pipelines
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=torch.float16)

# Commenting out safety checker initialization
# SAFETY_CHECKER_ENABLED = os.environ.get("SAFETY_CHECKER", "0") == "1"
# safety_checker = None
# feature_extractor = None
# if SAFETY_CHECKER_ENABLED:
#     safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to("cuda")
#     feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")

main_pipe = StableDiffusionControlNetPipeline.from_pretrained(
    BASE_MODEL,
    controlnet=controlnet,
    vae=vae,
    # Remove safety checker and feature extractor from pipeline components
    # safety_checker=safety_checker,
    # feature_extractor=feature_extractor,
    torch_dtype=torch.float16,
).to("cuda")

def inference(control_image: Image.Image, prompt: str, negative_prompt: str, guidance_scale: float = 8.0, controlnet_conditioning_scale: float = 1, upscaler_strength: float = 0.5, seed: int = -1):
    # Inference logic remains unchanged
    ...